mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-02-04 21:40:51 +08:00
Compare commits
39 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
a67b6811d1 | ||
|
|
35fdc4cfd3 | ||
|
|
3ebbab0a9a | ||
|
|
480cd714b2 | ||
|
|
41ee44432d | ||
|
|
1434bc38e5 | ||
|
|
0fd2abbc3b | ||
|
|
0ebb654019 | ||
|
|
08a1d2edf9 | ||
|
|
3409f4e336 | ||
|
|
9354b87e54 | ||
|
|
54e24110ec | ||
|
|
717c703bff | ||
|
|
1c6f4be8ae | ||
|
|
0de2560cee | ||
|
|
85eb926482 | ||
|
|
c52ef08e67 | ||
|
|
cb580cd083 | ||
|
|
75e278c7a5 | ||
|
|
73208c4e55 | ||
|
|
32d3809f8c | ||
|
|
a748e93fd9 | ||
|
|
54a9c4c3c7 | ||
|
|
18b5c35dea | ||
|
|
7b7871ede2 | ||
|
|
c4e3646b75 | ||
|
|
022aa81be1 | ||
|
|
c43f0ea7b1 | ||
|
|
6a191358af | ||
|
|
db1119dd78 | ||
|
|
33a5656235 | ||
|
|
2cd59806e2 | ||
|
|
5983e3ec87 | ||
|
|
f8cebb9343 | ||
|
|
72c7ef7647 | ||
|
|
d2e4639b2a | ||
|
|
08321223c4 | ||
|
|
7e30157590 | ||
|
|
e73cdf5cff |
5
.gitignore
vendored
5
.gitignore
vendored
@@ -15,6 +15,7 @@ pgstore/*
|
||||
gitstore/*
|
||||
objectstore/*
|
||||
static/*
|
||||
refs/*
|
||||
|
||||
# Authentication data
|
||||
auths/*
|
||||
@@ -30,3 +31,7 @@ GEMINI.md
|
||||
.vscode/*
|
||||
.claude/*
|
||||
.serena/*
|
||||
|
||||
# macOS
|
||||
.DS_Store
|
||||
._*
|
||||
|
||||
@@ -56,6 +56,7 @@ CLIProxyAPI includes integrated support for [Amp CLI](https://ampcode.com) and A
|
||||
- Provider route aliases for Amp's API patterns (`/api/provider/{provider}/v1...`)
|
||||
- Management proxy for OAuth authentication and account features
|
||||
- Smart model fallback with automatic routing
|
||||
- **Model mapping** to route unavailable models to alternatives (e.g., `claude-opus-4.5` → `claude-sonnet-4`)
|
||||
- Security-first design with localhost-only management endpoints
|
||||
|
||||
**→ [Complete Amp CLI Integration Guide](docs/amp-cli-integration.md)**
|
||||
@@ -90,6 +91,10 @@ Native macOS menu bar app to use your Claude Code & ChatGPT subscriptions with A
|
||||
|
||||
Browser-based tool to translate SRT subtitles using your Gemini subscription via CLIProxyAPI with automatic validation/error correction - no API keys needed
|
||||
|
||||
### [CCS (Claude Code Switch)](https://github.com/kaitranntt/ccs)
|
||||
|
||||
CLI wrapper for instant switching between multiple Claude accounts and alternative models (Gemini, Codex, Antigravity) via CLIProxyAPI OAuth - no API keys needed
|
||||
|
||||
> [!NOTE]
|
||||
> If you developed a project based on CLIProxyAPI, please open a PR to add it to this list.
|
||||
|
||||
|
||||
@@ -89,6 +89,10 @@ CLIProxyAPI 已内置对 [Amp CLI](https://ampcode.com) 和 Amp IDE 扩展的支
|
||||
|
||||
一款基于浏览器的 SRT 字幕翻译工具,可通过 CLI 代理 API 使用您的 Gemini 订阅。内置自动验证与错误修正功能,无需 API 密钥。
|
||||
|
||||
### [CCS (Claude Code Switch)](https://github.com/kaitranntt/ccs)
|
||||
|
||||
CLI 封装器,用于通过 CLIProxyAPI OAuth 即时切换多个 Claude 账户和替代模型(Gemini, Codex, Antigravity),无需 API 密钥。
|
||||
|
||||
> [!NOTE]
|
||||
> 如果你开发了基于 CLIProxyAPI 的项目,请提交一个 PR(拉取请求)将其添加到此列表中。
|
||||
|
||||
|
||||
@@ -55,6 +55,28 @@ quota-exceeded:
|
||||
# When true, enable authentication for the WebSocket API (/v1/ws).
|
||||
ws-auth: false
|
||||
|
||||
# Amp CLI Integration
|
||||
# Configure upstream URL for Amp CLI OAuth and management features
|
||||
#amp-upstream-url: "https://ampcode.com"
|
||||
|
||||
# Optional: Override API key for Amp upstream (otherwise uses env or file)
|
||||
#amp-upstream-api-key: ""
|
||||
|
||||
# Restrict Amp management routes (/api/auth, /api/user, etc.) to localhost only (recommended)
|
||||
#amp-restrict-management-to-localhost: true
|
||||
|
||||
# Amp Model Mappings
|
||||
# Route unavailable Amp models to alternative models available in your local proxy.
|
||||
# Useful when Amp CLI requests models you don't have access to (e.g., Claude Opus 4.5)
|
||||
# but you have a similar model available (e.g., Claude Sonnet 4).
|
||||
#amp-model-mappings:
|
||||
# - from: "claude-opus-4.5" # Model requested by Amp CLI
|
||||
# to: "claude-sonnet-4" # Route to this available model instead
|
||||
# - from: "gpt-5"
|
||||
# to: "gemini-2.5-pro"
|
||||
# - from: "claude-3-opus-20240229"
|
||||
# to: "claude-3-5-sonnet-20241022"
|
||||
|
||||
# Gemini API keys (preferred)
|
||||
#gemini-api-key:
|
||||
# - api-key: "AIzaSy...01"
|
||||
@@ -62,6 +84,11 @@ ws-auth: false
|
||||
# headers:
|
||||
# X-Custom-Header: "custom-value"
|
||||
# proxy-url: "socks5://proxy.example.com:1080"
|
||||
# excluded-models:
|
||||
# - "gemini-2.5-pro" # exclude specific models from this provider (exact match)
|
||||
# - "gemini-2.5-*" # wildcard matching prefix (e.g. gemini-2.5-flash, gemini-2.5-pro)
|
||||
# - "*-preview" # wildcard matching suffix (e.g. gemini-3-pro-preview)
|
||||
# - "*flash*" # wildcard matching substring (e.g. gemini-2.5-flash-lite)
|
||||
# - api-key: "AIzaSy...02"
|
||||
|
||||
# API keys for official Generative Language API (legacy compatibility)
|
||||
@@ -76,6 +103,11 @@ ws-auth: false
|
||||
# headers:
|
||||
# X-Custom-Header: "custom-value"
|
||||
# proxy-url: "socks5://proxy.example.com:1080" # optional: per-key proxy override
|
||||
# excluded-models:
|
||||
# - "gpt-5.1" # exclude specific models (exact match)
|
||||
# - "gpt-5-*" # wildcard matching prefix (e.g. gpt-5-medium, gpt-5-codex)
|
||||
# - "*-mini" # wildcard matching suffix (e.g. gpt-5-codex-mini)
|
||||
# - "*codex*" # wildcard matching substring (e.g. gpt-5-codex-low)
|
||||
|
||||
# Claude API keys
|
||||
#claude-api-key:
|
||||
@@ -88,6 +120,11 @@ ws-auth: false
|
||||
# models:
|
||||
# - name: "claude-3-5-sonnet-20241022" # upstream model name
|
||||
# alias: "claude-sonnet-latest" # client alias mapped to the upstream model
|
||||
# excluded-models:
|
||||
# - "claude-opus-4-5-20251101" # exclude specific models (exact match)
|
||||
# - "claude-3-*" # wildcard matching prefix (e.g. claude-3-7-sonnet-20250219)
|
||||
# - "*-think" # wildcard matching suffix (e.g. claude-opus-4-5-thinking)
|
||||
# - "*haiku*" # wildcard matching substring (e.g. claude-3-5-haiku-20241022)
|
||||
|
||||
# OpenAI compatibility providers
|
||||
#openai-compatibility:
|
||||
@@ -108,6 +145,19 @@ ws-auth: false
|
||||
# - name: "moonshotai/kimi-k2:free" # The actual model name.
|
||||
# alias: "kimi-k2" # The alias used in the API.
|
||||
|
||||
# Vertex API keys (Vertex-compatible endpoints, use API key + base URL)
|
||||
#vertex-api-key:
|
||||
# - api-key: "vk-123..." # x-goog-api-key header
|
||||
# base-url: "https://example.com/api" # e.g. https://zenmux.ai/api
|
||||
# proxy-url: "socks5://proxy.example.com:1080" # optional per-key proxy override
|
||||
# headers:
|
||||
# X-Custom-Header: "custom-value"
|
||||
# models: # optional: map aliases to upstream model names
|
||||
# - name: "gemini-2.0-flash" # upstream model name
|
||||
# alias: "vertex-flash" # client-visible alias
|
||||
# - name: "gemini-1.5-pro"
|
||||
# alias: "vertex-pro"
|
||||
|
||||
#payload: # Optional payload configuration
|
||||
# default: # Default rules only set parameters when they are missing in the payload.
|
||||
# - models:
|
||||
@@ -121,3 +171,25 @@ ws-auth: false
|
||||
# protocol: "codex" # restricts the rule to a specific protocol, options: openai, gemini, claude, codex
|
||||
# params: # JSON path (gjson/sjson syntax) -> value
|
||||
# "reasoning.effort": "high"
|
||||
|
||||
# OAuth provider excluded models
|
||||
#oauth-excluded-models:
|
||||
# gemini-cli:
|
||||
# - "gemini-2.5-pro" # exclude specific models (exact match)
|
||||
# - "gemini-2.5-*" # wildcard matching prefix (e.g. gemini-2.5-flash, gemini-2.5-pro)
|
||||
# - "*-preview" # wildcard matching suffix (e.g. gemini-3-pro-preview)
|
||||
# - "*flash*" # wildcard matching substring (e.g. gemini-2.5-flash-lite)
|
||||
# vertex:
|
||||
# - "gemini-3-pro-preview"
|
||||
# aistudio:
|
||||
# - "gemini-3-pro-preview"
|
||||
# antigravity:
|
||||
# - "gemini-3-pro-preview"
|
||||
# claude:
|
||||
# - "claude-3-5-haiku-20241022"
|
||||
# codex:
|
||||
# - "gpt-5-codex-mini"
|
||||
# qwen:
|
||||
# - "vision-model"
|
||||
# iflow:
|
||||
# - "tstars2.0"
|
||||
|
||||
@@ -8,6 +8,7 @@ This guide explains how to use CLIProxyAPI with Amp CLI and Amp IDE extensions,
|
||||
- [Which Providers Should You Authenticate?](#which-providers-should-you-authenticate)
|
||||
- [Architecture](#architecture)
|
||||
- [Configuration](#configuration)
|
||||
- [Model Mapping Configuration](#model-mapping-configuration)
|
||||
- [Setup](#setup)
|
||||
- [Usage](#usage)
|
||||
- [Troubleshooting](#troubleshooting)
|
||||
@@ -21,6 +22,7 @@ The Amp CLI integration adds specialized routing to support Amp's API patterns w
|
||||
- **Provider route aliases**: Maps Amp's `/api/provider/{provider}/v1...` patterns to CLIProxyAPI handlers
|
||||
- **Management proxy**: Forwards OAuth and account management requests to Amp's control plane
|
||||
- **Smart fallback**: Automatically routes unconfigured models to ampcode.com
|
||||
- **Model mapping**: Route unavailable models to alternatives you have access to (e.g., `claude-opus-4.5` → `claude-sonnet-4`)
|
||||
- **Secret management**: Configurable precedence (config > env > file) with 5-minute caching
|
||||
- **Security-first**: Management routes restricted to localhost by default
|
||||
- **Automatic gzip handling**: Decompresses responses from Amp upstream
|
||||
@@ -75,7 +77,10 @@ Amp CLI/IDE
|
||||
│ ↓
|
||||
│ ├─ Model configured locally?
|
||||
│ │ YES → Use local OAuth tokens (OpenAI/Claude/Gemini handlers)
|
||||
│ │ NO → Forward to ampcode.com (reverse proxy)
|
||||
│ │ NO ↓
|
||||
│ │ ├─ Model mapping configured?
|
||||
│ │ │ YES → Rewrite model → Use local handler (free)
|
||||
│ │ │ NO → Forward to ampcode.com (uses Amp credits)
|
||||
│ ↓
|
||||
│ Response
|
||||
│
|
||||
@@ -115,6 +120,49 @@ amp-upstream-url: "https://ampcode.com"
|
||||
amp-restrict-management-to-localhost: true
|
||||
```
|
||||
|
||||
### Model Mapping Configuration
|
||||
|
||||
When Amp CLI requests a model that you don't have access to, you can configure mappings to route those requests to alternative models that you DO have available. This avoids consuming Amp credits for models you could handle locally.
|
||||
|
||||
```yaml
|
||||
# Route unavailable models to alternatives
|
||||
amp-model-mappings:
|
||||
# Example: Route Claude Opus 4.5 requests to Claude Sonnet 4
|
||||
- from: "claude-opus-4.5"
|
||||
to: "claude-sonnet-4"
|
||||
|
||||
# Example: Route GPT-5 requests to Gemini 2.5 Pro
|
||||
- from: "gpt-5"
|
||||
to: "gemini-2.5-pro"
|
||||
|
||||
# Example: Map older model names to newer versions
|
||||
- from: "claude-3-opus-20240229"
|
||||
to: "claude-3-5-sonnet-20241022"
|
||||
```
|
||||
|
||||
**How it works:**
|
||||
|
||||
1. Amp CLI requests a model (e.g., `claude-opus-4.5`)
|
||||
2. CLIProxyAPI checks if a local provider is available for that model
|
||||
3. If not available, it checks the model mappings
|
||||
4. If a mapping exists, the request is rewritten to use the target model
|
||||
5. The request is then handled locally (free, using your OAuth subscription)
|
||||
|
||||
**Benefits:**
|
||||
- **Save Amp credits**: Use your local subscriptions instead of forwarding to ampcode.com
|
||||
- **Hot-reload**: Mappings can be updated without restarting the proxy
|
||||
- **Structured logging**: Clear logs show when mappings are applied
|
||||
|
||||
**Routing Decision Logs:**
|
||||
|
||||
The proxy logs each routing decision with structured fields:
|
||||
|
||||
```
|
||||
[AMP] Using local provider for model: gemini-2.5-pro # Local provider (free)
|
||||
[AMP] Model mapped: claude-opus-4.5 -> claude-sonnet-4 # Mapping applied (free)
|
||||
[AMP] Forwarding to ampcode.com (uses Amp credits) - model_id: gpt-5 # Fallback (costs credits)
|
||||
```
|
||||
|
||||
### Secret Resolution Precedence
|
||||
|
||||
The Amp module resolves API keys using this precedence order:
|
||||
@@ -301,11 +349,14 @@ When Amp requests a model:
|
||||
|
||||
1. **Check local configuration**: Does CLIProxyAPI have OAuth tokens for this model's provider?
|
||||
2. **If YES**: Route to local handler (use your OAuth subscription)
|
||||
3. **If NO**: Forward to ampcode.com (use Amp's default routing)
|
||||
3. **If NO**: Check if a model mapping exists
|
||||
4. **If mapping exists**: Rewrite request to mapped model → Route to local handler (free)
|
||||
5. **If no mapping**: Forward to ampcode.com (uses Amp credits)
|
||||
|
||||
This enables seamless mixed usage:
|
||||
- Models you've configured (Gemini, ChatGPT, Claude) → Your OAuth subscriptions
|
||||
- Models you haven't configured → Amp's default providers
|
||||
- Models with mappings configured → Routed to alternative local models (free)
|
||||
- Models you haven't configured and have no mapping → Amp's default providers (uses credits)
|
||||
|
||||
### Example API Calls
|
||||
|
||||
|
||||
@@ -891,6 +891,8 @@ func (h *Handler) RequestAnthropicToken(c *gin.Context) {
|
||||
|
||||
func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
|
||||
ctx := context.Background()
|
||||
proxyHTTPClient := util.SetProxy(&h.cfg.SDKConfig, &http.Client{})
|
||||
ctx = context.WithValue(ctx, oauth2.HTTPClient, proxyHTTPClient)
|
||||
|
||||
// Optional project ID from query
|
||||
projectID := c.Query("project_id")
|
||||
@@ -976,7 +978,7 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
|
||||
requestedProjectID := strings.TrimSpace(projectID)
|
||||
|
||||
// Create token storage (mirrors internal/auth/gemini createTokenStorage)
|
||||
httpClient := conf.Client(ctx, token)
|
||||
authHTTPClient := conf.Client(ctx, token)
|
||||
req, errNewRequest := http.NewRequestWithContext(ctx, "GET", "https://www.googleapis.com/oauth2/v1/userinfo?alt=json", nil)
|
||||
if errNewRequest != nil {
|
||||
log.Errorf("Could not get user info: %v", errNewRequest)
|
||||
@@ -986,7 +988,7 @@ func (h *Handler) RequestGeminiCLIToken(c *gin.Context) {
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token.AccessToken))
|
||||
|
||||
resp, errDo := httpClient.Do(req)
|
||||
resp, errDo := authHTTPClient.Do(req)
|
||||
if errDo != nil {
|
||||
log.Errorf("Failed to execute request: %v", errDo)
|
||||
oauthStatus[state] = "Failed to execute request"
|
||||
|
||||
@@ -223,6 +223,7 @@ func (h *Handler) PatchGeminiKey(c *gin.Context) {
|
||||
value.APIKey = strings.TrimSpace(value.APIKey)
|
||||
value.BaseURL = strings.TrimSpace(value.BaseURL)
|
||||
value.ProxyURL = strings.TrimSpace(value.ProxyURL)
|
||||
value.ExcludedModels = config.NormalizeExcludedModels(value.ExcludedModels)
|
||||
if value.APIKey == "" {
|
||||
// Treat empty API key as delete.
|
||||
if body.Index != nil && *body.Index >= 0 && *body.Index < len(h.cfg.GeminiKey) {
|
||||
@@ -504,6 +505,91 @@ func (h *Handler) DeleteOpenAICompat(c *gin.Context) {
|
||||
c.JSON(400, gin.H{"error": "missing name or index"})
|
||||
}
|
||||
|
||||
// oauth-excluded-models: map[string][]string
|
||||
func (h *Handler) GetOAuthExcludedModels(c *gin.Context) {
|
||||
c.JSON(200, gin.H{"oauth-excluded-models": config.NormalizeOAuthExcludedModels(h.cfg.OAuthExcludedModels)})
|
||||
}
|
||||
|
||||
func (h *Handler) PutOAuthExcludedModels(c *gin.Context) {
|
||||
data, err := c.GetRawData()
|
||||
if err != nil {
|
||||
c.JSON(400, gin.H{"error": "failed to read body"})
|
||||
return
|
||||
}
|
||||
var entries map[string][]string
|
||||
if err = json.Unmarshal(data, &entries); err != nil {
|
||||
var wrapper struct {
|
||||
Items map[string][]string `json:"items"`
|
||||
}
|
||||
if err2 := json.Unmarshal(data, &wrapper); err2 != nil {
|
||||
c.JSON(400, gin.H{"error": "invalid body"})
|
||||
return
|
||||
}
|
||||
entries = wrapper.Items
|
||||
}
|
||||
h.cfg.OAuthExcludedModels = config.NormalizeOAuthExcludedModels(entries)
|
||||
h.persist(c)
|
||||
}
|
||||
|
||||
func (h *Handler) PatchOAuthExcludedModels(c *gin.Context) {
|
||||
var body struct {
|
||||
Provider *string `json:"provider"`
|
||||
Models []string `json:"models"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&body); err != nil || body.Provider == nil {
|
||||
c.JSON(400, gin.H{"error": "invalid body"})
|
||||
return
|
||||
}
|
||||
provider := strings.ToLower(strings.TrimSpace(*body.Provider))
|
||||
if provider == "" {
|
||||
c.JSON(400, gin.H{"error": "invalid provider"})
|
||||
return
|
||||
}
|
||||
normalized := config.NormalizeExcludedModels(body.Models)
|
||||
if len(normalized) == 0 {
|
||||
if h.cfg.OAuthExcludedModels == nil {
|
||||
c.JSON(404, gin.H{"error": "provider not found"})
|
||||
return
|
||||
}
|
||||
if _, ok := h.cfg.OAuthExcludedModels[provider]; !ok {
|
||||
c.JSON(404, gin.H{"error": "provider not found"})
|
||||
return
|
||||
}
|
||||
delete(h.cfg.OAuthExcludedModels, provider)
|
||||
if len(h.cfg.OAuthExcludedModels) == 0 {
|
||||
h.cfg.OAuthExcludedModels = nil
|
||||
}
|
||||
h.persist(c)
|
||||
return
|
||||
}
|
||||
if h.cfg.OAuthExcludedModels == nil {
|
||||
h.cfg.OAuthExcludedModels = make(map[string][]string)
|
||||
}
|
||||
h.cfg.OAuthExcludedModels[provider] = normalized
|
||||
h.persist(c)
|
||||
}
|
||||
|
||||
func (h *Handler) DeleteOAuthExcludedModels(c *gin.Context) {
|
||||
provider := strings.ToLower(strings.TrimSpace(c.Query("provider")))
|
||||
if provider == "" {
|
||||
c.JSON(400, gin.H{"error": "missing provider"})
|
||||
return
|
||||
}
|
||||
if h.cfg.OAuthExcludedModels == nil {
|
||||
c.JSON(404, gin.H{"error": "provider not found"})
|
||||
return
|
||||
}
|
||||
if _, ok := h.cfg.OAuthExcludedModels[provider]; !ok {
|
||||
c.JSON(404, gin.H{"error": "provider not found"})
|
||||
return
|
||||
}
|
||||
delete(h.cfg.OAuthExcludedModels, provider)
|
||||
if len(h.cfg.OAuthExcludedModels) == 0 {
|
||||
h.cfg.OAuthExcludedModels = nil
|
||||
}
|
||||
h.persist(c)
|
||||
}
|
||||
|
||||
// codex-api-key: []CodexKey
|
||||
func (h *Handler) GetCodexKeys(c *gin.Context) {
|
||||
c.JSON(200, gin.H{"codex-api-key": h.cfg.CodexKey})
|
||||
@@ -533,6 +619,7 @@ func (h *Handler) PutCodexKeys(c *gin.Context) {
|
||||
entry.BaseURL = strings.TrimSpace(entry.BaseURL)
|
||||
entry.ProxyURL = strings.TrimSpace(entry.ProxyURL)
|
||||
entry.Headers = config.NormalizeHeaders(entry.Headers)
|
||||
entry.ExcludedModels = config.NormalizeExcludedModels(entry.ExcludedModels)
|
||||
if entry.BaseURL == "" {
|
||||
continue
|
||||
}
|
||||
@@ -557,6 +644,7 @@ func (h *Handler) PatchCodexKey(c *gin.Context) {
|
||||
value.BaseURL = strings.TrimSpace(value.BaseURL)
|
||||
value.ProxyURL = strings.TrimSpace(value.ProxyURL)
|
||||
value.Headers = config.NormalizeHeaders(value.Headers)
|
||||
value.ExcludedModels = config.NormalizeExcludedModels(value.ExcludedModels)
|
||||
// If base-url becomes empty, delete instead of update
|
||||
if value.BaseURL == "" {
|
||||
if body.Index != nil && *body.Index >= 0 && *body.Index < len(h.cfg.CodexKey) {
|
||||
@@ -694,6 +782,7 @@ func normalizeClaudeKey(entry *config.ClaudeKey) {
|
||||
entry.BaseURL = strings.TrimSpace(entry.BaseURL)
|
||||
entry.ProxyURL = strings.TrimSpace(entry.ProxyURL)
|
||||
entry.Headers = config.NormalizeHeaders(entry.Headers)
|
||||
entry.ExcludedModels = config.NormalizeExcludedModels(entry.ExcludedModels)
|
||||
if len(entry.Models) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
@@ -23,11 +23,13 @@ type Option func(*AmpModule)
|
||||
// - Reverse proxy to Amp control plane for OAuth/management
|
||||
// - Provider-specific route aliases (/api/provider/{provider}/...)
|
||||
// - Automatic gzip decompression for misconfigured upstreams
|
||||
// - Model mapping for routing unavailable models to alternatives
|
||||
type AmpModule struct {
|
||||
secretSource SecretSource
|
||||
proxy *httputil.ReverseProxy
|
||||
accessManager *sdkaccess.Manager
|
||||
authMiddleware_ gin.HandlerFunc
|
||||
modelMapper *DefaultModelMapper
|
||||
enabled bool
|
||||
registerOnce sync.Once
|
||||
}
|
||||
@@ -101,6 +103,9 @@ func (m *AmpModule) Register(ctx modules.Context) error {
|
||||
// Use registerOnce to ensure routes are only registered once
|
||||
var regErr error
|
||||
m.registerOnce.Do(func() {
|
||||
// Initialize model mapper from config (for routing unavailable models to alternatives)
|
||||
m.modelMapper = NewModelMapper(ctx.Config.AmpModelMappings)
|
||||
|
||||
// Always register provider aliases - these work without an upstream
|
||||
m.registerProviderAliases(ctx.Engine, ctx.BaseHandler, auth)
|
||||
|
||||
@@ -159,8 +164,16 @@ func (m *AmpModule) getAuthMiddleware(ctx modules.Context) gin.HandlerFunc {
|
||||
// OnConfigUpdated handles configuration updates.
|
||||
// Currently requires restart for URL changes (could be enhanced for dynamic updates).
|
||||
func (m *AmpModule) OnConfigUpdated(cfg *config.Config) error {
|
||||
// Update model mappings (hot-reload supported)
|
||||
if m.modelMapper != nil {
|
||||
log.Infof("amp config updated: reloading %d model mapping(s)", len(cfg.AmpModelMappings))
|
||||
m.modelMapper.UpdateMappings(cfg.AmpModelMappings)
|
||||
} else {
|
||||
log.Warnf("amp model mapper not initialized, skipping model mapping update")
|
||||
}
|
||||
|
||||
if !m.enabled {
|
||||
log.Debug("Amp routing not enabled, skipping config update")
|
||||
log.Debug("Amp routing not enabled, skipping other config updates")
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -181,3 +194,8 @@ func (m *AmpModule) OnConfigUpdated(cfg *config.Config) error {
|
||||
log.Debug("Amp config updated (restart required for URL changes)")
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetModelMapper returns the model mapper instance (for testing/debugging).
|
||||
func (m *AmpModule) GetModelMapper() *DefaultModelMapper {
|
||||
return m.modelMapper
|
||||
}
|
||||
|
||||
@@ -6,16 +6,75 @@ import (
|
||||
"io"
|
||||
"net/http/httputil"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
// AmpRouteType represents the type of routing decision made for an Amp request
|
||||
type AmpRouteType string
|
||||
|
||||
const (
|
||||
// RouteTypeLocalProvider indicates the request is handled by a local OAuth provider (free)
|
||||
RouteTypeLocalProvider AmpRouteType = "LOCAL_PROVIDER"
|
||||
// RouteTypeModelMapping indicates the request was remapped to another available model (free)
|
||||
RouteTypeModelMapping AmpRouteType = "MODEL_MAPPING"
|
||||
// RouteTypeAmpCredits indicates the request is forwarded to ampcode.com (uses Amp credits)
|
||||
RouteTypeAmpCredits AmpRouteType = "AMP_CREDITS"
|
||||
// RouteTypeNoProvider indicates no provider or fallback available
|
||||
RouteTypeNoProvider AmpRouteType = "NO_PROVIDER"
|
||||
)
|
||||
|
||||
// logAmpRouting logs the routing decision for an Amp request with structured fields
|
||||
func logAmpRouting(routeType AmpRouteType, requestedModel, resolvedModel, provider, path string) {
|
||||
fields := log.Fields{
|
||||
"component": "amp-routing",
|
||||
"route_type": string(routeType),
|
||||
"requested_model": requestedModel,
|
||||
"path": path,
|
||||
"timestamp": time.Now().Format(time.RFC3339),
|
||||
}
|
||||
|
||||
if resolvedModel != "" && resolvedModel != requestedModel {
|
||||
fields["resolved_model"] = resolvedModel
|
||||
}
|
||||
if provider != "" {
|
||||
fields["provider"] = provider
|
||||
}
|
||||
|
||||
switch routeType {
|
||||
case RouteTypeLocalProvider:
|
||||
fields["cost"] = "free"
|
||||
fields["source"] = "local_oauth"
|
||||
log.WithFields(fields).Infof("[AMP] Using local provider for model: %s", requestedModel)
|
||||
|
||||
case RouteTypeModelMapping:
|
||||
fields["cost"] = "free"
|
||||
fields["source"] = "local_oauth"
|
||||
fields["mapping"] = requestedModel + " -> " + resolvedModel
|
||||
log.WithFields(fields).Infof("[AMP] Model mapped: %s -> %s", requestedModel, resolvedModel)
|
||||
|
||||
case RouteTypeAmpCredits:
|
||||
fields["cost"] = "amp_credits"
|
||||
fields["source"] = "ampcode.com"
|
||||
fields["model_id"] = requestedModel // Explicit model_id for easy config reference
|
||||
log.WithFields(fields).Warnf("[AMP] Forwarding to ampcode.com (uses Amp credits) - model_id: %s | To use local proxy, add to config: amp-model-mappings: [{from: \"%s\", to: \"<your-local-model>\"}]", requestedModel, requestedModel)
|
||||
|
||||
case RouteTypeNoProvider:
|
||||
fields["cost"] = "none"
|
||||
fields["source"] = "error"
|
||||
fields["model_id"] = requestedModel // Explicit model_id for easy config reference
|
||||
log.WithFields(fields).Warnf("[AMP] No provider available for model_id: %s", requestedModel)
|
||||
}
|
||||
}
|
||||
|
||||
// FallbackHandler wraps a standard handler with fallback logic to ampcode.com
|
||||
// when the model's provider is not available in CLIProxyAPI
|
||||
type FallbackHandler struct {
|
||||
getProxy func() *httputil.ReverseProxy
|
||||
getProxy func() *httputil.ReverseProxy
|
||||
modelMapper ModelMapper
|
||||
}
|
||||
|
||||
// NewFallbackHandler creates a new fallback handler wrapper
|
||||
@@ -26,10 +85,25 @@ func NewFallbackHandler(getProxy func() *httputil.ReverseProxy) *FallbackHandler
|
||||
}
|
||||
}
|
||||
|
||||
// NewFallbackHandlerWithMapper creates a new fallback handler with model mapping support
|
||||
func NewFallbackHandlerWithMapper(getProxy func() *httputil.ReverseProxy, mapper ModelMapper) *FallbackHandler {
|
||||
return &FallbackHandler{
|
||||
getProxy: getProxy,
|
||||
modelMapper: mapper,
|
||||
}
|
||||
}
|
||||
|
||||
// SetModelMapper sets the model mapper for this handler (allows late binding)
|
||||
func (fh *FallbackHandler) SetModelMapper(mapper ModelMapper) {
|
||||
fh.modelMapper = mapper
|
||||
}
|
||||
|
||||
// WrapHandler wraps a gin.HandlerFunc with fallback logic
|
||||
// If the model's provider is not configured in CLIProxyAPI, it forwards to ampcode.com
|
||||
func (fh *FallbackHandler) WrapHandler(handler gin.HandlerFunc) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
requestPath := c.Request.URL.Path
|
||||
|
||||
// Read the request body to extract the model name
|
||||
bodyBytes, err := io.ReadAll(c.Request.Body)
|
||||
if err != nil {
|
||||
@@ -55,12 +129,33 @@ func (fh *FallbackHandler) WrapHandler(handler gin.HandlerFunc) gin.HandlerFunc
|
||||
// Check if we have providers for this model
|
||||
providers := util.GetProviderName(normalizedModel)
|
||||
|
||||
// Track resolved model for logging (may change if mapping is applied)
|
||||
resolvedModel := normalizedModel
|
||||
usedMapping := false
|
||||
|
||||
if len(providers) == 0 {
|
||||
// No providers configured - check if we have a proxy for fallback
|
||||
// No providers configured - check if we have a model mapping
|
||||
if fh.modelMapper != nil {
|
||||
if mappedModel := fh.modelMapper.MapModel(normalizedModel); mappedModel != "" {
|
||||
// Mapping found - rewrite the model in request body
|
||||
bodyBytes = rewriteModelInBody(bodyBytes, mappedModel)
|
||||
c.Request.Body = io.NopCloser(bytes.NewReader(bodyBytes))
|
||||
resolvedModel = mappedModel
|
||||
usedMapping = true
|
||||
|
||||
// Get providers for the mapped model
|
||||
providers = util.GetProviderName(mappedModel)
|
||||
|
||||
// Continue to handler with remapped model
|
||||
goto handleRequest
|
||||
}
|
||||
}
|
||||
|
||||
// No mapping found - check if we have a proxy for fallback
|
||||
proxy := fh.getProxy()
|
||||
if proxy != nil {
|
||||
// Fallback to ampcode.com
|
||||
log.Infof("amp fallback: model %s has no configured provider, forwarding to ampcode.com", modelName)
|
||||
// Log: Forwarding to ampcode.com (uses Amp credits)
|
||||
logAmpRouting(RouteTypeAmpCredits, modelName, "", "", requestPath)
|
||||
|
||||
// Restore body again for the proxy
|
||||
c.Request.Body = io.NopCloser(bytes.NewReader(bodyBytes))
|
||||
@@ -71,7 +166,23 @@ func (fh *FallbackHandler) WrapHandler(handler gin.HandlerFunc) gin.HandlerFunc
|
||||
}
|
||||
|
||||
// No proxy available, let the normal handler return the error
|
||||
log.Debugf("amp fallback: model %s has no configured provider and no proxy available", modelName)
|
||||
logAmpRouting(RouteTypeNoProvider, modelName, "", "", requestPath)
|
||||
}
|
||||
|
||||
handleRequest:
|
||||
|
||||
// Log the routing decision
|
||||
providerName := ""
|
||||
if len(providers) > 0 {
|
||||
providerName = providers[0]
|
||||
}
|
||||
|
||||
if usedMapping {
|
||||
// Log: Model was mapped to another model
|
||||
logAmpRouting(RouteTypeModelMapping, modelName, resolvedModel, providerName, requestPath)
|
||||
} else if len(providers) > 0 {
|
||||
// Log: Using local provider (free)
|
||||
logAmpRouting(RouteTypeLocalProvider, modelName, resolvedModel, providerName, requestPath)
|
||||
}
|
||||
|
||||
// Providers available or no proxy for fallback, restore body and use normal handler
|
||||
@@ -91,6 +202,27 @@ func (fh *FallbackHandler) WrapHandler(handler gin.HandlerFunc) gin.HandlerFunc
|
||||
}
|
||||
}
|
||||
|
||||
// rewriteModelInBody replaces the model name in a JSON request body
|
||||
func rewriteModelInBody(body []byte, newModel string) []byte {
|
||||
var payload map[string]interface{}
|
||||
if err := json.Unmarshal(body, &payload); err != nil {
|
||||
log.Warnf("amp model mapping: failed to parse body for rewrite: %v", err)
|
||||
return body
|
||||
}
|
||||
|
||||
if _, exists := payload["model"]; exists {
|
||||
payload["model"] = newModel
|
||||
newBody, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
log.Warnf("amp model mapping: failed to marshal rewritten body: %v", err)
|
||||
return body
|
||||
}
|
||||
return newBody
|
||||
}
|
||||
|
||||
return body
|
||||
}
|
||||
|
||||
// extractModelFromRequest attempts to extract the model name from various request formats
|
||||
func extractModelFromRequest(body []byte, c *gin.Context) string {
|
||||
// First try to parse from JSON body (OpenAI, Claude, etc.)
|
||||
|
||||
113
internal/api/modules/amp/model_mapping.go
Normal file
113
internal/api/modules/amp/model_mapping.go
Normal file
@@ -0,0 +1,113 @@
|
||||
// Package amp provides model mapping functionality for routing Amp CLI requests
|
||||
// to alternative models when the requested model is not available locally.
|
||||
package amp
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
// ModelMapper provides model name mapping/aliasing for Amp CLI requests.
|
||||
// When an Amp request comes in for a model that isn't available locally,
|
||||
// this mapper can redirect it to an alternative model that IS available.
|
||||
type ModelMapper interface {
|
||||
// MapModel returns the target model name if a mapping exists and the target
|
||||
// model has available providers. Returns empty string if no mapping applies.
|
||||
MapModel(requestedModel string) string
|
||||
|
||||
// UpdateMappings refreshes the mapping configuration (for hot-reload).
|
||||
UpdateMappings(mappings []config.AmpModelMapping)
|
||||
}
|
||||
|
||||
// DefaultModelMapper implements ModelMapper with thread-safe mapping storage.
|
||||
type DefaultModelMapper struct {
|
||||
mu sync.RWMutex
|
||||
mappings map[string]string // from -> to (normalized lowercase keys)
|
||||
}
|
||||
|
||||
// NewModelMapper creates a new model mapper with the given initial mappings.
|
||||
func NewModelMapper(mappings []config.AmpModelMapping) *DefaultModelMapper {
|
||||
m := &DefaultModelMapper{
|
||||
mappings: make(map[string]string),
|
||||
}
|
||||
m.UpdateMappings(mappings)
|
||||
return m
|
||||
}
|
||||
|
||||
// MapModel checks if a mapping exists for the requested model and if the
|
||||
// target model has available local providers. Returns the mapped model name
|
||||
// or empty string if no valid mapping exists.
|
||||
func (m *DefaultModelMapper) MapModel(requestedModel string) string {
|
||||
if requestedModel == "" {
|
||||
return ""
|
||||
}
|
||||
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
|
||||
// Normalize the requested model for lookup
|
||||
normalizedRequest := strings.ToLower(strings.TrimSpace(requestedModel))
|
||||
|
||||
// Check for direct mapping
|
||||
targetModel, exists := m.mappings[normalizedRequest]
|
||||
if !exists {
|
||||
return ""
|
||||
}
|
||||
|
||||
// Verify target model has available providers
|
||||
providers := util.GetProviderName(targetModel)
|
||||
if len(providers) == 0 {
|
||||
log.Debugf("amp model mapping: target model %s has no available providers, skipping mapping", targetModel)
|
||||
return ""
|
||||
}
|
||||
|
||||
// Note: Detailed routing log is handled by logAmpRouting in fallback_handlers.go
|
||||
log.Debugf("amp model mapping: resolved %s -> %s", requestedModel, targetModel)
|
||||
return targetModel
|
||||
}
|
||||
|
||||
// UpdateMappings refreshes the mapping configuration from config.
|
||||
// This is called during initialization and on config hot-reload.
|
||||
func (m *DefaultModelMapper) UpdateMappings(mappings []config.AmpModelMapping) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
// Clear and rebuild mappings
|
||||
m.mappings = make(map[string]string, len(mappings))
|
||||
|
||||
for _, mapping := range mappings {
|
||||
from := strings.TrimSpace(mapping.From)
|
||||
to := strings.TrimSpace(mapping.To)
|
||||
|
||||
if from == "" || to == "" {
|
||||
log.Warnf("amp model mapping: skipping invalid mapping (from=%q, to=%q)", from, to)
|
||||
continue
|
||||
}
|
||||
|
||||
// Store with normalized lowercase key for case-insensitive lookup
|
||||
normalizedFrom := strings.ToLower(from)
|
||||
m.mappings[normalizedFrom] = to
|
||||
|
||||
log.Debugf("amp model mapping registered: %s -> %s", from, to)
|
||||
}
|
||||
|
||||
if len(m.mappings) > 0 {
|
||||
log.Infof("amp model mapping: loaded %d mapping(s)", len(m.mappings))
|
||||
}
|
||||
}
|
||||
|
||||
// GetMappings returns a copy of current mappings (for debugging/status).
|
||||
func (m *DefaultModelMapper) GetMappings() map[string]string {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
|
||||
result := make(map[string]string, len(m.mappings))
|
||||
for k, v := range m.mappings {
|
||||
result[k] = v
|
||||
}
|
||||
return result
|
||||
}
|
||||
186
internal/api/modules/amp/model_mapping_test.go
Normal file
186
internal/api/modules/amp/model_mapping_test.go
Normal file
@@ -0,0 +1,186 @@
|
||||
package amp
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
|
||||
)
|
||||
|
||||
func TestNewModelMapper(t *testing.T) {
|
||||
mappings := []config.AmpModelMapping{
|
||||
{From: "claude-opus-4.5", To: "claude-sonnet-4"},
|
||||
{From: "gpt-5", To: "gemini-2.5-pro"},
|
||||
}
|
||||
|
||||
mapper := NewModelMapper(mappings)
|
||||
if mapper == nil {
|
||||
t.Fatal("Expected non-nil mapper")
|
||||
}
|
||||
|
||||
result := mapper.GetMappings()
|
||||
if len(result) != 2 {
|
||||
t.Errorf("Expected 2 mappings, got %d", len(result))
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewModelMapper_Empty(t *testing.T) {
|
||||
mapper := NewModelMapper(nil)
|
||||
if mapper == nil {
|
||||
t.Fatal("Expected non-nil mapper")
|
||||
}
|
||||
|
||||
result := mapper.GetMappings()
|
||||
if len(result) != 0 {
|
||||
t.Errorf("Expected 0 mappings, got %d", len(result))
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelMapper_MapModel_NoProvider(t *testing.T) {
|
||||
mappings := []config.AmpModelMapping{
|
||||
{From: "claude-opus-4.5", To: "claude-sonnet-4"},
|
||||
}
|
||||
|
||||
mapper := NewModelMapper(mappings)
|
||||
|
||||
// Without a registered provider for the target, mapping should return empty
|
||||
result := mapper.MapModel("claude-opus-4.5")
|
||||
if result != "" {
|
||||
t.Errorf("Expected empty result when target has no provider, got %s", result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelMapper_MapModel_WithProvider(t *testing.T) {
|
||||
// Register a mock provider for the target model
|
||||
reg := registry.GetGlobalRegistry()
|
||||
reg.RegisterClient("test-client", "claude", []*registry.ModelInfo{
|
||||
{ID: "claude-sonnet-4", OwnedBy: "anthropic", Type: "claude"},
|
||||
})
|
||||
defer reg.UnregisterClient("test-client")
|
||||
|
||||
mappings := []config.AmpModelMapping{
|
||||
{From: "claude-opus-4.5", To: "claude-sonnet-4"},
|
||||
}
|
||||
|
||||
mapper := NewModelMapper(mappings)
|
||||
|
||||
// With a registered provider, mapping should work
|
||||
result := mapper.MapModel("claude-opus-4.5")
|
||||
if result != "claude-sonnet-4" {
|
||||
t.Errorf("Expected claude-sonnet-4, got %s", result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelMapper_MapModel_CaseInsensitive(t *testing.T) {
|
||||
reg := registry.GetGlobalRegistry()
|
||||
reg.RegisterClient("test-client2", "claude", []*registry.ModelInfo{
|
||||
{ID: "claude-sonnet-4", OwnedBy: "anthropic", Type: "claude"},
|
||||
})
|
||||
defer reg.UnregisterClient("test-client2")
|
||||
|
||||
mappings := []config.AmpModelMapping{
|
||||
{From: "Claude-Opus-4.5", To: "claude-sonnet-4"},
|
||||
}
|
||||
|
||||
mapper := NewModelMapper(mappings)
|
||||
|
||||
// Should match case-insensitively
|
||||
result := mapper.MapModel("claude-opus-4.5")
|
||||
if result != "claude-sonnet-4" {
|
||||
t.Errorf("Expected claude-sonnet-4, got %s", result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelMapper_MapModel_NotFound(t *testing.T) {
|
||||
mappings := []config.AmpModelMapping{
|
||||
{From: "claude-opus-4.5", To: "claude-sonnet-4"},
|
||||
}
|
||||
|
||||
mapper := NewModelMapper(mappings)
|
||||
|
||||
// Unknown model should return empty
|
||||
result := mapper.MapModel("unknown-model")
|
||||
if result != "" {
|
||||
t.Errorf("Expected empty for unknown model, got %s", result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelMapper_MapModel_EmptyInput(t *testing.T) {
|
||||
mappings := []config.AmpModelMapping{
|
||||
{From: "claude-opus-4.5", To: "claude-sonnet-4"},
|
||||
}
|
||||
|
||||
mapper := NewModelMapper(mappings)
|
||||
|
||||
result := mapper.MapModel("")
|
||||
if result != "" {
|
||||
t.Errorf("Expected empty for empty input, got %s", result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelMapper_UpdateMappings(t *testing.T) {
|
||||
mapper := NewModelMapper(nil)
|
||||
|
||||
// Initially empty
|
||||
if len(mapper.GetMappings()) != 0 {
|
||||
t.Error("Expected 0 initial mappings")
|
||||
}
|
||||
|
||||
// Update with new mappings
|
||||
mapper.UpdateMappings([]config.AmpModelMapping{
|
||||
{From: "model-a", To: "model-b"},
|
||||
{From: "model-c", To: "model-d"},
|
||||
})
|
||||
|
||||
result := mapper.GetMappings()
|
||||
if len(result) != 2 {
|
||||
t.Errorf("Expected 2 mappings after update, got %d", len(result))
|
||||
}
|
||||
|
||||
// Update again should replace, not append
|
||||
mapper.UpdateMappings([]config.AmpModelMapping{
|
||||
{From: "model-x", To: "model-y"},
|
||||
})
|
||||
|
||||
result = mapper.GetMappings()
|
||||
if len(result) != 1 {
|
||||
t.Errorf("Expected 1 mapping after second update, got %d", len(result))
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelMapper_UpdateMappings_SkipsInvalid(t *testing.T) {
|
||||
mapper := NewModelMapper(nil)
|
||||
|
||||
mapper.UpdateMappings([]config.AmpModelMapping{
|
||||
{From: "", To: "model-b"}, // Invalid: empty from
|
||||
{From: "model-a", To: ""}, // Invalid: empty to
|
||||
{From: " ", To: "model-b"}, // Invalid: whitespace from
|
||||
{From: "model-c", To: "model-d"}, // Valid
|
||||
})
|
||||
|
||||
result := mapper.GetMappings()
|
||||
if len(result) != 1 {
|
||||
t.Errorf("Expected 1 valid mapping, got %d", len(result))
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelMapper_GetMappings_ReturnsCopy(t *testing.T) {
|
||||
mappings := []config.AmpModelMapping{
|
||||
{From: "model-a", To: "model-b"},
|
||||
}
|
||||
|
||||
mapper := NewModelMapper(mappings)
|
||||
|
||||
// Get mappings and modify the returned map
|
||||
result := mapper.GetMappings()
|
||||
result["new-key"] = "new-value"
|
||||
|
||||
// Original should be unchanged
|
||||
original := mapper.GetMappings()
|
||||
if len(original) != 1 {
|
||||
t.Errorf("Expected original to have 1 mapping, got %d", len(original))
|
||||
}
|
||||
if _, exists := original["new-key"]; exists {
|
||||
t.Error("Original map was modified")
|
||||
}
|
||||
}
|
||||
@@ -111,6 +111,14 @@ func (m *AmpModule) registerManagementRoutes(engine *gin.Engine, baseHandler *ha
|
||||
ampAPI.Any("/otel", proxyHandler)
|
||||
ampAPI.Any("/otel/*path", proxyHandler)
|
||||
|
||||
// Root-level routes that AMP CLI expects without /api prefix
|
||||
// These need the same security middleware as the /api/* routes
|
||||
rootMiddleware := []gin.HandlerFunc{noCORSMiddleware()}
|
||||
if restrictToLocalhost {
|
||||
rootMiddleware = append(rootMiddleware, localhostOnlyMiddleware())
|
||||
}
|
||||
engine.GET("/threads.rss", append(rootMiddleware, proxyHandler)...)
|
||||
|
||||
// Google v1beta1 passthrough with OAuth fallback
|
||||
// AMP CLI uses non-standard paths like /publishers/google/models/...
|
||||
// We bridge these to our standard Gemini handler to enable local OAuth.
|
||||
@@ -162,9 +170,10 @@ func (m *AmpModule) registerProviderAliases(engine *gin.Engine, baseHandler *han
|
||||
|
||||
// Create fallback handler wrapper that forwards to ampcode.com when provider not found
|
||||
// Uses lazy evaluation to access proxy (which is created after routes are registered)
|
||||
fallbackHandler := NewFallbackHandler(func() *httputil.ReverseProxy {
|
||||
// Also includes model mapping support for routing unavailable models to alternatives
|
||||
fallbackHandler := NewFallbackHandlerWithMapper(func() *httputil.ReverseProxy {
|
||||
return m.proxy
|
||||
})
|
||||
}, m.modelMapper)
|
||||
|
||||
// Provider-specific routes under /api/provider/:provider
|
||||
ampProviders := engine.Group("/api/provider")
|
||||
|
||||
@@ -37,6 +37,7 @@ func TestRegisterManagementRoutes(t *testing.T) {
|
||||
{"/api/meta", http.MethodGet},
|
||||
{"/api/telemetry", http.MethodGet},
|
||||
{"/api/threads", http.MethodGet},
|
||||
{"/threads.rss", http.MethodGet}, // Root-level route (no /api prefix)
|
||||
{"/api/otel", http.MethodGet},
|
||||
// Google v1beta1 bridge should still proxy non-model requests (GET) and allow POST
|
||||
{"/api/provider/google/v1beta1/models", http.MethodGet},
|
||||
|
||||
@@ -150,6 +150,9 @@ type Server struct {
|
||||
// management handler
|
||||
mgmt *managementHandlers.Handler
|
||||
|
||||
// ampModule is the Amp routing module for model mapping hot-reload
|
||||
ampModule *ampmodule.AmpModule
|
||||
|
||||
// managementRoutesRegistered tracks whether the management routes have been attached to the engine.
|
||||
managementRoutesRegistered atomic.Bool
|
||||
// managementRoutesEnabled controls whether management endpoints serve real handlers.
|
||||
@@ -268,14 +271,14 @@ func NewServer(cfg *config.Config, authManager *auth.Manager, accessManager *sdk
|
||||
s.setupRoutes()
|
||||
|
||||
// Register Amp module using V2 interface with Context
|
||||
ampModule := ampmodule.NewLegacy(accessManager, AuthMiddleware(accessManager))
|
||||
s.ampModule = ampmodule.NewLegacy(accessManager, AuthMiddleware(accessManager))
|
||||
ctx := modules.Context{
|
||||
Engine: engine,
|
||||
BaseHandler: s.handlers,
|
||||
Config: cfg,
|
||||
AuthMiddleware: AuthMiddleware(accessManager),
|
||||
}
|
||||
if err := modules.RegisterModule(ctx, ampModule); err != nil {
|
||||
if err := modules.RegisterModule(ctx, s.ampModule); err != nil {
|
||||
log.Errorf("Failed to register Amp module: %v", err)
|
||||
}
|
||||
|
||||
@@ -543,6 +546,11 @@ func (s *Server) registerManagementRoutes() {
|
||||
mgmt.PATCH("/openai-compatibility", s.mgmt.PatchOpenAICompat)
|
||||
mgmt.DELETE("/openai-compatibility", s.mgmt.DeleteOpenAICompat)
|
||||
|
||||
mgmt.GET("/oauth-excluded-models", s.mgmt.GetOAuthExcludedModels)
|
||||
mgmt.PUT("/oauth-excluded-models", s.mgmt.PutOAuthExcludedModels)
|
||||
mgmt.PATCH("/oauth-excluded-models", s.mgmt.PatchOAuthExcludedModels)
|
||||
mgmt.DELETE("/oauth-excluded-models", s.mgmt.DeleteOAuthExcludedModels)
|
||||
|
||||
mgmt.GET("/auth-files", s.mgmt.ListAuthFiles)
|
||||
mgmt.GET("/auth-files/download", s.mgmt.DownloadAuthFile)
|
||||
mgmt.POST("/auth-files", s.mgmt.UploadAuthFile)
|
||||
@@ -911,11 +919,22 @@ func (s *Server) UpdateClients(cfg *config.Config) {
|
||||
s.mgmt.SetAuthManager(s.handlers.AuthManager)
|
||||
}
|
||||
|
||||
// Notify Amp module of config changes (for model mapping hot-reload)
|
||||
if s.ampModule != nil {
|
||||
log.Debugf("triggering amp module config update")
|
||||
if err := s.ampModule.OnConfigUpdated(cfg); err != nil {
|
||||
log.Errorf("failed to update Amp module config: %v", err)
|
||||
}
|
||||
} else {
|
||||
log.Warnf("amp module is nil, skipping config update")
|
||||
}
|
||||
|
||||
// Count client sources from configuration and auth directory
|
||||
authFiles := util.CountAuthFiles(cfg.AuthDir)
|
||||
geminiAPIKeyCount := len(cfg.GeminiKey)
|
||||
claudeAPIKeyCount := len(cfg.ClaudeKey)
|
||||
codexAPIKeyCount := len(cfg.CodexKey)
|
||||
vertexAICompatCount := len(cfg.VertexCompatAPIKey)
|
||||
openAICompatCount := 0
|
||||
for i := range cfg.OpenAICompatibility {
|
||||
entry := cfg.OpenAICompatibility[i]
|
||||
@@ -926,13 +945,14 @@ func (s *Server) UpdateClients(cfg *config.Config) {
|
||||
openAICompatCount += len(entry.APIKeys)
|
||||
}
|
||||
|
||||
total := authFiles + geminiAPIKeyCount + claudeAPIKeyCount + codexAPIKeyCount + openAICompatCount
|
||||
fmt.Printf("server clients and configuration updated: %d clients (%d auth files + %d Gemini API keys + %d Claude API keys + %d Codex keys + %d OpenAI-compat)\n",
|
||||
total := authFiles + geminiAPIKeyCount + claudeAPIKeyCount + codexAPIKeyCount + vertexAICompatCount + openAICompatCount
|
||||
fmt.Printf("server clients and configuration updated: %d clients (%d auth files + %d Gemini API keys + %d Claude API keys + %d Codex keys + %d Vertex-compat + %d OpenAI-compat)\n",
|
||||
total,
|
||||
authFiles,
|
||||
geminiAPIKeyCount,
|
||||
claudeAPIKeyCount,
|
||||
codexAPIKeyCount,
|
||||
vertexAICompatCount,
|
||||
openAICompatCount,
|
||||
)
|
||||
}
|
||||
|
||||
@@ -37,6 +37,12 @@ type Config struct {
|
||||
// browser attacks and remote access to management endpoints. Default: true (recommended).
|
||||
AmpRestrictManagementToLocalhost bool `yaml:"amp-restrict-management-to-localhost" json:"amp-restrict-management-to-localhost"`
|
||||
|
||||
// AmpModelMappings defines model name mappings for Amp CLI requests.
|
||||
// When Amp requests a model that isn't available locally, these mappings
|
||||
// allow routing to an alternative model that IS available.
|
||||
// Example: Map "claude-opus-4.5" -> "claude-sonnet-4" when opus isn't available.
|
||||
AmpModelMappings []AmpModelMapping `yaml:"amp-model-mappings" json:"amp-model-mappings"`
|
||||
|
||||
// AuthDir is the directory where authentication token files are stored.
|
||||
AuthDir string `yaml:"auth-dir" json:"-"`
|
||||
|
||||
@@ -64,6 +70,10 @@ type Config struct {
|
||||
// GeminiKey defines Gemini API key configurations with optional routing overrides.
|
||||
GeminiKey []GeminiKey `yaml:"gemini-api-key" json:"gemini-api-key"`
|
||||
|
||||
// VertexCompatAPIKey defines Vertex AI-compatible API key configurations for third-party providers.
|
||||
// Used for services that use Vertex AI-style paths but with simple API key authentication.
|
||||
VertexCompatAPIKey []VertexCompatKey `yaml:"vertex-api-key" json:"vertex-api-key"`
|
||||
|
||||
// RequestRetry defines the retry times when the request failed.
|
||||
RequestRetry int `yaml:"request-retry" json:"request-retry"`
|
||||
// MaxRetryInterval defines the maximum wait time in seconds before retrying a cooled-down credential.
|
||||
@@ -83,6 +93,9 @@ type Config struct {
|
||||
|
||||
// Payload defines default and override rules for provider payload parameters.
|
||||
Payload PayloadConfig `yaml:"payload" json:"payload"`
|
||||
|
||||
// OAuthExcludedModels defines per-provider global model exclusions applied to OAuth/file-backed auth entries.
|
||||
OAuthExcludedModels map[string][]string `yaml:"oauth-excluded-models,omitempty" json:"oauth-excluded-models,omitempty"`
|
||||
}
|
||||
|
||||
// TLSConfig holds HTTPS server settings.
|
||||
@@ -115,6 +128,18 @@ type QuotaExceeded struct {
|
||||
SwitchPreviewModel bool `yaml:"switch-preview-model" json:"switch-preview-model"`
|
||||
}
|
||||
|
||||
// AmpModelMapping defines a model name mapping for Amp CLI requests.
|
||||
// When Amp requests a model that isn't available locally, this mapping
|
||||
// allows routing to an alternative model that IS available.
|
||||
type AmpModelMapping struct {
|
||||
// From is the model name that Amp CLI requests (e.g., "claude-opus-4.5").
|
||||
From string `yaml:"from" json:"from"`
|
||||
|
||||
// To is the target model name to route to (e.g., "claude-sonnet-4").
|
||||
// The target model must have available providers in the registry.
|
||||
To string `yaml:"to" json:"to"`
|
||||
}
|
||||
|
||||
// PayloadConfig defines default and override parameter rules applied to provider payloads.
|
||||
type PayloadConfig struct {
|
||||
// Default defines rules that only set parameters when they are missing in the payload.
|
||||
@@ -157,6 +182,9 @@ type ClaudeKey struct {
|
||||
|
||||
// Headers optionally adds extra HTTP headers for requests sent with this key.
|
||||
Headers map[string]string `yaml:"headers,omitempty" json:"headers,omitempty"`
|
||||
|
||||
// ExcludedModels lists model IDs that should be excluded for this provider.
|
||||
ExcludedModels []string `yaml:"excluded-models,omitempty" json:"excluded-models,omitempty"`
|
||||
}
|
||||
|
||||
// ClaudeModel describes a mapping between an alias and the actual upstream model name.
|
||||
@@ -183,6 +211,9 @@ type CodexKey struct {
|
||||
|
||||
// Headers optionally adds extra HTTP headers for requests sent with this key.
|
||||
Headers map[string]string `yaml:"headers,omitempty" json:"headers,omitempty"`
|
||||
|
||||
// ExcludedModels lists model IDs that should be excluded for this provider.
|
||||
ExcludedModels []string `yaml:"excluded-models,omitempty" json:"excluded-models,omitempty"`
|
||||
}
|
||||
|
||||
// GeminiKey represents the configuration for a Gemini API key,
|
||||
@@ -199,6 +230,9 @@ type GeminiKey struct {
|
||||
|
||||
// Headers optionally adds extra HTTP headers for requests sent with this key.
|
||||
Headers map[string]string `yaml:"headers,omitempty" json:"headers,omitempty"`
|
||||
|
||||
// ExcludedModels lists model IDs that should be excluded for this provider.
|
||||
ExcludedModels []string `yaml:"excluded-models,omitempty" json:"excluded-models,omitempty"`
|
||||
}
|
||||
|
||||
// OpenAICompatibility represents the configuration for OpenAI API compatibility
|
||||
@@ -313,6 +347,9 @@ func LoadConfigOptional(configFile string, optional bool) (*Config, error) {
|
||||
// Sanitize Gemini API key configuration and migrate legacy entries.
|
||||
cfg.SanitizeGeminiKeys()
|
||||
|
||||
// Sanitize Vertex-compatible API keys: drop entries without base-url
|
||||
cfg.SanitizeVertexCompatKeys()
|
||||
|
||||
// Sanitize Codex keys: drop entries without base-url
|
||||
cfg.SanitizeCodexKeys()
|
||||
|
||||
@@ -322,6 +359,9 @@ func LoadConfigOptional(configFile string, optional bool) (*Config, error) {
|
||||
// Sanitize OpenAI compatibility providers: drop entries without base-url
|
||||
cfg.SanitizeOpenAICompatibility()
|
||||
|
||||
// Normalize OAuth provider model exclusion map.
|
||||
cfg.OAuthExcludedModels = NormalizeOAuthExcludedModels(cfg.OAuthExcludedModels)
|
||||
|
||||
// Return the populated configuration struct.
|
||||
return &cfg, nil
|
||||
}
|
||||
@@ -359,6 +399,7 @@ func (cfg *Config) SanitizeCodexKeys() {
|
||||
e := cfg.CodexKey[i]
|
||||
e.BaseURL = strings.TrimSpace(e.BaseURL)
|
||||
e.Headers = NormalizeHeaders(e.Headers)
|
||||
e.ExcludedModels = NormalizeExcludedModels(e.ExcludedModels)
|
||||
if e.BaseURL == "" {
|
||||
continue
|
||||
}
|
||||
@@ -375,6 +416,7 @@ func (cfg *Config) SanitizeClaudeKeys() {
|
||||
for i := range cfg.ClaudeKey {
|
||||
entry := &cfg.ClaudeKey[i]
|
||||
entry.Headers = NormalizeHeaders(entry.Headers)
|
||||
entry.ExcludedModels = NormalizeExcludedModels(entry.ExcludedModels)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -395,6 +437,7 @@ func (cfg *Config) SanitizeGeminiKeys() {
|
||||
entry.BaseURL = strings.TrimSpace(entry.BaseURL)
|
||||
entry.ProxyURL = strings.TrimSpace(entry.ProxyURL)
|
||||
entry.Headers = NormalizeHeaders(entry.Headers)
|
||||
entry.ExcludedModels = NormalizeExcludedModels(entry.ExcludedModels)
|
||||
if _, exists := seen[entry.APIKey]; exists {
|
||||
continue
|
||||
}
|
||||
@@ -457,6 +500,55 @@ func NormalizeHeaders(headers map[string]string) map[string]string {
|
||||
return clean
|
||||
}
|
||||
|
||||
// NormalizeExcludedModels trims, lowercases, and deduplicates model exclusion patterns.
|
||||
// It preserves the order of first occurrences and drops empty entries.
|
||||
func NormalizeExcludedModels(models []string) []string {
|
||||
if len(models) == 0 {
|
||||
return nil
|
||||
}
|
||||
seen := make(map[string]struct{}, len(models))
|
||||
out := make([]string, 0, len(models))
|
||||
for _, raw := range models {
|
||||
trimmed := strings.ToLower(strings.TrimSpace(raw))
|
||||
if trimmed == "" {
|
||||
continue
|
||||
}
|
||||
if _, exists := seen[trimmed]; exists {
|
||||
continue
|
||||
}
|
||||
seen[trimmed] = struct{}{}
|
||||
out = append(out, trimmed)
|
||||
}
|
||||
if len(out) == 0 {
|
||||
return nil
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// NormalizeOAuthExcludedModels cleans provider -> excluded models mappings by normalizing provider keys
|
||||
// and applying model exclusion normalization to each entry.
|
||||
func NormalizeOAuthExcludedModels(entries map[string][]string) map[string][]string {
|
||||
if len(entries) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make(map[string][]string, len(entries))
|
||||
for provider, models := range entries {
|
||||
key := strings.ToLower(strings.TrimSpace(provider))
|
||||
if key == "" {
|
||||
continue
|
||||
}
|
||||
normalized := NormalizeExcludedModels(models)
|
||||
if len(normalized) == 0 {
|
||||
continue
|
||||
}
|
||||
out[key] = normalized
|
||||
}
|
||||
if len(out) == 0 {
|
||||
return nil
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// hashSecret hashes the given secret using bcrypt.
|
||||
func hashSecret(secret string) (string, error) {
|
||||
// Use default cost for simplicity.
|
||||
@@ -507,6 +599,7 @@ func SaveConfigPreserveComments(configFile string, cfg *Config) error {
|
||||
// Remove deprecated auth block before merging to avoid persisting it again.
|
||||
removeMapKey(original.Content[0], "auth")
|
||||
removeLegacyOpenAICompatAPIKeys(original.Content[0])
|
||||
pruneMappingToGeneratedKeys(original.Content[0], generated.Content[0], "oauth-excluded-models")
|
||||
|
||||
// Merge generated into original in-place, preserving comments/order of existing nodes.
|
||||
mergeMappingPreserve(original.Content[0], generated.Content[0])
|
||||
@@ -705,6 +798,10 @@ func mergeNodePreserve(dst, src *yaml.Node) {
|
||||
continue
|
||||
}
|
||||
mergeNodePreserve(dst.Content[i], src.Content[i])
|
||||
if dst.Content[i] != nil && src.Content[i] != nil &&
|
||||
dst.Content[i].Kind == yaml.MappingNode && src.Content[i].Kind == yaml.MappingNode {
|
||||
pruneMissingMapKeys(dst.Content[i], src.Content[i])
|
||||
}
|
||||
}
|
||||
// Append any extra items from src
|
||||
for i := len(dst.Content); i < len(src.Content); i++ {
|
||||
@@ -746,6 +843,7 @@ func shouldSkipEmptyCollectionOnPersist(key string, node *yaml.Node) bool {
|
||||
switch key {
|
||||
case "generative-language-api-key",
|
||||
"gemini-api-key",
|
||||
"vertex-api-key",
|
||||
"claude-api-key",
|
||||
"codex-api-key",
|
||||
"openai-compatibility":
|
||||
@@ -982,6 +1080,73 @@ func removeLegacyOpenAICompatAPIKeys(root *yaml.Node) {
|
||||
}
|
||||
}
|
||||
|
||||
func pruneMappingToGeneratedKeys(dstRoot, srcRoot *yaml.Node, key string) {
|
||||
if key == "" || dstRoot == nil || srcRoot == nil {
|
||||
return
|
||||
}
|
||||
if dstRoot.Kind != yaml.MappingNode || srcRoot.Kind != yaml.MappingNode {
|
||||
return
|
||||
}
|
||||
dstIdx := findMapKeyIndex(dstRoot, key)
|
||||
if dstIdx < 0 || dstIdx+1 >= len(dstRoot.Content) {
|
||||
return
|
||||
}
|
||||
srcIdx := findMapKeyIndex(srcRoot, key)
|
||||
if srcIdx < 0 {
|
||||
removeMapKey(dstRoot, key)
|
||||
return
|
||||
}
|
||||
if srcIdx+1 >= len(srcRoot.Content) {
|
||||
return
|
||||
}
|
||||
srcVal := srcRoot.Content[srcIdx+1]
|
||||
dstVal := dstRoot.Content[dstIdx+1]
|
||||
if srcVal == nil {
|
||||
dstRoot.Content[dstIdx+1] = nil
|
||||
return
|
||||
}
|
||||
if srcVal.Kind != yaml.MappingNode {
|
||||
dstRoot.Content[dstIdx+1] = deepCopyNode(srcVal)
|
||||
return
|
||||
}
|
||||
if dstVal == nil || dstVal.Kind != yaml.MappingNode {
|
||||
dstRoot.Content[dstIdx+1] = deepCopyNode(srcVal)
|
||||
return
|
||||
}
|
||||
pruneMissingMapKeys(dstVal, srcVal)
|
||||
}
|
||||
|
||||
func pruneMissingMapKeys(dstMap, srcMap *yaml.Node) {
|
||||
if dstMap == nil || srcMap == nil || dstMap.Kind != yaml.MappingNode || srcMap.Kind != yaml.MappingNode {
|
||||
return
|
||||
}
|
||||
keep := make(map[string]struct{}, len(srcMap.Content)/2)
|
||||
for i := 0; i+1 < len(srcMap.Content); i += 2 {
|
||||
keyNode := srcMap.Content[i]
|
||||
if keyNode == nil {
|
||||
continue
|
||||
}
|
||||
key := strings.TrimSpace(keyNode.Value)
|
||||
if key == "" {
|
||||
continue
|
||||
}
|
||||
keep[key] = struct{}{}
|
||||
}
|
||||
for i := 0; i+1 < len(dstMap.Content); {
|
||||
keyNode := dstMap.Content[i]
|
||||
if keyNode == nil {
|
||||
i += 2
|
||||
continue
|
||||
}
|
||||
key := strings.TrimSpace(keyNode.Value)
|
||||
if _, ok := keep[key]; !ok {
|
||||
dstMap.Content = append(dstMap.Content[:i], dstMap.Content[i+2:]...)
|
||||
continue
|
||||
}
|
||||
i += 2
|
||||
}
|
||||
}
|
||||
|
||||
// normalizeCollectionNodeStyles forces YAML collections to use block notation, keeping
|
||||
// lists and maps readable. Empty sequences retain flow style ([]) so empty list markers
|
||||
// remain compact.
|
||||
|
||||
84
internal/config/vertex_compat.go
Normal file
84
internal/config/vertex_compat.go
Normal file
@@ -0,0 +1,84 @@
|
||||
package config
|
||||
|
||||
import "strings"
|
||||
|
||||
// VertexCompatKey represents the configuration for Vertex AI-compatible API keys.
|
||||
// This supports third-party services that use Vertex AI-style endpoint paths
|
||||
// (/publishers/google/models/{model}:streamGenerateContent) but authenticate
|
||||
// with simple API keys instead of Google Cloud service account credentials.
|
||||
//
|
||||
// Example services: zenmux.ai and similar Vertex-compatible providers.
|
||||
type VertexCompatKey struct {
|
||||
// APIKey is the authentication key for accessing the Vertex-compatible API.
|
||||
// Maps to the x-goog-api-key header.
|
||||
APIKey string `yaml:"api-key" json:"api-key"`
|
||||
|
||||
// BaseURL is the base URL for the Vertex-compatible API endpoint.
|
||||
// The executor will append "/v1/publishers/google/models/{model}:action" to this.
|
||||
// Example: "https://zenmux.ai/api" becomes "https://zenmux.ai/api/v1/publishers/google/models/..."
|
||||
BaseURL string `yaml:"base-url,omitempty" json:"base-url,omitempty"`
|
||||
|
||||
// ProxyURL optionally overrides the global proxy for this API key.
|
||||
ProxyURL string `yaml:"proxy-url,omitempty" json:"proxy-url,omitempty"`
|
||||
|
||||
// Headers optionally adds extra HTTP headers for requests sent with this key.
|
||||
// Commonly used for cookies, user-agent, and other authentication headers.
|
||||
Headers map[string]string `yaml:"headers,omitempty" json:"headers,omitempty"`
|
||||
|
||||
// Models defines the model configurations including aliases for routing.
|
||||
Models []VertexCompatModel `yaml:"models,omitempty" json:"models,omitempty"`
|
||||
}
|
||||
|
||||
// VertexCompatModel represents a model configuration for Vertex compatibility,
|
||||
// including the actual model name and its alias for API routing.
|
||||
type VertexCompatModel struct {
|
||||
// Name is the actual model name used by the external provider.
|
||||
Name string `yaml:"name" json:"name"`
|
||||
|
||||
// Alias is the model name alias that clients will use to reference this model.
|
||||
Alias string `yaml:"alias" json:"alias"`
|
||||
}
|
||||
|
||||
// SanitizeVertexCompatKeys deduplicates and normalizes Vertex-compatible API key credentials.
|
||||
func (cfg *Config) SanitizeVertexCompatKeys() {
|
||||
if cfg == nil {
|
||||
return
|
||||
}
|
||||
|
||||
seen := make(map[string]struct{}, len(cfg.VertexCompatAPIKey))
|
||||
out := cfg.VertexCompatAPIKey[:0]
|
||||
for i := range cfg.VertexCompatAPIKey {
|
||||
entry := cfg.VertexCompatAPIKey[i]
|
||||
entry.APIKey = strings.TrimSpace(entry.APIKey)
|
||||
if entry.APIKey == "" {
|
||||
continue
|
||||
}
|
||||
entry.BaseURL = strings.TrimSpace(entry.BaseURL)
|
||||
if entry.BaseURL == "" {
|
||||
// BaseURL is required for Vertex API key entries
|
||||
continue
|
||||
}
|
||||
entry.ProxyURL = strings.TrimSpace(entry.ProxyURL)
|
||||
entry.Headers = NormalizeHeaders(entry.Headers)
|
||||
|
||||
// Sanitize models: remove entries without valid alias
|
||||
sanitizedModels := make([]VertexCompatModel, 0, len(entry.Models))
|
||||
for _, model := range entry.Models {
|
||||
model.Alias = strings.TrimSpace(model.Alias)
|
||||
model.Name = strings.TrimSpace(model.Name)
|
||||
if model.Alias != "" && model.Name != "" {
|
||||
sanitizedModels = append(sanitizedModels, model)
|
||||
}
|
||||
}
|
||||
entry.Models = sanitizedModels
|
||||
|
||||
// Use API key + base URL as uniqueness key
|
||||
uniqueKey := entry.APIKey + "|" + entry.BaseURL
|
||||
if _, exists := seen[uniqueKey]; exists {
|
||||
continue
|
||||
}
|
||||
seen[uniqueKey] = struct{}{}
|
||||
out = append(out, entry)
|
||||
}
|
||||
cfg.VertexCompatAPIKey = out
|
||||
}
|
||||
@@ -8,65 +8,79 @@ func GetClaudeModels() []*ModelInfo {
|
||||
return []*ModelInfo{
|
||||
|
||||
{
|
||||
ID: "claude-haiku-4-5-20251001",
|
||||
Object: "model",
|
||||
Created: 1759276800, // 2025-10-01
|
||||
OwnedBy: "anthropic",
|
||||
Type: "claude",
|
||||
DisplayName: "Claude 4.5 Haiku",
|
||||
ID: "claude-haiku-4-5-20251001",
|
||||
Object: "model",
|
||||
Created: 1759276800, // 2025-10-01
|
||||
OwnedBy: "anthropic",
|
||||
Type: "claude",
|
||||
DisplayName: "Claude 4.5 Haiku",
|
||||
ContextLength: 200000,
|
||||
MaxCompletionTokens: 64000,
|
||||
},
|
||||
{
|
||||
ID: "claude-sonnet-4-5-20250929",
|
||||
Object: "model",
|
||||
Created: 1759104000, // 2025-09-29
|
||||
OwnedBy: "anthropic",
|
||||
Type: "claude",
|
||||
DisplayName: "Claude 4.5 Sonnet",
|
||||
ID: "claude-sonnet-4-5-20250929",
|
||||
Object: "model",
|
||||
Created: 1759104000, // 2025-09-29
|
||||
OwnedBy: "anthropic",
|
||||
Type: "claude",
|
||||
DisplayName: "Claude 4.5 Sonnet",
|
||||
ContextLength: 200000,
|
||||
MaxCompletionTokens: 64000,
|
||||
},
|
||||
{
|
||||
ID: "claude-sonnet-4-5-thinking",
|
||||
Object: "model",
|
||||
Created: 1759104000, // 2025-09-29
|
||||
OwnedBy: "anthropic",
|
||||
Type: "claude",
|
||||
DisplayName: "Claude 4.5 Sonnet Thinking",
|
||||
Thinking: &ThinkingSupport{Min: 1024, Max: 100000, ZeroAllowed: false, DynamicAllowed: true},
|
||||
ID: "claude-sonnet-4-5-thinking",
|
||||
Object: "model",
|
||||
Created: 1759104000, // 2025-09-29
|
||||
OwnedBy: "anthropic",
|
||||
Type: "claude",
|
||||
DisplayName: "Claude 4.5 Sonnet Thinking",
|
||||
ContextLength: 200000,
|
||||
MaxCompletionTokens: 64000,
|
||||
Thinking: &ThinkingSupport{Min: 1024, Max: 100000, ZeroAllowed: false, DynamicAllowed: true},
|
||||
},
|
||||
{
|
||||
ID: "claude-opus-4-5-thinking",
|
||||
Object: "model",
|
||||
Created: 1761955200, // 2025-11-01
|
||||
OwnedBy: "anthropic",
|
||||
Type: "claude",
|
||||
DisplayName: "Claude 4.5 Opus Thinking",
|
||||
Thinking: &ThinkingSupport{Min: 1024, Max: 100000, ZeroAllowed: false, DynamicAllowed: true},
|
||||
ID: "claude-opus-4-5-thinking",
|
||||
Object: "model",
|
||||
Created: 1761955200, // 2025-11-01
|
||||
OwnedBy: "anthropic",
|
||||
Type: "claude",
|
||||
DisplayName: "Claude 4.5 Opus Thinking",
|
||||
ContextLength: 200000,
|
||||
MaxCompletionTokens: 64000,
|
||||
Thinking: &ThinkingSupport{Min: 1024, Max: 100000, ZeroAllowed: false, DynamicAllowed: true},
|
||||
},
|
||||
{
|
||||
ID: "claude-opus-4-5-thinking-low",
|
||||
Object: "model",
|
||||
Created: 1761955200, // 2025-11-01
|
||||
OwnedBy: "anthropic",
|
||||
Type: "claude",
|
||||
DisplayName: "Claude 4.5 Opus Thinking Low",
|
||||
Thinking: &ThinkingSupport{Min: 1024, Max: 100000, ZeroAllowed: false, DynamicAllowed: true},
|
||||
ID: "claude-opus-4-5-thinking-low",
|
||||
Object: "model",
|
||||
Created: 1761955200, // 2025-11-01
|
||||
OwnedBy: "anthropic",
|
||||
Type: "claude",
|
||||
DisplayName: "Claude 4.5 Opus Thinking Low",
|
||||
ContextLength: 200000,
|
||||
MaxCompletionTokens: 64000,
|
||||
Thinking: &ThinkingSupport{Min: 1024, Max: 100000, ZeroAllowed: false, DynamicAllowed: true},
|
||||
},
|
||||
{
|
||||
ID: "claude-opus-4-5-thinking-medium",
|
||||
Object: "model",
|
||||
Created: 1761955200, // 2025-11-01
|
||||
OwnedBy: "anthropic",
|
||||
Type: "claude",
|
||||
DisplayName: "Claude 4.5 Opus Thinking Medium",
|
||||
Thinking: &ThinkingSupport{Min: 1024, Max: 100000, ZeroAllowed: false, DynamicAllowed: true},
|
||||
ID: "claude-opus-4-5-thinking-medium",
|
||||
Object: "model",
|
||||
Created: 1761955200, // 2025-11-01
|
||||
OwnedBy: "anthropic",
|
||||
Type: "claude",
|
||||
DisplayName: "Claude 4.5 Opus Thinking Medium",
|
||||
ContextLength: 200000,
|
||||
MaxCompletionTokens: 64000,
|
||||
Thinking: &ThinkingSupport{Min: 1024, Max: 100000, ZeroAllowed: false, DynamicAllowed: true},
|
||||
},
|
||||
{
|
||||
ID: "claude-opus-4-5-thinking-high",
|
||||
Object: "model",
|
||||
Created: 1761955200, // 2025-11-01
|
||||
OwnedBy: "anthropic",
|
||||
Type: "claude",
|
||||
DisplayName: "Claude 4.5 Opus Thinking High",
|
||||
Thinking: &ThinkingSupport{Min: 1024, Max: 100000, ZeroAllowed: false, DynamicAllowed: true},
|
||||
ID: "claude-opus-4-5-thinking-high",
|
||||
Object: "model",
|
||||
Created: 1761955200, // 2025-11-01
|
||||
OwnedBy: "anthropic",
|
||||
Type: "claude",
|
||||
DisplayName: "Claude 4.5 Opus Thinking High",
|
||||
ContextLength: 200000,
|
||||
MaxCompletionTokens: 64000,
|
||||
Thinking: &ThinkingSupport{Min: 1024, Max: 100000, ZeroAllowed: false, DynamicAllowed: true},
|
||||
},
|
||||
{
|
||||
ID: "claude-opus-4-5-20251101",
|
||||
@@ -80,44 +94,54 @@ func GetClaudeModels() []*ModelInfo {
|
||||
MaxCompletionTokens: 64000,
|
||||
},
|
||||
{
|
||||
ID: "claude-opus-4-1-20250805",
|
||||
Object: "model",
|
||||
Created: 1722945600, // 2025-08-05
|
||||
OwnedBy: "anthropic",
|
||||
Type: "claude",
|
||||
DisplayName: "Claude 4.1 Opus",
|
||||
ID: "claude-opus-4-1-20250805",
|
||||
Object: "model",
|
||||
Created: 1722945600, // 2025-08-05
|
||||
OwnedBy: "anthropic",
|
||||
Type: "claude",
|
||||
DisplayName: "Claude 4.1 Opus",
|
||||
ContextLength: 200000,
|
||||
MaxCompletionTokens: 32000,
|
||||
},
|
||||
{
|
||||
ID: "claude-opus-4-20250514",
|
||||
Object: "model",
|
||||
Created: 1715644800, // 2025-05-14
|
||||
OwnedBy: "anthropic",
|
||||
Type: "claude",
|
||||
DisplayName: "Claude 4 Opus",
|
||||
ID: "claude-opus-4-20250514",
|
||||
Object: "model",
|
||||
Created: 1715644800, // 2025-05-14
|
||||
OwnedBy: "anthropic",
|
||||
Type: "claude",
|
||||
DisplayName: "Claude 4 Opus",
|
||||
ContextLength: 200000,
|
||||
MaxCompletionTokens: 32000,
|
||||
},
|
||||
{
|
||||
ID: "claude-sonnet-4-20250514",
|
||||
Object: "model",
|
||||
Created: 1715644800, // 2025-05-14
|
||||
OwnedBy: "anthropic",
|
||||
Type: "claude",
|
||||
DisplayName: "Claude 4 Sonnet",
|
||||
ID: "claude-sonnet-4-20250514",
|
||||
Object: "model",
|
||||
Created: 1715644800, // 2025-05-14
|
||||
OwnedBy: "anthropic",
|
||||
Type: "claude",
|
||||
DisplayName: "Claude 4 Sonnet",
|
||||
ContextLength: 200000,
|
||||
MaxCompletionTokens: 64000,
|
||||
},
|
||||
{
|
||||
ID: "claude-3-7-sonnet-20250219",
|
||||
Object: "model",
|
||||
Created: 1708300800, // 2025-02-19
|
||||
OwnedBy: "anthropic",
|
||||
Type: "claude",
|
||||
DisplayName: "Claude 3.7 Sonnet",
|
||||
ID: "claude-3-7-sonnet-20250219",
|
||||
Object: "model",
|
||||
Created: 1708300800, // 2025-02-19
|
||||
OwnedBy: "anthropic",
|
||||
Type: "claude",
|
||||
DisplayName: "Claude 3.7 Sonnet",
|
||||
ContextLength: 128000,
|
||||
MaxCompletionTokens: 8192,
|
||||
},
|
||||
{
|
||||
ID: "claude-3-5-haiku-20241022",
|
||||
Object: "model",
|
||||
Created: 1729555200, // 2024-10-22
|
||||
OwnedBy: "anthropic",
|
||||
Type: "claude",
|
||||
DisplayName: "Claude 3.5 Haiku",
|
||||
ID: "claude-3-5-haiku-20241022",
|
||||
Object: "model",
|
||||
Created: 1729555200, // 2024-10-22
|
||||
OwnedBy: "anthropic",
|
||||
Type: "claude",
|
||||
DisplayName: "Claude 3.5 Haiku",
|
||||
ContextLength: 128000,
|
||||
MaxCompletionTokens: 8192,
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -198,6 +222,7 @@ func GetGeminiModels() []*ModelInfo {
|
||||
InputTokenLimit: 1048576,
|
||||
OutputTokenLimit: 65536,
|
||||
SupportedGenerationMethods: []string{"generateContent", "countTokens", "createCachedContent", "batchGenerateContent"},
|
||||
Thinking: &ThinkingSupport{Min: 128, Max: 32768, ZeroAllowed: false, DynamicAllowed: true},
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -277,6 +302,7 @@ func GetGeminiVertexModels() []*ModelInfo {
|
||||
InputTokenLimit: 1048576,
|
||||
OutputTokenLimit: 65536,
|
||||
SupportedGenerationMethods: []string{"generateContent", "countTokens", "createCachedContent", "batchGenerateContent"},
|
||||
Thinking: &ThinkingSupport{Min: 128, Max: 32768, ZeroAllowed: false, DynamicAllowed: true},
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -928,7 +954,6 @@ func GetIFlowModels() []*ModelInfo {
|
||||
}{
|
||||
{ID: "tstars2.0", DisplayName: "TStars-2.0", Description: "iFlow TStars-2.0 multimodal assistant", Created: 1746489600},
|
||||
{ID: "qwen3-coder-plus", DisplayName: "Qwen3-Coder-Plus", Description: "Qwen3 Coder Plus code generation", Created: 1753228800},
|
||||
{ID: "qwen3-coder", DisplayName: "Qwen3-Coder-480B-A35B", Description: "Qwen3 Coder 480B A35B", Created: 1753228800},
|
||||
{ID: "qwen3-max", DisplayName: "Qwen3-Max", Description: "Qwen3 flagship model", Created: 1758672000},
|
||||
{ID: "qwen3-vl-plus", DisplayName: "Qwen3-VL-Plus", Description: "Qwen3 multimodal vision-language", Created: 1758672000},
|
||||
{ID: "qwen3-max-preview", DisplayName: "Qwen3-Max-Preview", Description: "Qwen3 Max preview build", Created: 1757030400},
|
||||
|
||||
@@ -308,13 +308,7 @@ func (e *AIStudioExecutor) translateRequest(req cliproxyexecutor.Request, opts c
|
||||
from := opts.SourceFormat
|
||||
to := sdktranslator.FromString("gemini")
|
||||
payload := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), stream)
|
||||
if budgetOverride, includeOverride, ok := util.GeminiThinkingFromMetadata(req.Metadata); ok && util.ModelSupportsThinking(req.Model) {
|
||||
if budgetOverride != nil {
|
||||
norm := util.NormalizeThinkingBudget(req.Model, *budgetOverride)
|
||||
budgetOverride = &norm
|
||||
}
|
||||
payload = util.ApplyGeminiThinkingConfig(payload, budgetOverride, includeOverride)
|
||||
}
|
||||
payload = applyThinkingMetadata(payload, req.Metadata, req.Model)
|
||||
payload = util.ConvertThinkingLevelToBudget(payload)
|
||||
payload = util.StripThinkingConfigIfUnsupported(req.Model, payload)
|
||||
payload = fixGeminiImageAspectRatio(req.Model, payload)
|
||||
|
||||
@@ -75,6 +75,8 @@ func (e *AntigravityExecutor) Execute(ctx context.Context, auth *cliproxyauth.Au
|
||||
to := sdktranslator.FromString("antigravity")
|
||||
translated := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)
|
||||
|
||||
translated = applyThinkingMetadataCLI(translated, req.Metadata, req.Model)
|
||||
|
||||
baseURLs := antigravityBaseURLFallbackOrder(auth)
|
||||
httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0)
|
||||
|
||||
@@ -166,6 +168,8 @@ func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxya
|
||||
to := sdktranslator.FromString("antigravity")
|
||||
translated := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), true)
|
||||
|
||||
translated = applyThinkingMetadataCLI(translated, req.Metadata, req.Model)
|
||||
|
||||
baseURLs := antigravityBaseURLFallbackOrder(auth)
|
||||
httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0)
|
||||
|
||||
|
||||
@@ -17,6 +17,7 @@ import (
|
||||
claudeauth "github.com/router-for-me/CLIProxyAPI/v6/internal/auth/claude"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/misc"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
|
||||
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
|
||||
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
|
||||
@@ -66,6 +67,9 @@ func (e *ClaudeExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r
|
||||
}
|
||||
body = applyPayloadConfig(e.cfg, req.Model, body)
|
||||
|
||||
// Ensure max_tokens > thinking.budget_tokens when thinking is enabled
|
||||
body = ensureMaxTokensForThinking(req.Model, body)
|
||||
|
||||
// Extract betas from body and convert to header
|
||||
var extraBetas []string
|
||||
extraBetas, body = extractAndRemoveBetas(body)
|
||||
@@ -165,6 +169,9 @@ func (e *ClaudeExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A
|
||||
body = checkSystemInstructions(body)
|
||||
body = applyPayloadConfig(e.cfg, req.Model, body)
|
||||
|
||||
// Ensure max_tokens > thinking.budget_tokens when thinking is enabled
|
||||
body = ensureMaxTokensForThinking(req.Model, body)
|
||||
|
||||
// Extract betas from body and convert to header
|
||||
var extraBetas []string
|
||||
extraBetas, body = extractAndRemoveBetas(body)
|
||||
@@ -447,6 +454,42 @@ func (e *ClaudeExecutor) injectThinkingConfig(modelName string, body []byte) []b
|
||||
return body
|
||||
}
|
||||
|
||||
// ensureMaxTokensForThinking ensures max_tokens > thinking.budget_tokens when thinking is enabled.
|
||||
// Anthropic API requires this constraint; violating it returns a 400 error.
|
||||
// This function should be called after all thinking configuration is finalized.
|
||||
// It looks up the model's MaxCompletionTokens from the registry to use as the cap.
|
||||
func ensureMaxTokensForThinking(modelName string, body []byte) []byte {
|
||||
thinkingType := gjson.GetBytes(body, "thinking.type").String()
|
||||
if thinkingType != "enabled" {
|
||||
return body
|
||||
}
|
||||
|
||||
budgetTokens := gjson.GetBytes(body, "thinking.budget_tokens").Int()
|
||||
if budgetTokens <= 0 {
|
||||
return body
|
||||
}
|
||||
|
||||
maxTokens := gjson.GetBytes(body, "max_tokens").Int()
|
||||
|
||||
// Look up the model's max completion tokens from the registry
|
||||
maxCompletionTokens := 0
|
||||
if modelInfo := registry.GetGlobalRegistry().GetModelInfo(modelName); modelInfo != nil {
|
||||
maxCompletionTokens = modelInfo.MaxCompletionTokens
|
||||
}
|
||||
|
||||
// Fall back to budget + buffer if registry lookup fails or returns 0
|
||||
const fallbackBuffer = 4000
|
||||
requiredMaxTokens := budgetTokens + fallbackBuffer
|
||||
if maxCompletionTokens > 0 {
|
||||
requiredMaxTokens = int64(maxCompletionTokens)
|
||||
}
|
||||
|
||||
if maxTokens < requiredMaxTokens {
|
||||
body, _ = sjson.SetBytes(body, "max_tokens", requiredMaxTokens)
|
||||
}
|
||||
return body
|
||||
}
|
||||
|
||||
func (e *ClaudeExecutor) resolveUpstreamModel(alias string, auth *cliproxyauth.Auth) string {
|
||||
if alias == "" {
|
||||
return ""
|
||||
|
||||
@@ -62,15 +62,8 @@ func (e *GeminiCLIExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth
|
||||
|
||||
from := opts.SourceFormat
|
||||
to := sdktranslator.FromString("gemini-cli")
|
||||
budgetOverride, includeOverride, hasOverride := util.GeminiThinkingFromMetadata(req.Metadata)
|
||||
basePayload := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)
|
||||
if hasOverride && util.ModelSupportsThinking(req.Model) {
|
||||
if budgetOverride != nil {
|
||||
norm := util.NormalizeThinkingBudget(req.Model, *budgetOverride)
|
||||
budgetOverride = &norm
|
||||
}
|
||||
basePayload = util.ApplyGeminiCLIThinkingConfig(basePayload, budgetOverride, includeOverride)
|
||||
}
|
||||
basePayload = applyThinkingMetadataCLI(basePayload, req.Metadata, req.Model)
|
||||
basePayload = util.StripThinkingConfigIfUnsupported(req.Model, basePayload)
|
||||
basePayload = fixGeminiCLIImageAspectRatio(req.Model, basePayload)
|
||||
basePayload = applyPayloadConfigWithRoot(e.cfg, req.Model, "gemini", "request", basePayload)
|
||||
@@ -204,15 +197,8 @@ func (e *GeminiCLIExecutor) ExecuteStream(ctx context.Context, auth *cliproxyaut
|
||||
|
||||
from := opts.SourceFormat
|
||||
to := sdktranslator.FromString("gemini-cli")
|
||||
budgetOverride, includeOverride, hasOverride := util.GeminiThinkingFromMetadata(req.Metadata)
|
||||
basePayload := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), true)
|
||||
if hasOverride && util.ModelSupportsThinking(req.Model) {
|
||||
if budgetOverride != nil {
|
||||
norm := util.NormalizeThinkingBudget(req.Model, *budgetOverride)
|
||||
budgetOverride = &norm
|
||||
}
|
||||
basePayload = util.ApplyGeminiCLIThinkingConfig(basePayload, budgetOverride, includeOverride)
|
||||
}
|
||||
basePayload = applyThinkingMetadataCLI(basePayload, req.Metadata, req.Model)
|
||||
basePayload = util.StripThinkingConfigIfUnsupported(req.Model, basePayload)
|
||||
basePayload = fixGeminiCLIImageAspectRatio(req.Model, basePayload)
|
||||
basePayload = applyPayloadConfigWithRoot(e.cfg, req.Model, "gemini", "request", basePayload)
|
||||
@@ -408,16 +394,9 @@ func (e *GeminiCLIExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.
|
||||
var lastStatus int
|
||||
var lastBody []byte
|
||||
|
||||
budgetOverride, includeOverride, hasOverride := util.GeminiThinkingFromMetadata(req.Metadata)
|
||||
for _, attemptModel := range models {
|
||||
payload := sdktranslator.TranslateRequest(from, to, attemptModel, bytes.Clone(req.Payload), false)
|
||||
if hasOverride && util.ModelSupportsThinking(req.Model) {
|
||||
if budgetOverride != nil {
|
||||
norm := util.NormalizeThinkingBudget(req.Model, *budgetOverride)
|
||||
budgetOverride = &norm
|
||||
}
|
||||
payload = util.ApplyGeminiCLIThinkingConfig(payload, budgetOverride, includeOverride)
|
||||
}
|
||||
payload = applyThinkingMetadataCLI(payload, req.Metadata, req.Model)
|
||||
payload = deleteJSONField(payload, "project")
|
||||
payload = deleteJSONField(payload, "model")
|
||||
payload = deleteJSONField(payload, "request.safetySettings")
|
||||
|
||||
@@ -79,13 +79,7 @@ func (e *GeminiExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, r
|
||||
from := opts.SourceFormat
|
||||
to := sdktranslator.FromString("gemini")
|
||||
body := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)
|
||||
if budgetOverride, includeOverride, ok := util.GeminiThinkingFromMetadata(req.Metadata); ok && util.ModelSupportsThinking(req.Model) {
|
||||
if budgetOverride != nil {
|
||||
norm := util.NormalizeThinkingBudget(req.Model, *budgetOverride)
|
||||
budgetOverride = &norm
|
||||
}
|
||||
body = util.ApplyGeminiThinkingConfig(body, budgetOverride, includeOverride)
|
||||
}
|
||||
body = applyThinkingMetadata(body, req.Metadata, req.Model)
|
||||
body = util.StripThinkingConfigIfUnsupported(req.Model, body)
|
||||
body = fixGeminiImageAspectRatio(req.Model, body)
|
||||
body = applyPayloadConfig(e.cfg, req.Model, body)
|
||||
@@ -174,13 +168,7 @@ func (e *GeminiExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.A
|
||||
from := opts.SourceFormat
|
||||
to := sdktranslator.FromString("gemini")
|
||||
body := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), true)
|
||||
if budgetOverride, includeOverride, ok := util.GeminiThinkingFromMetadata(req.Metadata); ok && util.ModelSupportsThinking(req.Model) {
|
||||
if budgetOverride != nil {
|
||||
norm := util.NormalizeThinkingBudget(req.Model, *budgetOverride)
|
||||
budgetOverride = &norm
|
||||
}
|
||||
body = util.ApplyGeminiThinkingConfig(body, budgetOverride, includeOverride)
|
||||
}
|
||||
body = applyThinkingMetadata(body, req.Metadata, req.Model)
|
||||
body = util.StripThinkingConfigIfUnsupported(req.Model, body)
|
||||
body = fixGeminiImageAspectRatio(req.Model, body)
|
||||
body = applyPayloadConfig(e.cfg, req.Model, body)
|
||||
@@ -288,13 +276,7 @@ func (e *GeminiExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Aut
|
||||
from := opts.SourceFormat
|
||||
to := sdktranslator.FromString("gemini")
|
||||
translatedReq := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)
|
||||
if budgetOverride, includeOverride, ok := util.GeminiThinkingFromMetadata(req.Metadata); ok && util.ModelSupportsThinking(req.Model) {
|
||||
if budgetOverride != nil {
|
||||
norm := util.NormalizeThinkingBudget(req.Model, *budgetOverride)
|
||||
budgetOverride = &norm
|
||||
}
|
||||
translatedReq = util.ApplyGeminiThinkingConfig(translatedReq, budgetOverride, includeOverride)
|
||||
}
|
||||
translatedReq = applyThinkingMetadata(translatedReq, req.Metadata, req.Model)
|
||||
translatedReq = util.StripThinkingConfigIfUnsupported(req.Model, translatedReq)
|
||||
translatedReq = fixGeminiImageAspectRatio(req.Model, translatedReq)
|
||||
respCtx := context.WithValue(ctx, "alt", opts.Alt)
|
||||
|
||||
@@ -51,11 +51,238 @@ func (e *GeminiVertexExecutor) PrepareRequest(_ *http.Request, _ *cliproxyauth.A
|
||||
|
||||
// Execute handles non-streaming requests.
|
||||
func (e *GeminiVertexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (resp cliproxyexecutor.Response, err error) {
|
||||
projectID, location, saJSON, errCreds := vertexCreds(auth)
|
||||
if errCreds != nil {
|
||||
return resp, errCreds
|
||||
// Try API key authentication first
|
||||
apiKey, baseURL := vertexAPICreds(auth)
|
||||
|
||||
// If no API key found, fall back to service account authentication
|
||||
if apiKey == "" {
|
||||
projectID, location, saJSON, errCreds := vertexCreds(auth)
|
||||
if errCreds != nil {
|
||||
return resp, errCreds
|
||||
}
|
||||
return e.executeWithServiceAccount(ctx, auth, req, opts, projectID, location, saJSON)
|
||||
}
|
||||
|
||||
// Use API key authentication
|
||||
return e.executeWithAPIKey(ctx, auth, req, opts, apiKey, baseURL)
|
||||
}
|
||||
|
||||
// ExecuteStream handles SSE streaming for Vertex.
|
||||
func (e *GeminiVertexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (stream <-chan cliproxyexecutor.StreamChunk, err error) {
|
||||
// Try API key authentication first
|
||||
apiKey, baseURL := vertexAPICreds(auth)
|
||||
|
||||
// If no API key found, fall back to service account authentication
|
||||
if apiKey == "" {
|
||||
projectID, location, saJSON, errCreds := vertexCreds(auth)
|
||||
if errCreds != nil {
|
||||
return nil, errCreds
|
||||
}
|
||||
return e.executeStreamWithServiceAccount(ctx, auth, req, opts, projectID, location, saJSON)
|
||||
}
|
||||
|
||||
// Use API key authentication
|
||||
return e.executeStreamWithAPIKey(ctx, auth, req, opts, apiKey, baseURL)
|
||||
}
|
||||
|
||||
// CountTokens calls Vertex countTokens endpoint.
|
||||
func (e *GeminiVertexExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
||||
// Try API key authentication first
|
||||
apiKey, baseURL := vertexAPICreds(auth)
|
||||
|
||||
// If no API key found, fall back to service account authentication
|
||||
if apiKey == "" {
|
||||
projectID, location, saJSON, errCreds := vertexCreds(auth)
|
||||
if errCreds != nil {
|
||||
return cliproxyexecutor.Response{}, errCreds
|
||||
}
|
||||
return e.countTokensWithServiceAccount(ctx, auth, req, opts, projectID, location, saJSON)
|
||||
}
|
||||
|
||||
// Use API key authentication
|
||||
return e.countTokensWithAPIKey(ctx, auth, req, opts, apiKey, baseURL)
|
||||
}
|
||||
|
||||
// countTokensWithServiceAccount handles token counting using service account credentials.
|
||||
func (e *GeminiVertexExecutor) countTokensWithServiceAccount(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, projectID, location string, saJSON []byte) (cliproxyexecutor.Response, error) {
|
||||
from := opts.SourceFormat
|
||||
to := sdktranslator.FromString("gemini")
|
||||
translatedReq := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)
|
||||
if budgetOverride, includeOverride, ok := util.GeminiThinkingFromMetadata(req.Metadata); ok && util.ModelSupportsThinking(req.Model) {
|
||||
if budgetOverride != nil {
|
||||
norm := util.NormalizeThinkingBudget(req.Model, *budgetOverride)
|
||||
budgetOverride = &norm
|
||||
}
|
||||
translatedReq = util.ApplyGeminiThinkingConfig(translatedReq, budgetOverride, includeOverride)
|
||||
}
|
||||
translatedReq = util.StripThinkingConfigIfUnsupported(req.Model, translatedReq)
|
||||
translatedReq = fixGeminiImageAspectRatio(req.Model, translatedReq)
|
||||
respCtx := context.WithValue(ctx, "alt", opts.Alt)
|
||||
translatedReq, _ = sjson.DeleteBytes(translatedReq, "tools")
|
||||
translatedReq, _ = sjson.DeleteBytes(translatedReq, "generationConfig")
|
||||
translatedReq, _ = sjson.DeleteBytes(translatedReq, "safetySettings")
|
||||
|
||||
baseURL := vertexBaseURL(location)
|
||||
url := fmt.Sprintf("%s/%s/projects/%s/locations/%s/publishers/google/models/%s:%s", baseURL, vertexAPIVersion, projectID, location, req.Model, "countTokens")
|
||||
|
||||
httpReq, errNewReq := http.NewRequestWithContext(respCtx, http.MethodPost, url, bytes.NewReader(translatedReq))
|
||||
if errNewReq != nil {
|
||||
return cliproxyexecutor.Response{}, errNewReq
|
||||
}
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
if token, errTok := vertexAccessToken(ctx, e.cfg, auth, saJSON); errTok == nil && token != "" {
|
||||
httpReq.Header.Set("Authorization", "Bearer "+token)
|
||||
} else if errTok != nil {
|
||||
log.Errorf("vertex executor: access token error: %v", errTok)
|
||||
return cliproxyexecutor.Response{}, statusErr{code: 500, msg: "internal server error"}
|
||||
}
|
||||
applyGeminiHeaders(httpReq, auth)
|
||||
|
||||
var authID, authLabel, authType, authValue string
|
||||
if auth != nil {
|
||||
authID = auth.ID
|
||||
authLabel = auth.Label
|
||||
authType, authValue = auth.AccountInfo()
|
||||
}
|
||||
recordAPIRequest(ctx, e.cfg, upstreamRequestLog{
|
||||
URL: url,
|
||||
Method: http.MethodPost,
|
||||
Headers: httpReq.Header.Clone(),
|
||||
Body: translatedReq,
|
||||
Provider: e.Identifier(),
|
||||
AuthID: authID,
|
||||
AuthLabel: authLabel,
|
||||
AuthType: authType,
|
||||
AuthValue: authValue,
|
||||
})
|
||||
|
||||
httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0)
|
||||
httpResp, errDo := httpClient.Do(httpReq)
|
||||
if errDo != nil {
|
||||
recordAPIResponseError(ctx, e.cfg, errDo)
|
||||
return cliproxyexecutor.Response{}, errDo
|
||||
}
|
||||
defer func() {
|
||||
if errClose := httpResp.Body.Close(); errClose != nil {
|
||||
log.Errorf("vertex executor: close response body error: %v", errClose)
|
||||
}
|
||||
}()
|
||||
recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())
|
||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||
b, _ := io.ReadAll(httpResp.Body)
|
||||
appendAPIResponseChunk(ctx, e.cfg, b)
|
||||
log.Debugf("request error, error status: %d, error body: %s", httpResp.StatusCode, summarizeErrorBody(httpResp.Header.Get("Content-Type"), b))
|
||||
return cliproxyexecutor.Response{}, statusErr{code: httpResp.StatusCode, msg: string(b)}
|
||||
}
|
||||
data, errRead := io.ReadAll(httpResp.Body)
|
||||
if errRead != nil {
|
||||
recordAPIResponseError(ctx, e.cfg, errRead)
|
||||
return cliproxyexecutor.Response{}, errRead
|
||||
}
|
||||
appendAPIResponseChunk(ctx, e.cfg, data)
|
||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||
log.Debugf("request error, error status: %d, error body: %s", httpResp.StatusCode, summarizeErrorBody(httpResp.Header.Get("Content-Type"), data))
|
||||
return cliproxyexecutor.Response{}, statusErr{code: httpResp.StatusCode, msg: string(data)}
|
||||
}
|
||||
count := gjson.GetBytes(data, "totalTokens").Int()
|
||||
out := sdktranslator.TranslateTokenCount(ctx, to, from, count, data)
|
||||
return cliproxyexecutor.Response{Payload: []byte(out)}, nil
|
||||
}
|
||||
|
||||
// countTokensWithAPIKey handles token counting using API key credentials.
|
||||
func (e *GeminiVertexExecutor) countTokensWithAPIKey(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, apiKey, baseURL string) (cliproxyexecutor.Response, error) {
|
||||
from := opts.SourceFormat
|
||||
to := sdktranslator.FromString("gemini")
|
||||
translatedReq := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)
|
||||
if budgetOverride, includeOverride, ok := util.GeminiThinkingFromMetadata(req.Metadata); ok && util.ModelSupportsThinking(req.Model) {
|
||||
if budgetOverride != nil {
|
||||
norm := util.NormalizeThinkingBudget(req.Model, *budgetOverride)
|
||||
budgetOverride = &norm
|
||||
}
|
||||
translatedReq = util.ApplyGeminiThinkingConfig(translatedReq, budgetOverride, includeOverride)
|
||||
}
|
||||
translatedReq = util.StripThinkingConfigIfUnsupported(req.Model, translatedReq)
|
||||
translatedReq = fixGeminiImageAspectRatio(req.Model, translatedReq)
|
||||
respCtx := context.WithValue(ctx, "alt", opts.Alt)
|
||||
translatedReq, _ = sjson.DeleteBytes(translatedReq, "tools")
|
||||
translatedReq, _ = sjson.DeleteBytes(translatedReq, "generationConfig")
|
||||
translatedReq, _ = sjson.DeleteBytes(translatedReq, "safetySettings")
|
||||
|
||||
// For API key auth, use simpler URL format without project/location
|
||||
if baseURL == "" {
|
||||
baseURL = "https://generativelanguage.googleapis.com"
|
||||
}
|
||||
url := fmt.Sprintf("%s/%s/publishers/google/models/%s:%s", baseURL, vertexAPIVersion, req.Model, "countTokens")
|
||||
|
||||
httpReq, errNewReq := http.NewRequestWithContext(respCtx, http.MethodPost, url, bytes.NewReader(translatedReq))
|
||||
if errNewReq != nil {
|
||||
return cliproxyexecutor.Response{}, errNewReq
|
||||
}
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
if apiKey != "" {
|
||||
httpReq.Header.Set("x-goog-api-key", apiKey)
|
||||
}
|
||||
applyGeminiHeaders(httpReq, auth)
|
||||
|
||||
var authID, authLabel, authType, authValue string
|
||||
if auth != nil {
|
||||
authID = auth.ID
|
||||
authLabel = auth.Label
|
||||
authType, authValue = auth.AccountInfo()
|
||||
}
|
||||
recordAPIRequest(ctx, e.cfg, upstreamRequestLog{
|
||||
URL: url,
|
||||
Method: http.MethodPost,
|
||||
Headers: httpReq.Header.Clone(),
|
||||
Body: translatedReq,
|
||||
Provider: e.Identifier(),
|
||||
AuthID: authID,
|
||||
AuthLabel: authLabel,
|
||||
AuthType: authType,
|
||||
AuthValue: authValue,
|
||||
})
|
||||
|
||||
httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0)
|
||||
httpResp, errDo := httpClient.Do(httpReq)
|
||||
if errDo != nil {
|
||||
recordAPIResponseError(ctx, e.cfg, errDo)
|
||||
return cliproxyexecutor.Response{}, errDo
|
||||
}
|
||||
defer func() {
|
||||
if errClose := httpResp.Body.Close(); errClose != nil {
|
||||
log.Errorf("vertex executor: close response body error: %v", errClose)
|
||||
}
|
||||
}()
|
||||
recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())
|
||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||
b, _ := io.ReadAll(httpResp.Body)
|
||||
appendAPIResponseChunk(ctx, e.cfg, b)
|
||||
log.Debugf("request error, error status: %d, error body: %s", httpResp.StatusCode, summarizeErrorBody(httpResp.Header.Get("Content-Type"), b))
|
||||
return cliproxyexecutor.Response{}, statusErr{code: httpResp.StatusCode, msg: string(b)}
|
||||
}
|
||||
data, errRead := io.ReadAll(httpResp.Body)
|
||||
if errRead != nil {
|
||||
recordAPIResponseError(ctx, e.cfg, errRead)
|
||||
return cliproxyexecutor.Response{}, errRead
|
||||
}
|
||||
appendAPIResponseChunk(ctx, e.cfg, data)
|
||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||
log.Debugf("request error, error status: %d, error body: %s", httpResp.StatusCode, summarizeErrorBody(httpResp.Header.Get("Content-Type"), data))
|
||||
return cliproxyexecutor.Response{}, statusErr{code: httpResp.StatusCode, msg: string(data)}
|
||||
}
|
||||
count := gjson.GetBytes(data, "totalTokens").Int()
|
||||
out := sdktranslator.TranslateTokenCount(ctx, to, from, count, data)
|
||||
return cliproxyexecutor.Response{Payload: []byte(out)}, nil
|
||||
}
|
||||
|
||||
// Refresh is a no-op for service account based credentials.
|
||||
func (e *GeminiVertexExecutor) Refresh(_ context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) {
|
||||
return auth, nil
|
||||
}
|
||||
|
||||
// executeWithServiceAccount handles authentication using service account credentials.
|
||||
// This method contains the original service account authentication logic.
|
||||
func (e *GeminiVertexExecutor) executeWithServiceAccount(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, projectID, location string, saJSON []byte) (resp cliproxyexecutor.Response, err error) {
|
||||
reporter := newUsageReporter(ctx, e.Identifier(), req.Model, auth)
|
||||
defer reporter.trackFailure(ctx, &err)
|
||||
|
||||
@@ -149,13 +376,104 @@ func (e *GeminiVertexExecutor) Execute(ctx context.Context, auth *cliproxyauth.A
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
// ExecuteStream handles SSE streaming for Vertex.
|
||||
func (e *GeminiVertexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (stream <-chan cliproxyexecutor.StreamChunk, err error) {
|
||||
projectID, location, saJSON, errCreds := vertexCreds(auth)
|
||||
if errCreds != nil {
|
||||
return nil, errCreds
|
||||
// executeWithAPIKey handles authentication using API key credentials.
|
||||
func (e *GeminiVertexExecutor) executeWithAPIKey(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, apiKey, baseURL string) (resp cliproxyexecutor.Response, err error) {
|
||||
reporter := newUsageReporter(ctx, e.Identifier(), req.Model, auth)
|
||||
defer reporter.trackFailure(ctx, &err)
|
||||
|
||||
from := opts.SourceFormat
|
||||
to := sdktranslator.FromString("gemini")
|
||||
body := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)
|
||||
if budgetOverride, includeOverride, ok := util.GeminiThinkingFromMetadata(req.Metadata); ok && util.ModelSupportsThinking(req.Model) {
|
||||
if budgetOverride != nil {
|
||||
norm := util.NormalizeThinkingBudget(req.Model, *budgetOverride)
|
||||
budgetOverride = &norm
|
||||
}
|
||||
body = util.ApplyGeminiThinkingConfig(body, budgetOverride, includeOverride)
|
||||
}
|
||||
body = util.StripThinkingConfigIfUnsupported(req.Model, body)
|
||||
body = fixGeminiImageAspectRatio(req.Model, body)
|
||||
body = applyPayloadConfig(e.cfg, req.Model, body)
|
||||
|
||||
action := "generateContent"
|
||||
if req.Metadata != nil {
|
||||
if a, _ := req.Metadata["action"].(string); a == "countTokens" {
|
||||
action = "countTokens"
|
||||
}
|
||||
}
|
||||
|
||||
// For API key auth, use simpler URL format without project/location
|
||||
if baseURL == "" {
|
||||
baseURL = "https://generativelanguage.googleapis.com"
|
||||
}
|
||||
url := fmt.Sprintf("%s/%s/publishers/google/models/%s:%s", baseURL, vertexAPIVersion, req.Model, action)
|
||||
if opts.Alt != "" && action != "countTokens" {
|
||||
url = url + fmt.Sprintf("?$alt=%s", opts.Alt)
|
||||
}
|
||||
body, _ = sjson.DeleteBytes(body, "session_id")
|
||||
|
||||
httpReq, errNewReq := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
|
||||
if errNewReq != nil {
|
||||
return resp, errNewReq
|
||||
}
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
if apiKey != "" {
|
||||
httpReq.Header.Set("x-goog-api-key", apiKey)
|
||||
}
|
||||
applyGeminiHeaders(httpReq, auth)
|
||||
|
||||
var authID, authLabel, authType, authValue string
|
||||
if auth != nil {
|
||||
authID = auth.ID
|
||||
authLabel = auth.Label
|
||||
authType, authValue = auth.AccountInfo()
|
||||
}
|
||||
recordAPIRequest(ctx, e.cfg, upstreamRequestLog{
|
||||
URL: url,
|
||||
Method: http.MethodPost,
|
||||
Headers: httpReq.Header.Clone(),
|
||||
Body: body,
|
||||
Provider: e.Identifier(),
|
||||
AuthID: authID,
|
||||
AuthLabel: authLabel,
|
||||
AuthType: authType,
|
||||
AuthValue: authValue,
|
||||
})
|
||||
|
||||
httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0)
|
||||
httpResp, errDo := httpClient.Do(httpReq)
|
||||
if errDo != nil {
|
||||
recordAPIResponseError(ctx, e.cfg, errDo)
|
||||
return resp, errDo
|
||||
}
|
||||
defer func() {
|
||||
if errClose := httpResp.Body.Close(); errClose != nil {
|
||||
log.Errorf("vertex executor: close response body error: %v", errClose)
|
||||
}
|
||||
}()
|
||||
recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())
|
||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||
b, _ := io.ReadAll(httpResp.Body)
|
||||
appendAPIResponseChunk(ctx, e.cfg, b)
|
||||
log.Debugf("request error, error status: %d, error body: %s", httpResp.StatusCode, summarizeErrorBody(httpResp.Header.Get("Content-Type"), b))
|
||||
err = statusErr{code: httpResp.StatusCode, msg: string(b)}
|
||||
return resp, err
|
||||
}
|
||||
data, errRead := io.ReadAll(httpResp.Body)
|
||||
if errRead != nil {
|
||||
recordAPIResponseError(ctx, e.cfg, errRead)
|
||||
return resp, errRead
|
||||
}
|
||||
appendAPIResponseChunk(ctx, e.cfg, data)
|
||||
reporter.publish(ctx, parseGeminiUsage(data))
|
||||
var param any
|
||||
out := sdktranslator.TranslateNonStream(ctx, to, from, req.Model, bytes.Clone(opts.OriginalRequest), body, data, ¶m)
|
||||
resp = cliproxyexecutor.Response{Payload: []byte(out)}
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
// executeStreamWithServiceAccount handles streaming authentication using service account credentials.
|
||||
func (e *GeminiVertexExecutor) executeStreamWithServiceAccount(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, projectID, location string, saJSON []byte) (stream <-chan cliproxyexecutor.StreamChunk, err error) {
|
||||
reporter := newUsageReporter(ctx, e.Identifier(), req.Model, auth)
|
||||
defer reporter.trackFailure(ctx, &err)
|
||||
|
||||
@@ -266,42 +584,44 @@ func (e *GeminiVertexExecutor) ExecuteStream(ctx context.Context, auth *cliproxy
|
||||
return stream, nil
|
||||
}
|
||||
|
||||
// CountTokens calls Vertex countTokens endpoint.
|
||||
func (e *GeminiVertexExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
||||
projectID, location, saJSON, errCreds := vertexCreds(auth)
|
||||
if errCreds != nil {
|
||||
return cliproxyexecutor.Response{}, errCreds
|
||||
}
|
||||
// executeStreamWithAPIKey handles streaming authentication using API key credentials.
|
||||
func (e *GeminiVertexExecutor) executeStreamWithAPIKey(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, apiKey, baseURL string) (stream <-chan cliproxyexecutor.StreamChunk, err error) {
|
||||
reporter := newUsageReporter(ctx, e.Identifier(), req.Model, auth)
|
||||
defer reporter.trackFailure(ctx, &err)
|
||||
|
||||
from := opts.SourceFormat
|
||||
to := sdktranslator.FromString("gemini")
|
||||
translatedReq := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)
|
||||
body := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), true)
|
||||
if budgetOverride, includeOverride, ok := util.GeminiThinkingFromMetadata(req.Metadata); ok && util.ModelSupportsThinking(req.Model) {
|
||||
if budgetOverride != nil {
|
||||
norm := util.NormalizeThinkingBudget(req.Model, *budgetOverride)
|
||||
budgetOverride = &norm
|
||||
}
|
||||
translatedReq = util.ApplyGeminiThinkingConfig(translatedReq, budgetOverride, includeOverride)
|
||||
body = util.ApplyGeminiThinkingConfig(body, budgetOverride, includeOverride)
|
||||
}
|
||||
translatedReq = util.StripThinkingConfigIfUnsupported(req.Model, translatedReq)
|
||||
translatedReq = fixGeminiImageAspectRatio(req.Model, translatedReq)
|
||||
respCtx := context.WithValue(ctx, "alt", opts.Alt)
|
||||
translatedReq, _ = sjson.DeleteBytes(translatedReq, "tools")
|
||||
translatedReq, _ = sjson.DeleteBytes(translatedReq, "generationConfig")
|
||||
translatedReq, _ = sjson.DeleteBytes(translatedReq, "safetySettings")
|
||||
body = util.StripThinkingConfigIfUnsupported(req.Model, body)
|
||||
body = fixGeminiImageAspectRatio(req.Model, body)
|
||||
body = applyPayloadConfig(e.cfg, req.Model, body)
|
||||
|
||||
baseURL := vertexBaseURL(location)
|
||||
url := fmt.Sprintf("%s/%s/projects/%s/locations/%s/publishers/google/models/%s:%s", baseURL, vertexAPIVersion, projectID, location, req.Model, "countTokens")
|
||||
// For API key auth, use simpler URL format without project/location
|
||||
if baseURL == "" {
|
||||
baseURL = "https://generativelanguage.googleapis.com"
|
||||
}
|
||||
url := fmt.Sprintf("%s/%s/publishers/google/models/%s:%s", baseURL, vertexAPIVersion, req.Model, "streamGenerateContent")
|
||||
if opts.Alt == "" {
|
||||
url = url + "?alt=sse"
|
||||
} else {
|
||||
url = url + fmt.Sprintf("?$alt=%s", opts.Alt)
|
||||
}
|
||||
body, _ = sjson.DeleteBytes(body, "session_id")
|
||||
|
||||
httpReq, errNewReq := http.NewRequestWithContext(respCtx, http.MethodPost, url, bytes.NewReader(translatedReq))
|
||||
httpReq, errNewReq := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
|
||||
if errNewReq != nil {
|
||||
return cliproxyexecutor.Response{}, errNewReq
|
||||
return nil, errNewReq
|
||||
}
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
if token, errTok := vertexAccessToken(ctx, e.cfg, auth, saJSON); errTok == nil && token != "" {
|
||||
httpReq.Header.Set("Authorization", "Bearer "+token)
|
||||
} else if errTok != nil {
|
||||
log.Errorf("vertex executor: access token error: %v", errTok)
|
||||
return cliproxyexecutor.Response{}, statusErr{code: 500, msg: "internal server error"}
|
||||
if apiKey != "" {
|
||||
httpReq.Header.Set("x-goog-api-key", apiKey)
|
||||
}
|
||||
applyGeminiHeaders(httpReq, auth)
|
||||
|
||||
@@ -315,7 +635,7 @@ func (e *GeminiVertexExecutor) CountTokens(ctx context.Context, auth *cliproxyau
|
||||
URL: url,
|
||||
Method: http.MethodPost,
|
||||
Headers: httpReq.Header.Clone(),
|
||||
Body: translatedReq,
|
||||
Body: body,
|
||||
Provider: e.Identifier(),
|
||||
AuthID: authID,
|
||||
AuthLabel: authLabel,
|
||||
@@ -327,38 +647,53 @@ func (e *GeminiVertexExecutor) CountTokens(ctx context.Context, auth *cliproxyau
|
||||
httpResp, errDo := httpClient.Do(httpReq)
|
||||
if errDo != nil {
|
||||
recordAPIResponseError(ctx, e.cfg, errDo)
|
||||
return cliproxyexecutor.Response{}, errDo
|
||||
return nil, errDo
|
||||
}
|
||||
defer func() {
|
||||
if errClose := httpResp.Body.Close(); errClose != nil {
|
||||
log.Errorf("vertex executor: close response body error: %v", errClose)
|
||||
}
|
||||
}()
|
||||
recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())
|
||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||
b, _ := io.ReadAll(httpResp.Body)
|
||||
appendAPIResponseChunk(ctx, e.cfg, b)
|
||||
log.Debugf("request error, error status: %d, error body: %s", httpResp.StatusCode, summarizeErrorBody(httpResp.Header.Get("Content-Type"), b))
|
||||
return cliproxyexecutor.Response{}, statusErr{code: httpResp.StatusCode, msg: string(b)}
|
||||
if errClose := httpResp.Body.Close(); errClose != nil {
|
||||
log.Errorf("vertex executor: close response body error: %v", errClose)
|
||||
}
|
||||
return nil, statusErr{code: httpResp.StatusCode, msg: string(b)}
|
||||
}
|
||||
data, errRead := io.ReadAll(httpResp.Body)
|
||||
if errRead != nil {
|
||||
recordAPIResponseError(ctx, e.cfg, errRead)
|
||||
return cliproxyexecutor.Response{}, errRead
|
||||
}
|
||||
appendAPIResponseChunk(ctx, e.cfg, data)
|
||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||
log.Debugf("request error, error status: %d, error body: %s", httpResp.StatusCode, summarizeErrorBody(httpResp.Header.Get("Content-Type"), data))
|
||||
return cliproxyexecutor.Response{}, statusErr{code: httpResp.StatusCode, msg: string(data)}
|
||||
}
|
||||
count := gjson.GetBytes(data, "totalTokens").Int()
|
||||
out := sdktranslator.TranslateTokenCount(ctx, to, from, count, data)
|
||||
return cliproxyexecutor.Response{Payload: []byte(out)}, nil
|
||||
}
|
||||
|
||||
// Refresh is a no-op for service account based credentials.
|
||||
func (e *GeminiVertexExecutor) Refresh(_ context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) {
|
||||
return auth, nil
|
||||
out := make(chan cliproxyexecutor.StreamChunk)
|
||||
stream = out
|
||||
go func() {
|
||||
defer close(out)
|
||||
defer func() {
|
||||
if errClose := httpResp.Body.Close(); errClose != nil {
|
||||
log.Errorf("vertex executor: close response body error: %v", errClose)
|
||||
}
|
||||
}()
|
||||
scanner := bufio.NewScanner(httpResp.Body)
|
||||
scanner.Buffer(nil, 20_971_520)
|
||||
var param any
|
||||
for scanner.Scan() {
|
||||
line := scanner.Bytes()
|
||||
appendAPIResponseChunk(ctx, e.cfg, line)
|
||||
if detail, ok := parseGeminiStreamUsage(line); ok {
|
||||
reporter.publish(ctx, detail)
|
||||
}
|
||||
lines := sdktranslator.TranslateStream(ctx, to, from, req.Model, bytes.Clone(opts.OriginalRequest), body, bytes.Clone(line), ¶m)
|
||||
for i := range lines {
|
||||
out <- cliproxyexecutor.StreamChunk{Payload: []byte(lines[i])}
|
||||
}
|
||||
}
|
||||
lines := sdktranslator.TranslateStream(ctx, to, from, req.Model, bytes.Clone(opts.OriginalRequest), body, []byte("[DONE]"), ¶m)
|
||||
for i := range lines {
|
||||
out <- cliproxyexecutor.StreamChunk{Payload: []byte(lines[i])}
|
||||
}
|
||||
if errScan := scanner.Err(); errScan != nil {
|
||||
recordAPIResponseError(ctx, e.cfg, errScan)
|
||||
reporter.publishFailure(ctx)
|
||||
out <- cliproxyexecutor.StreamChunk{Err: errScan}
|
||||
}
|
||||
}()
|
||||
return stream, nil
|
||||
}
|
||||
|
||||
// vertexCreds extracts project, location and raw service account JSON from auth metadata.
|
||||
@@ -401,6 +736,23 @@ func vertexCreds(a *cliproxyauth.Auth) (projectID, location string, serviceAccou
|
||||
return projectID, location, saJSON, nil
|
||||
}
|
||||
|
||||
// vertexAPICreds extracts API key and base URL from auth attributes following the claudeCreds pattern.
|
||||
func vertexAPICreds(a *cliproxyauth.Auth) (apiKey, baseURL string) {
|
||||
if a == nil {
|
||||
return "", ""
|
||||
}
|
||||
if a.Attributes != nil {
|
||||
apiKey = a.Attributes["api_key"]
|
||||
baseURL = a.Attributes["base_url"]
|
||||
}
|
||||
if apiKey == "" && a.Metadata != nil {
|
||||
if v, ok := a.Metadata["access_token"].(string); ok {
|
||||
apiKey = v
|
||||
}
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func vertexBaseURL(location string) string {
|
||||
loc := strings.TrimSpace(location)
|
||||
if loc == "" {
|
||||
|
||||
@@ -4,10 +4,42 @@ import (
|
||||
"strings"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
|
||||
"github.com/tidwall/gjson"
|
||||
"github.com/tidwall/sjson"
|
||||
)
|
||||
|
||||
// applyThinkingMetadata applies thinking config from model suffix metadata (e.g., -reasoning, -thinking-N)
|
||||
// for standard Gemini format payloads. It normalizes the budget when the model supports thinking.
|
||||
func applyThinkingMetadata(payload []byte, metadata map[string]any, model string) []byte {
|
||||
budgetOverride, includeOverride, ok := util.GeminiThinkingFromMetadata(metadata)
|
||||
if !ok {
|
||||
return payload
|
||||
}
|
||||
if !util.ModelSupportsThinking(model) {
|
||||
return payload
|
||||
}
|
||||
if budgetOverride != nil {
|
||||
norm := util.NormalizeThinkingBudget(model, *budgetOverride)
|
||||
budgetOverride = &norm
|
||||
}
|
||||
return util.ApplyGeminiThinkingConfig(payload, budgetOverride, includeOverride)
|
||||
}
|
||||
|
||||
// applyThinkingMetadataCLI applies thinking config from model suffix metadata (e.g., -reasoning, -thinking-N)
|
||||
// for Gemini CLI format payloads (nested under "request"). It normalizes the budget when the model supports thinking.
|
||||
func applyThinkingMetadataCLI(payload []byte, metadata map[string]any, model string) []byte {
|
||||
budgetOverride, includeOverride, ok := util.GeminiThinkingFromMetadata(metadata)
|
||||
if !ok {
|
||||
return payload
|
||||
}
|
||||
if budgetOverride != nil && util.ModelSupportsThinking(model) {
|
||||
norm := util.NormalizeThinkingBudget(model, *budgetOverride)
|
||||
budgetOverride = &norm
|
||||
}
|
||||
return util.ApplyGeminiCLIThinkingConfig(payload, budgetOverride, includeOverride)
|
||||
}
|
||||
|
||||
// applyPayloadConfig applies payload default and override rules from configuration
|
||||
// to the given JSON payload for the specified model.
|
||||
// Defaults only fill missing fields, while overrides always overwrite existing values.
|
||||
|
||||
@@ -37,7 +37,7 @@ func newUsageReporter(ctx context.Context, provider, model string, auth *cliprox
|
||||
}
|
||||
if auth != nil {
|
||||
reporter.authID = auth.ID
|
||||
reporter.authIndex = auth.Index
|
||||
reporter.authIndex = auth.EnsureIndex()
|
||||
}
|
||||
return reporter
|
||||
}
|
||||
|
||||
@@ -271,7 +271,15 @@ func ConvertOpenAIRequestToAntigravity(modelName string, inputRawJSON []byte, _
|
||||
if resp == "" {
|
||||
resp = "{}"
|
||||
}
|
||||
toolNode, _ = sjson.SetBytes(toolNode, "parts."+itoa(pp)+".functionResponse.response.result", []byte(resp))
|
||||
// Handle non-JSON output gracefully (matches dev branch approach)
|
||||
if resp != "null" {
|
||||
parsed := gjson.Parse(resp)
|
||||
if parsed.Type == gjson.JSON {
|
||||
toolNode, _ = sjson.SetRawBytes(toolNode, "parts."+itoa(pp)+".functionResponse.response.result", []byte(parsed.Raw))
|
||||
} else {
|
||||
toolNode, _ = sjson.SetBytes(toolNode, "parts."+itoa(pp)+".functionResponse.response.result", resp)
|
||||
}
|
||||
}
|
||||
pp++
|
||||
}
|
||||
}
|
||||
|
||||
@@ -91,6 +91,11 @@ func ConvertGeminiRequestToGemini(_ string, inputRawJSON []byte, _ bool) []byte
|
||||
return true
|
||||
})
|
||||
|
||||
if gjson.GetBytes(rawJSON, "generationConfig.responseSchema").Exists() {
|
||||
strJson, _ := util.RenameKey(string(out), "generationConfig.responseSchema", "generationConfig.responseJsonSchema")
|
||||
out = []byte(strJson)
|
||||
}
|
||||
|
||||
out = common.AttachDefaultSafetySettings(out, "safetySettings")
|
||||
return out
|
||||
}
|
||||
|
||||
@@ -202,6 +202,8 @@ func ConvertOpenAIResponsesRequestToOpenAIChatCompletions(modelName string, inpu
|
||||
out, _ = sjson.Set(out, "reasoning_effort", "medium")
|
||||
case "high":
|
||||
out, _ = sjson.Set(out, "reasoning_effort", "high")
|
||||
case "xhigh":
|
||||
out, _ = sjson.Set(out, "reasoning_effort", "xhigh")
|
||||
default:
|
||||
out, _ = sjson.Set(out, "reasoning_effort", "auto")
|
||||
}
|
||||
|
||||
@@ -34,6 +34,15 @@ func ParseGeminiThinkingSuffix(model string) (string, *int, *bool, bool) {
|
||||
return base, &budgetValue, &include, true
|
||||
}
|
||||
|
||||
// Handle "-reasoning" suffix: enables thinking with dynamic budget (-1)
|
||||
// Maps: gemini-2.5-flash-reasoning -> gemini-2.5-flash with thinkingBudget=-1
|
||||
if strings.HasSuffix(lower, "-reasoning") {
|
||||
base := model[:len(model)-len("-reasoning")]
|
||||
budgetValue := -1 // Dynamic budget
|
||||
include := true
|
||||
return base, &budgetValue, &include, true
|
||||
}
|
||||
|
||||
idx := strings.LastIndex(lower, "-thinking-")
|
||||
if idx == -1 {
|
||||
return model, nil, nil, false
|
||||
|
||||
@@ -30,6 +30,16 @@ import (
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
func matchProvider(provider string, targets []string) (string, bool) {
|
||||
p := strings.ToLower(strings.TrimSpace(provider))
|
||||
for _, t := range targets {
|
||||
if strings.EqualFold(p, strings.TrimSpace(t)) {
|
||||
return p, true
|
||||
}
|
||||
}
|
||||
return p, false
|
||||
}
|
||||
|
||||
// storePersister captures persistence-capable token store methods used by the watcher.
|
||||
type storePersister interface {
|
||||
PersistConfig(ctx context.Context) error
|
||||
@@ -54,6 +64,7 @@ type Watcher struct {
|
||||
lastConfigHash string
|
||||
authQueue chan<- AuthUpdate
|
||||
currentAuths map[string]*coreauth.Auth
|
||||
runtimeAuths map[string]*coreauth.Auth
|
||||
dispatchMu sync.Mutex
|
||||
dispatchCond *sync.Cond
|
||||
pendingUpdates map[string]AuthUpdate
|
||||
@@ -169,7 +180,7 @@ func (w *Watcher) Start(ctx context.Context) error {
|
||||
go w.processEvents(ctx)
|
||||
|
||||
// Perform an initial full reload based on current config and auth dir
|
||||
w.reloadClients(true)
|
||||
w.reloadClients(true, nil)
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -221,9 +232,57 @@ func (w *Watcher) SetAuthUpdateQueue(queue chan<- AuthUpdate) {
|
||||
}
|
||||
}
|
||||
|
||||
// DispatchRuntimeAuthUpdate allows external runtime providers (e.g., websocket-driven auths)
|
||||
// to push auth updates through the same queue used by file/config watchers.
|
||||
// Returns true if the update was enqueued; false if no queue is configured.
|
||||
func (w *Watcher) DispatchRuntimeAuthUpdate(update AuthUpdate) bool {
|
||||
if w == nil {
|
||||
return false
|
||||
}
|
||||
w.clientsMutex.Lock()
|
||||
if w.runtimeAuths == nil {
|
||||
w.runtimeAuths = make(map[string]*coreauth.Auth)
|
||||
}
|
||||
switch update.Action {
|
||||
case AuthUpdateActionAdd, AuthUpdateActionModify:
|
||||
if update.Auth != nil && update.Auth.ID != "" {
|
||||
clone := update.Auth.Clone()
|
||||
w.runtimeAuths[clone.ID] = clone
|
||||
if w.currentAuths == nil {
|
||||
w.currentAuths = make(map[string]*coreauth.Auth)
|
||||
}
|
||||
w.currentAuths[clone.ID] = clone.Clone()
|
||||
}
|
||||
case AuthUpdateActionDelete:
|
||||
id := update.ID
|
||||
if id == "" && update.Auth != nil {
|
||||
id = update.Auth.ID
|
||||
}
|
||||
if id != "" {
|
||||
delete(w.runtimeAuths, id)
|
||||
if w.currentAuths != nil {
|
||||
delete(w.currentAuths, id)
|
||||
}
|
||||
}
|
||||
}
|
||||
w.clientsMutex.Unlock()
|
||||
if w.getAuthQueue() == nil {
|
||||
return false
|
||||
}
|
||||
w.dispatchAuthUpdates([]AuthUpdate{update})
|
||||
return true
|
||||
}
|
||||
|
||||
func (w *Watcher) refreshAuthState() {
|
||||
auths := w.SnapshotCoreAuths()
|
||||
w.clientsMutex.Lock()
|
||||
if len(w.runtimeAuths) > 0 {
|
||||
for _, a := range w.runtimeAuths {
|
||||
if a != nil {
|
||||
auths = append(auths, a.Clone())
|
||||
}
|
||||
}
|
||||
}
|
||||
updates := w.prepareAuthUpdatesLocked(auths)
|
||||
w.clientsMutex.Unlock()
|
||||
w.dispatchAuthUpdates(updates)
|
||||
@@ -437,6 +496,18 @@ func computeOpenAICompatModelsHash(models []config.OpenAICompatibilityModel) str
|
||||
return hex.EncodeToString(sum[:])
|
||||
}
|
||||
|
||||
func computeVertexCompatModelsHash(models []config.VertexCompatModel) string {
|
||||
if len(models) == 0 {
|
||||
return ""
|
||||
}
|
||||
data, err := json.Marshal(models)
|
||||
if err != nil || len(data) == 0 {
|
||||
return ""
|
||||
}
|
||||
sum := sha256.Sum256(data)
|
||||
return hex.EncodeToString(sum[:])
|
||||
}
|
||||
|
||||
// computeClaudeModelsHash returns a stable hash for Claude model aliases.
|
||||
func computeClaudeModelsHash(models []config.ClaudeModel) string {
|
||||
if len(models) == 0 {
|
||||
@@ -450,6 +521,142 @@ func computeClaudeModelsHash(models []config.ClaudeModel) string {
|
||||
return hex.EncodeToString(sum[:])
|
||||
}
|
||||
|
||||
func computeExcludedModelsHash(excluded []string) string {
|
||||
if len(excluded) == 0 {
|
||||
return ""
|
||||
}
|
||||
normalized := make([]string, 0, len(excluded))
|
||||
for _, entry := range excluded {
|
||||
if trimmed := strings.TrimSpace(entry); trimmed != "" {
|
||||
normalized = append(normalized, strings.ToLower(trimmed))
|
||||
}
|
||||
}
|
||||
if len(normalized) == 0 {
|
||||
return ""
|
||||
}
|
||||
sort.Strings(normalized)
|
||||
data, err := json.Marshal(normalized)
|
||||
if err != nil || len(data) == 0 {
|
||||
return ""
|
||||
}
|
||||
sum := sha256.Sum256(data)
|
||||
return hex.EncodeToString(sum[:])
|
||||
}
|
||||
|
||||
type excludedModelsSummary struct {
|
||||
hash string
|
||||
count int
|
||||
}
|
||||
|
||||
func summarizeExcludedModels(list []string) excludedModelsSummary {
|
||||
if len(list) == 0 {
|
||||
return excludedModelsSummary{}
|
||||
}
|
||||
seen := make(map[string]struct{}, len(list))
|
||||
normalized := make([]string, 0, len(list))
|
||||
for _, entry := range list {
|
||||
if trimmed := strings.ToLower(strings.TrimSpace(entry)); trimmed != "" {
|
||||
if _, exists := seen[trimmed]; exists {
|
||||
continue
|
||||
}
|
||||
seen[trimmed] = struct{}{}
|
||||
normalized = append(normalized, trimmed)
|
||||
}
|
||||
}
|
||||
sort.Strings(normalized)
|
||||
return excludedModelsSummary{
|
||||
hash: computeExcludedModelsHash(normalized),
|
||||
count: len(normalized),
|
||||
}
|
||||
}
|
||||
|
||||
func summarizeOAuthExcludedModels(entries map[string][]string) map[string]excludedModelsSummary {
|
||||
if len(entries) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make(map[string]excludedModelsSummary, len(entries))
|
||||
for k, v := range entries {
|
||||
key := strings.ToLower(strings.TrimSpace(k))
|
||||
if key == "" {
|
||||
continue
|
||||
}
|
||||
out[key] = summarizeExcludedModels(v)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func diffOAuthExcludedModelChanges(oldMap, newMap map[string][]string) ([]string, []string) {
|
||||
oldSummary := summarizeOAuthExcludedModels(oldMap)
|
||||
newSummary := summarizeOAuthExcludedModels(newMap)
|
||||
keys := make(map[string]struct{}, len(oldSummary)+len(newSummary))
|
||||
for k := range oldSummary {
|
||||
keys[k] = struct{}{}
|
||||
}
|
||||
for k := range newSummary {
|
||||
keys[k] = struct{}{}
|
||||
}
|
||||
changes := make([]string, 0, len(keys))
|
||||
affected := make([]string, 0, len(keys))
|
||||
for key := range keys {
|
||||
oldInfo, okOld := oldSummary[key]
|
||||
newInfo, okNew := newSummary[key]
|
||||
switch {
|
||||
case okOld && !okNew:
|
||||
changes = append(changes, fmt.Sprintf("oauth-excluded-models[%s]: removed", key))
|
||||
affected = append(affected, key)
|
||||
case !okOld && okNew:
|
||||
changes = append(changes, fmt.Sprintf("oauth-excluded-models[%s]: added (%d entries)", key, newInfo.count))
|
||||
affected = append(affected, key)
|
||||
case okOld && okNew && oldInfo.hash != newInfo.hash:
|
||||
changes = append(changes, fmt.Sprintf("oauth-excluded-models[%s]: updated (%d -> %d entries)", key, oldInfo.count, newInfo.count))
|
||||
affected = append(affected, key)
|
||||
}
|
||||
}
|
||||
sort.Strings(changes)
|
||||
sort.Strings(affected)
|
||||
return changes, affected
|
||||
}
|
||||
|
||||
func applyAuthExcludedModelsMeta(auth *coreauth.Auth, cfg *config.Config, perKey []string, authKind string) {
|
||||
if auth == nil || cfg == nil {
|
||||
return
|
||||
}
|
||||
authKindKey := strings.ToLower(strings.TrimSpace(authKind))
|
||||
seen := make(map[string]struct{})
|
||||
add := func(list []string) {
|
||||
for _, entry := range list {
|
||||
if trimmed := strings.TrimSpace(entry); trimmed != "" {
|
||||
key := strings.ToLower(trimmed)
|
||||
if _, exists := seen[key]; exists {
|
||||
continue
|
||||
}
|
||||
seen[key] = struct{}{}
|
||||
}
|
||||
}
|
||||
}
|
||||
if authKindKey == "apikey" {
|
||||
add(perKey)
|
||||
} else if cfg.OAuthExcludedModels != nil {
|
||||
providerKey := strings.ToLower(strings.TrimSpace(auth.Provider))
|
||||
add(cfg.OAuthExcludedModels[providerKey])
|
||||
}
|
||||
combined := make([]string, 0, len(seen))
|
||||
for k := range seen {
|
||||
combined = append(combined, k)
|
||||
}
|
||||
sort.Strings(combined)
|
||||
hash := computeExcludedModelsHash(combined)
|
||||
if auth.Attributes == nil {
|
||||
auth.Attributes = make(map[string]string)
|
||||
}
|
||||
if hash != "" {
|
||||
auth.Attributes["excluded_models_hash"] = hash
|
||||
}
|
||||
if authKind != "" {
|
||||
auth.Attributes["auth_kind"] = authKind
|
||||
}
|
||||
}
|
||||
|
||||
// SetClients sets the file-based clients.
|
||||
// SetClients removed
|
||||
// SetAPIKeyClients removed
|
||||
@@ -634,6 +841,11 @@ func (w *Watcher) reloadConfig() bool {
|
||||
w.config = newConfig
|
||||
w.clientsMutex.Unlock()
|
||||
|
||||
var affectedOAuthProviders []string
|
||||
if oldConfig != nil {
|
||||
_, affectedOAuthProviders = diffOAuthExcludedModelChanges(oldConfig.OAuthExcludedModels, newConfig.OAuthExcludedModels)
|
||||
}
|
||||
|
||||
// Always apply the current log level based on the latest config.
|
||||
// This ensures logrus reflects the desired level even if change detection misses.
|
||||
util.SetLogLevel(newConfig)
|
||||
@@ -659,12 +871,12 @@ func (w *Watcher) reloadConfig() bool {
|
||||
|
||||
log.Infof("config successfully reloaded, triggering client reload")
|
||||
// Reload clients with new config
|
||||
w.reloadClients(authDirChanged)
|
||||
w.reloadClients(authDirChanged, affectedOAuthProviders)
|
||||
return true
|
||||
}
|
||||
|
||||
// reloadClients performs a full scan and reload of all clients.
|
||||
func (w *Watcher) reloadClients(rescanAuth bool) {
|
||||
func (w *Watcher) reloadClients(rescanAuth bool, affectedOAuthProviders []string) {
|
||||
log.Debugf("starting full client load process")
|
||||
|
||||
w.clientsMutex.RLock()
|
||||
@@ -676,12 +888,34 @@ func (w *Watcher) reloadClients(rescanAuth bool) {
|
||||
return
|
||||
}
|
||||
|
||||
if len(affectedOAuthProviders) > 0 {
|
||||
w.clientsMutex.Lock()
|
||||
if w.currentAuths != nil {
|
||||
filtered := make(map[string]*coreauth.Auth, len(w.currentAuths))
|
||||
for id, auth := range w.currentAuths {
|
||||
if auth == nil {
|
||||
continue
|
||||
}
|
||||
provider := strings.ToLower(strings.TrimSpace(auth.Provider))
|
||||
if _, match := matchProvider(provider, affectedOAuthProviders); match {
|
||||
continue
|
||||
}
|
||||
filtered[id] = auth
|
||||
}
|
||||
w.currentAuths = filtered
|
||||
log.Debugf("applying oauth-excluded-models to providers %v", affectedOAuthProviders)
|
||||
} else {
|
||||
w.currentAuths = nil
|
||||
}
|
||||
w.clientsMutex.Unlock()
|
||||
}
|
||||
|
||||
// Unregister all old API key clients before creating new ones
|
||||
// no legacy clients to unregister
|
||||
|
||||
// Create new API key clients based on the new config
|
||||
geminiAPIKeyCount, claudeAPIKeyCount, codexAPIKeyCount, openAICompatCount := BuildAPIKeyClients(cfg)
|
||||
totalAPIKeyClients := geminiAPIKeyCount + claudeAPIKeyCount + codexAPIKeyCount + openAICompatCount
|
||||
geminiAPIKeyCount, vertexCompatAPIKeyCount, claudeAPIKeyCount, codexAPIKeyCount, openAICompatCount := BuildAPIKeyClients(cfg)
|
||||
totalAPIKeyClients := geminiAPIKeyCount + vertexCompatAPIKeyCount + claudeAPIKeyCount + codexAPIKeyCount + openAICompatCount
|
||||
log.Debugf("loaded %d API key clients", totalAPIKeyClients)
|
||||
|
||||
var authFileCount int
|
||||
@@ -724,7 +958,7 @@ func (w *Watcher) reloadClients(rescanAuth bool) {
|
||||
w.clientsMutex.Unlock()
|
||||
}
|
||||
|
||||
totalNewClients := authFileCount + geminiAPIKeyCount + claudeAPIKeyCount + codexAPIKeyCount + openAICompatCount
|
||||
totalNewClients := authFileCount + geminiAPIKeyCount + vertexCompatAPIKeyCount + claudeAPIKeyCount + codexAPIKeyCount + openAICompatCount
|
||||
|
||||
// Ensure consumers observe the new configuration before auth updates dispatch.
|
||||
if w.reloadCallback != nil {
|
||||
@@ -734,10 +968,11 @@ func (w *Watcher) reloadClients(rescanAuth bool) {
|
||||
|
||||
w.refreshAuthState()
|
||||
|
||||
log.Infof("full client load complete - %d clients (%d auth files + %d Gemini API keys + %d Claude API keys + %d Codex keys + %d OpenAI-compat)",
|
||||
log.Infof("full client load complete - %d clients (%d auth files + %d Gemini API keys + %d Vertex API keys + %d Claude API keys + %d Codex keys + %d OpenAI-compat)",
|
||||
totalNewClients,
|
||||
authFileCount,
|
||||
geminiAPIKeyCount,
|
||||
vertexCompatAPIKeyCount,
|
||||
claudeAPIKeyCount,
|
||||
codexAPIKeyCount,
|
||||
openAICompatCount,
|
||||
@@ -849,8 +1084,10 @@ func (w *Watcher) SnapshotCoreAuths() []*coreauth.Auth {
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
}
|
||||
applyAuthExcludedModelsMeta(a, cfg, entry.ExcludedModels, "apikey")
|
||||
out = append(out, a)
|
||||
}
|
||||
|
||||
// Claude API keys -> synthesize auths
|
||||
for i := range cfg.ClaudeKey {
|
||||
ck := cfg.ClaudeKey[i]
|
||||
@@ -882,6 +1119,7 @@ func (w *Watcher) SnapshotCoreAuths() []*coreauth.Auth {
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
}
|
||||
applyAuthExcludedModelsMeta(a, cfg, ck.ExcludedModels, "apikey")
|
||||
out = append(out, a)
|
||||
}
|
||||
// Codex API keys -> synthesize auths
|
||||
@@ -911,6 +1149,7 @@ func (w *Watcher) SnapshotCoreAuths() []*coreauth.Auth {
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
}
|
||||
applyAuthExcludedModelsMeta(a, cfg, ck.ExcludedModels, "apikey")
|
||||
out = append(out, a)
|
||||
}
|
||||
for i := range cfg.OpenAICompatibility {
|
||||
@@ -1015,6 +1254,43 @@ func (w *Watcher) SnapshotCoreAuths() []*coreauth.Auth {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Process Vertex API key providers (Vertex-compatible endpoints)
|
||||
for i := range cfg.VertexCompatAPIKey {
|
||||
compat := &cfg.VertexCompatAPIKey[i]
|
||||
providerName := "vertex"
|
||||
base := strings.TrimSpace(compat.BaseURL)
|
||||
|
||||
key := strings.TrimSpace(compat.APIKey)
|
||||
proxyURL := strings.TrimSpace(compat.ProxyURL)
|
||||
idKind := fmt.Sprintf("vertex:apikey:%s", base)
|
||||
id, token := idGen.next(idKind, key, base, proxyURL)
|
||||
attrs := map[string]string{
|
||||
"source": fmt.Sprintf("config:vertex-apikey[%s]", token),
|
||||
"base_url": base,
|
||||
"provider_key": providerName,
|
||||
}
|
||||
if key != "" {
|
||||
attrs["api_key"] = key
|
||||
}
|
||||
if hash := computeVertexCompatModelsHash(compat.Models); hash != "" {
|
||||
attrs["models_hash"] = hash
|
||||
}
|
||||
addConfigHeadersToAttrs(compat.Headers, attrs)
|
||||
a := &coreauth.Auth{
|
||||
ID: id,
|
||||
Provider: providerName,
|
||||
Label: "vertex-apikey",
|
||||
Status: coreauth.StatusActive,
|
||||
ProxyURL: proxyURL,
|
||||
Attributes: attrs,
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
}
|
||||
applyAuthExcludedModelsMeta(a, cfg, nil, "apikey")
|
||||
out = append(out, a)
|
||||
}
|
||||
|
||||
// Also synthesize auth entries directly from auth files (for OAuth/file-backed providers)
|
||||
entries, _ := os.ReadDir(w.authDir)
|
||||
for _, e := range entries {
|
||||
@@ -1071,8 +1347,12 @@ func (w *Watcher) SnapshotCoreAuths() []*coreauth.Auth {
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
}
|
||||
applyAuthExcludedModelsMeta(a, cfg, nil, "oauth")
|
||||
if provider == "gemini-cli" {
|
||||
if virtuals := synthesizeGeminiVirtualAuths(a, metadata, now); len(virtuals) > 0 {
|
||||
for _, v := range virtuals {
|
||||
applyAuthExcludedModelsMeta(v, cfg, nil, "oauth")
|
||||
}
|
||||
out = append(out, a)
|
||||
out = append(out, virtuals...)
|
||||
continue
|
||||
@@ -1227,8 +1507,9 @@ func (w *Watcher) loadFileClients(cfg *config.Config) int {
|
||||
return authFileCount
|
||||
}
|
||||
|
||||
func BuildAPIKeyClients(cfg *config.Config) (int, int, int, int) {
|
||||
func BuildAPIKeyClients(cfg *config.Config) (int, int, int, int, int) {
|
||||
geminiAPIKeyCount := 0
|
||||
vertexCompatAPIKeyCount := 0
|
||||
claudeAPIKeyCount := 0
|
||||
codexAPIKeyCount := 0
|
||||
openAICompatCount := 0
|
||||
@@ -1237,6 +1518,9 @@ func BuildAPIKeyClients(cfg *config.Config) (int, int, int, int) {
|
||||
// Stateless executor handles Gemini API keys; avoid constructing legacy clients.
|
||||
geminiAPIKeyCount += len(cfg.GeminiKey)
|
||||
}
|
||||
if len(cfg.VertexCompatAPIKey) > 0 {
|
||||
vertexCompatAPIKeyCount += len(cfg.VertexCompatAPIKey)
|
||||
}
|
||||
if len(cfg.ClaudeKey) > 0 {
|
||||
claudeAPIKeyCount += len(cfg.ClaudeKey)
|
||||
}
|
||||
@@ -1254,7 +1538,7 @@ func BuildAPIKeyClients(cfg *config.Config) (int, int, int, int) {
|
||||
}
|
||||
}
|
||||
}
|
||||
return geminiAPIKeyCount, claudeAPIKeyCount, codexAPIKeyCount, openAICompatCount
|
||||
return geminiAPIKeyCount, vertexCompatAPIKeyCount, claudeAPIKeyCount, codexAPIKeyCount, openAICompatCount
|
||||
}
|
||||
|
||||
func diffOpenAICompatibility(oldList, newList []config.OpenAICompatibility) []string {
|
||||
@@ -1464,6 +1748,11 @@ func buildConfigChangeDetails(oldCfg, newCfg *config.Config) []string {
|
||||
if !equalStringMap(o.Headers, n.Headers) {
|
||||
changes = append(changes, fmt.Sprintf("gemini[%d].headers: updated", i))
|
||||
}
|
||||
oldExcluded := summarizeExcludedModels(o.ExcludedModels)
|
||||
newExcluded := summarizeExcludedModels(n.ExcludedModels)
|
||||
if oldExcluded.hash != newExcluded.hash {
|
||||
changes = append(changes, fmt.Sprintf("gemini[%d].excluded-models: updated (%d -> %d entries)", i, oldExcluded.count, newExcluded.count))
|
||||
}
|
||||
}
|
||||
if !reflect.DeepEqual(trimStrings(oldCfg.GlAPIKey), trimStrings(newCfg.GlAPIKey)) {
|
||||
changes = append(changes, "generative-language-api-key: values updated (legacy view, redacted)")
|
||||
@@ -1492,6 +1781,11 @@ func buildConfigChangeDetails(oldCfg, newCfg *config.Config) []string {
|
||||
if !equalStringMap(o.Headers, n.Headers) {
|
||||
changes = append(changes, fmt.Sprintf("claude[%d].headers: updated", i))
|
||||
}
|
||||
oldExcluded := summarizeExcludedModels(o.ExcludedModels)
|
||||
newExcluded := summarizeExcludedModels(n.ExcludedModels)
|
||||
if oldExcluded.hash != newExcluded.hash {
|
||||
changes = append(changes, fmt.Sprintf("claude[%d].excluded-models: updated (%d -> %d entries)", i, oldExcluded.count, newExcluded.count))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1517,9 +1811,18 @@ func buildConfigChangeDetails(oldCfg, newCfg *config.Config) []string {
|
||||
if !equalStringMap(o.Headers, n.Headers) {
|
||||
changes = append(changes, fmt.Sprintf("codex[%d].headers: updated", i))
|
||||
}
|
||||
oldExcluded := summarizeExcludedModels(o.ExcludedModels)
|
||||
newExcluded := summarizeExcludedModels(n.ExcludedModels)
|
||||
if oldExcluded.hash != newExcluded.hash {
|
||||
changes = append(changes, fmt.Sprintf("codex[%d].excluded-models: updated (%d -> %d entries)", i, oldExcluded.count, newExcluded.count))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if entries, _ := diffOAuthExcludedModelChanges(oldCfg.OAuthExcludedModels, newCfg.OAuthExcludedModels); len(entries) > 0 {
|
||||
changes = append(changes, entries...)
|
||||
}
|
||||
|
||||
// Remote management (never print the key)
|
||||
if oldCfg.RemoteManagement.AllowRemote != newCfg.RemoteManagement.AllowRemote {
|
||||
changes = append(changes, fmt.Sprintf("remote-management.allow-remote: %t -> %t", oldCfg.RemoteManagement.AllowRemote, newCfg.RemoteManagement.AllowRemote))
|
||||
|
||||
@@ -1118,6 +1118,14 @@ func (m *Manager) pickNext(ctx context.Context, provider, model string, opts cli
|
||||
}
|
||||
authCopy := selected.Clone()
|
||||
m.mu.RUnlock()
|
||||
if !selected.indexAssigned {
|
||||
m.mu.Lock()
|
||||
if current := m.auths[authCopy.ID]; current != nil && !current.indexAssigned {
|
||||
current.EnsureIndex()
|
||||
authCopy = current.Clone()
|
||||
}
|
||||
m.mu.Unlock()
|
||||
}
|
||||
return authCopy, executor, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -29,7 +29,7 @@ func NewAPIKeyClientProvider() APIKeyClientProvider {
|
||||
type apiKeyClientProvider struct{}
|
||||
|
||||
func (p *apiKeyClientProvider) Load(ctx context.Context, cfg *config.Config) (*APIKeyClientResult, error) {
|
||||
geminiCount, claudeCount, codexCount, openAICompat := watcher.BuildAPIKeyClients(cfg)
|
||||
geminiCount, vertexCompatCount, claudeCount, codexCount, openAICompat := watcher.BuildAPIKeyClients(cfg)
|
||||
if ctx != nil {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
@@ -38,9 +38,10 @@ func (p *apiKeyClientProvider) Load(ctx context.Context, cfg *config.Config) (*A
|
||||
}
|
||||
}
|
||||
return &APIKeyClientResult{
|
||||
GeminiKeyCount: geminiCount,
|
||||
ClaudeKeyCount: claudeCount,
|
||||
CodexKeyCount: codexCount,
|
||||
OpenAICompatCount: openAICompat,
|
||||
GeminiKeyCount: geminiCount,
|
||||
VertexCompatKeyCount: vertexCompatCount,
|
||||
ClaudeKeyCount: claudeCount,
|
||||
CodexKeyCount: codexCount,
|
||||
OpenAICompatCount: openAICompat,
|
||||
}, nil
|
||||
}
|
||||
|
||||
@@ -146,6 +146,27 @@ func (s *Service) consumeAuthUpdates(ctx context.Context) {
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) emitAuthUpdate(ctx context.Context, update watcher.AuthUpdate) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
if s.watcher != nil && s.watcher.DispatchRuntimeAuthUpdate(update) {
|
||||
return
|
||||
}
|
||||
if s.authUpdates != nil {
|
||||
select {
|
||||
case s.authUpdates <- update:
|
||||
return
|
||||
default:
|
||||
log.Debugf("auth update queue saturated, applying inline action=%v id=%s", update.Action, update.ID)
|
||||
}
|
||||
}
|
||||
s.handleAuthUpdate(ctx, update)
|
||||
}
|
||||
|
||||
func (s *Service) handleAuthUpdate(ctx context.Context, update watcher.AuthUpdate) {
|
||||
if s == nil {
|
||||
return
|
||||
@@ -220,7 +241,11 @@ func (s *Service) wsOnConnected(channelID string) {
|
||||
Metadata: map[string]any{"email": channelID}, // metadata drives logging and usage tracking
|
||||
}
|
||||
log.Infof("websocket provider connected: %s", channelID)
|
||||
s.applyCoreAuthAddOrUpdate(context.Background(), auth)
|
||||
s.emitAuthUpdate(context.Background(), watcher.AuthUpdate{
|
||||
Action: watcher.AuthUpdateActionAdd,
|
||||
ID: auth.ID,
|
||||
Auth: auth,
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Service) wsOnDisconnected(channelID string, reason error) {
|
||||
@@ -237,7 +262,10 @@ func (s *Service) wsOnDisconnected(channelID string, reason error) {
|
||||
log.Infof("websocket provider disconnected: %s", channelID)
|
||||
}
|
||||
ctx := context.Background()
|
||||
s.applyCoreAuthRemoval(ctx, channelID)
|
||||
s.emitAuthUpdate(ctx, watcher.AuthUpdate{
|
||||
Action: watcher.AuthUpdateActionDelete,
|
||||
ID: channelID,
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Service) applyCoreAuthAddOrUpdate(ctx context.Context, auth *coreauth.Auth) {
|
||||
@@ -296,7 +324,7 @@ func openAICompatInfoFromAuth(a *coreauth.Auth) (providerKey string, compatName
|
||||
if len(a.Attributes) > 0 {
|
||||
providerKey = strings.TrimSpace(a.Attributes["provider_key"])
|
||||
compatName = strings.TrimSpace(a.Attributes["compat_name"])
|
||||
if providerKey != "" || compatName != "" {
|
||||
if compatName != "" {
|
||||
if providerKey == "" {
|
||||
providerKey = compatName
|
||||
}
|
||||
@@ -470,7 +498,7 @@ func (s *Service) Run(ctx context.Context) error {
|
||||
}()
|
||||
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
fmt.Println("API server started successfully")
|
||||
fmt.Printf("API server started successfully on: %d\n", s.cfg.Port)
|
||||
|
||||
if s.hooks.OnAfterStart != nil {
|
||||
s.hooks.OnAfterStart(s)
|
||||
@@ -617,6 +645,7 @@ func (s *Service) registerModelsForAuth(a *coreauth.Auth) {
|
||||
if a == nil || a.ID == "" {
|
||||
return
|
||||
}
|
||||
authKind := strings.ToLower(strings.TrimSpace(a.Attributes["auth_kind"]))
|
||||
if a.Attributes != nil {
|
||||
if v := strings.TrimSpace(a.Attributes["gemini_virtual_primary"]); strings.EqualFold(v, "true") {
|
||||
GlobalModelRegistry().UnregisterClient(a.ID)
|
||||
@@ -636,32 +665,62 @@ func (s *Service) registerModelsForAuth(a *coreauth.Auth) {
|
||||
if compatDetected {
|
||||
provider = "openai-compatibility"
|
||||
}
|
||||
excluded := s.oauthExcludedModels(provider, authKind)
|
||||
var models []*ModelInfo
|
||||
switch provider {
|
||||
case "gemini":
|
||||
models = registry.GetGeminiModels()
|
||||
if entry := s.resolveConfigGeminiKey(a); entry != nil {
|
||||
if authKind == "apikey" {
|
||||
excluded = entry.ExcludedModels
|
||||
}
|
||||
}
|
||||
models = applyExcludedModels(models, excluded)
|
||||
case "vertex":
|
||||
// Vertex AI Gemini supports the same model identifiers as Gemini.
|
||||
models = registry.GetGeminiVertexModels()
|
||||
if authKind == "apikey" {
|
||||
if entry := s.resolveConfigVertexCompatKey(a); entry != nil && len(entry.Models) > 0 {
|
||||
models = buildVertexCompatConfigModels(entry)
|
||||
}
|
||||
}
|
||||
models = applyExcludedModels(models, excluded)
|
||||
case "gemini-cli":
|
||||
models = registry.GetGeminiCLIModels()
|
||||
models = applyExcludedModels(models, excluded)
|
||||
case "aistudio":
|
||||
models = registry.GetAIStudioModels()
|
||||
models = applyExcludedModels(models, excluded)
|
||||
case "antigravity":
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
|
||||
models = executor.FetchAntigravityModels(ctx, a, s.cfg)
|
||||
cancel()
|
||||
models = applyExcludedModels(models, excluded)
|
||||
case "claude":
|
||||
models = registry.GetClaudeModels()
|
||||
if entry := s.resolveConfigClaudeKey(a); entry != nil && len(entry.Models) > 0 {
|
||||
models = buildClaudeConfigModels(entry)
|
||||
if entry := s.resolveConfigClaudeKey(a); entry != nil {
|
||||
if len(entry.Models) > 0 {
|
||||
models = buildClaudeConfigModels(entry)
|
||||
}
|
||||
if authKind == "apikey" {
|
||||
excluded = entry.ExcludedModels
|
||||
}
|
||||
}
|
||||
models = applyExcludedModels(models, excluded)
|
||||
case "codex":
|
||||
models = registry.GetOpenAIModels()
|
||||
if entry := s.resolveConfigCodexKey(a); entry != nil {
|
||||
if authKind == "apikey" {
|
||||
excluded = entry.ExcludedModels
|
||||
}
|
||||
}
|
||||
models = applyExcludedModels(models, excluded)
|
||||
case "qwen":
|
||||
models = registry.GetQwenModels()
|
||||
models = applyExcludedModels(models, excluded)
|
||||
case "iflow":
|
||||
models = registry.GetIFlowModels()
|
||||
models = applyExcludedModels(models, excluded)
|
||||
default:
|
||||
// Handle OpenAI-compatibility providers by name using config
|
||||
if s.cfg != nil {
|
||||
@@ -749,7 +808,10 @@ func (s *Service) registerModelsForAuth(a *coreauth.Auth) {
|
||||
key = strings.ToLower(strings.TrimSpace(a.Provider))
|
||||
}
|
||||
GlobalModelRegistry().RegisterClient(a.ID, key, models)
|
||||
return
|
||||
}
|
||||
|
||||
GlobalModelRegistry().UnregisterClient(a.ID)
|
||||
}
|
||||
|
||||
func (s *Service) resolveConfigClaudeKey(auth *coreauth.Auth) *config.ClaudeKey {
|
||||
@@ -791,6 +853,222 @@ func (s *Service) resolveConfigClaudeKey(auth *coreauth.Auth) *config.ClaudeKey
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Service) resolveConfigGeminiKey(auth *coreauth.Auth) *config.GeminiKey {
|
||||
if auth == nil || s.cfg == nil {
|
||||
return nil
|
||||
}
|
||||
var attrKey, attrBase string
|
||||
if auth.Attributes != nil {
|
||||
attrKey = strings.TrimSpace(auth.Attributes["api_key"])
|
||||
attrBase = strings.TrimSpace(auth.Attributes["base_url"])
|
||||
}
|
||||
for i := range s.cfg.GeminiKey {
|
||||
entry := &s.cfg.GeminiKey[i]
|
||||
cfgKey := strings.TrimSpace(entry.APIKey)
|
||||
cfgBase := strings.TrimSpace(entry.BaseURL)
|
||||
if attrKey != "" && strings.EqualFold(cfgKey, attrKey) {
|
||||
if cfgBase == "" || strings.EqualFold(cfgBase, attrBase) {
|
||||
return entry
|
||||
}
|
||||
continue
|
||||
}
|
||||
if attrKey == "" && attrBase != "" && strings.EqualFold(cfgBase, attrBase) {
|
||||
return entry
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Service) resolveConfigVertexCompatKey(auth *coreauth.Auth) *config.VertexCompatKey {
|
||||
if auth == nil || s.cfg == nil {
|
||||
return nil
|
||||
}
|
||||
var attrKey, attrBase string
|
||||
if auth.Attributes != nil {
|
||||
attrKey = strings.TrimSpace(auth.Attributes["api_key"])
|
||||
attrBase = strings.TrimSpace(auth.Attributes["base_url"])
|
||||
}
|
||||
for i := range s.cfg.VertexCompatAPIKey {
|
||||
entry := &s.cfg.VertexCompatAPIKey[i]
|
||||
cfgKey := strings.TrimSpace(entry.APIKey)
|
||||
cfgBase := strings.TrimSpace(entry.BaseURL)
|
||||
if attrKey != "" && strings.EqualFold(cfgKey, attrKey) {
|
||||
if cfgBase == "" || strings.EqualFold(cfgBase, attrBase) {
|
||||
return entry
|
||||
}
|
||||
continue
|
||||
}
|
||||
if attrKey == "" && attrBase != "" && strings.EqualFold(cfgBase, attrBase) {
|
||||
return entry
|
||||
}
|
||||
}
|
||||
if attrKey != "" {
|
||||
for i := range s.cfg.VertexCompatAPIKey {
|
||||
entry := &s.cfg.VertexCompatAPIKey[i]
|
||||
if strings.EqualFold(strings.TrimSpace(entry.APIKey), attrKey) {
|
||||
return entry
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Service) resolveConfigCodexKey(auth *coreauth.Auth) *config.CodexKey {
|
||||
if auth == nil || s.cfg == nil {
|
||||
return nil
|
||||
}
|
||||
var attrKey, attrBase string
|
||||
if auth.Attributes != nil {
|
||||
attrKey = strings.TrimSpace(auth.Attributes["api_key"])
|
||||
attrBase = strings.TrimSpace(auth.Attributes["base_url"])
|
||||
}
|
||||
for i := range s.cfg.CodexKey {
|
||||
entry := &s.cfg.CodexKey[i]
|
||||
cfgKey := strings.TrimSpace(entry.APIKey)
|
||||
cfgBase := strings.TrimSpace(entry.BaseURL)
|
||||
if attrKey != "" && strings.EqualFold(cfgKey, attrKey) {
|
||||
if cfgBase == "" || strings.EqualFold(cfgBase, attrBase) {
|
||||
return entry
|
||||
}
|
||||
continue
|
||||
}
|
||||
if attrKey == "" && attrBase != "" && strings.EqualFold(cfgBase, attrBase) {
|
||||
return entry
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Service) oauthExcludedModels(provider, authKind string) []string {
|
||||
cfg := s.cfg
|
||||
if cfg == nil {
|
||||
return nil
|
||||
}
|
||||
authKindKey := strings.ToLower(strings.TrimSpace(authKind))
|
||||
providerKey := strings.ToLower(strings.TrimSpace(provider))
|
||||
if authKindKey == "apikey" {
|
||||
return nil
|
||||
}
|
||||
return cfg.OAuthExcludedModels[providerKey]
|
||||
}
|
||||
|
||||
func applyExcludedModels(models []*ModelInfo, excluded []string) []*ModelInfo {
|
||||
if len(models) == 0 || len(excluded) == 0 {
|
||||
return models
|
||||
}
|
||||
|
||||
patterns := make([]string, 0, len(excluded))
|
||||
for _, item := range excluded {
|
||||
if trimmed := strings.TrimSpace(item); trimmed != "" {
|
||||
patterns = append(patterns, strings.ToLower(trimmed))
|
||||
}
|
||||
}
|
||||
if len(patterns) == 0 {
|
||||
return models
|
||||
}
|
||||
|
||||
filtered := make([]*ModelInfo, 0, len(models))
|
||||
for _, model := range models {
|
||||
if model == nil {
|
||||
continue
|
||||
}
|
||||
modelID := strings.ToLower(strings.TrimSpace(model.ID))
|
||||
blocked := false
|
||||
for _, pattern := range patterns {
|
||||
if matchWildcard(pattern, modelID) {
|
||||
blocked = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !blocked {
|
||||
filtered = append(filtered, model)
|
||||
}
|
||||
}
|
||||
return filtered
|
||||
}
|
||||
|
||||
// matchWildcard performs case-insensitive wildcard matching where '*' matches any substring.
|
||||
func matchWildcard(pattern, value string) bool {
|
||||
if pattern == "" {
|
||||
return false
|
||||
}
|
||||
|
||||
// Fast path for exact match (no wildcard present).
|
||||
if !strings.Contains(pattern, "*") {
|
||||
return pattern == value
|
||||
}
|
||||
|
||||
parts := strings.Split(pattern, "*")
|
||||
// Handle prefix.
|
||||
if prefix := parts[0]; prefix != "" {
|
||||
if !strings.HasPrefix(value, prefix) {
|
||||
return false
|
||||
}
|
||||
value = value[len(prefix):]
|
||||
}
|
||||
|
||||
// Handle suffix.
|
||||
if suffix := parts[len(parts)-1]; suffix != "" {
|
||||
if !strings.HasSuffix(value, suffix) {
|
||||
return false
|
||||
}
|
||||
value = value[:len(value)-len(suffix)]
|
||||
}
|
||||
|
||||
// Handle middle segments in order.
|
||||
for i := 1; i < len(parts)-1; i++ {
|
||||
segment := parts[i]
|
||||
if segment == "" {
|
||||
continue
|
||||
}
|
||||
idx := strings.Index(value, segment)
|
||||
if idx < 0 {
|
||||
return false
|
||||
}
|
||||
value = value[idx+len(segment):]
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
func buildVertexCompatConfigModels(entry *config.VertexCompatKey) []*ModelInfo {
|
||||
if entry == nil || len(entry.Models) == 0 {
|
||||
return nil
|
||||
}
|
||||
now := time.Now().Unix()
|
||||
out := make([]*ModelInfo, 0, len(entry.Models))
|
||||
seen := make(map[string]struct{}, len(entry.Models))
|
||||
for i := range entry.Models {
|
||||
model := entry.Models[i]
|
||||
name := strings.TrimSpace(model.Name)
|
||||
alias := strings.TrimSpace(model.Alias)
|
||||
if alias == "" {
|
||||
alias = name
|
||||
}
|
||||
if alias == "" {
|
||||
continue
|
||||
}
|
||||
key := strings.ToLower(alias)
|
||||
if _, exists := seen[key]; exists {
|
||||
continue
|
||||
}
|
||||
seen[key] = struct{}{}
|
||||
display := name
|
||||
if display == "" {
|
||||
display = alias
|
||||
}
|
||||
out = append(out, &ModelInfo{
|
||||
ID: alias,
|
||||
Object: "model",
|
||||
Created: now,
|
||||
OwnedBy: "vertex",
|
||||
Type: "vertex",
|
||||
DisplayName: display,
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func buildClaudeConfigModels(entry *config.ClaudeKey) []*ModelInfo {
|
||||
if entry == nil || len(entry.Models) == 0 {
|
||||
return nil
|
||||
|
||||
@@ -49,19 +49,21 @@ type APIKeyClientProvider interface {
|
||||
Load(ctx context.Context, cfg *config.Config) (*APIKeyClientResult, error)
|
||||
}
|
||||
|
||||
// APIKeyClientResult contains API key based clients along with type counts.
|
||||
// It provides metadata about the number of clients loaded for each provider type.
|
||||
// APIKeyClientResult is returned by APIKeyClientProvider.Load()
|
||||
type APIKeyClientResult struct {
|
||||
// GeminiKeyCount is the number of Gemini API key clients loaded.
|
||||
// GeminiKeyCount is the number of Gemini API keys loaded
|
||||
GeminiKeyCount int
|
||||
|
||||
// ClaudeKeyCount is the number of Claude API key clients loaded.
|
||||
// VertexCompatKeyCount is the number of Vertex-compatible API keys loaded
|
||||
VertexCompatKeyCount int
|
||||
|
||||
// ClaudeKeyCount is the number of Claude API keys loaded
|
||||
ClaudeKeyCount int
|
||||
|
||||
// CodexKeyCount is the number of Codex API key clients loaded.
|
||||
// CodexKeyCount is the number of Codex API keys loaded
|
||||
CodexKeyCount int
|
||||
|
||||
// OpenAICompatCount is the number of OpenAI-compatible API key clients loaded.
|
||||
// OpenAICompatCount is the number of OpenAI compatibility API keys loaded
|
||||
OpenAICompatCount int
|
||||
}
|
||||
|
||||
@@ -83,9 +85,10 @@ type WatcherWrapper struct {
|
||||
start func(ctx context.Context) error
|
||||
stop func() error
|
||||
|
||||
setConfig func(cfg *config.Config)
|
||||
snapshotAuths func() []*coreauth.Auth
|
||||
setUpdateQueue func(queue chan<- watcher.AuthUpdate)
|
||||
setConfig func(cfg *config.Config)
|
||||
snapshotAuths func() []*coreauth.Auth
|
||||
setUpdateQueue func(queue chan<- watcher.AuthUpdate)
|
||||
dispatchRuntimeUpdate func(update watcher.AuthUpdate) bool
|
||||
}
|
||||
|
||||
// Start proxies to the underlying watcher Start implementation.
|
||||
@@ -112,6 +115,16 @@ func (w *WatcherWrapper) SetConfig(cfg *config.Config) {
|
||||
w.setConfig(cfg)
|
||||
}
|
||||
|
||||
// DispatchRuntimeAuthUpdate forwards runtime auth updates (e.g., websocket providers)
|
||||
// into the watcher-managed auth update queue when available.
|
||||
// Returns true if the update was enqueued successfully.
|
||||
func (w *WatcherWrapper) DispatchRuntimeAuthUpdate(update watcher.AuthUpdate) bool {
|
||||
if w == nil || w.dispatchRuntimeUpdate == nil {
|
||||
return false
|
||||
}
|
||||
return w.dispatchRuntimeUpdate(update)
|
||||
}
|
||||
|
||||
// SetClients updates the watcher file-backed clients registry.
|
||||
// SetClients and SetAPIKeyClients removed; watcher manages its own caches
|
||||
|
||||
|
||||
@@ -28,5 +28,8 @@ func defaultWatcherFactory(configPath, authDir string, reload func(*config.Confi
|
||||
setUpdateQueue: func(queue chan<- watcher.AuthUpdate) {
|
||||
w.SetAuthUpdateQueue(queue)
|
||||
},
|
||||
dispatchRuntimeUpdate: func(update watcher.AuthUpdate) bool {
|
||||
return w.DispatchRuntimeAuthUpdate(update)
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user