i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-01 17:00:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: 新增 Gemini CLI provider adapter

Browse Source 
Closes #216

Co-authored-by: Entropy.Xu <entropy.xu@cloudhabitatsh.com>
This commit is contained in:
fawney19
2026-03-10 23:15:05 +08:00
parent 85aa66c76d 1e39ab3c2e
commit a816235efb
21 changed files with 1317 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
frontend/src/features/providers/components/ProviderFormDialog.vue
View File

@@ -51,6 +51,9 @@
<SelectItem value="codex">
Codex
</SelectItem>
<SelectItem value="gemini_cli">
Gemini CLI
</SelectItem>
<SelectItem value="kiro">
Kiro
</SelectItem>
@@ -73,7 +76,7 @@
Codex
</SelectItem>
<SelectItem value="gemini_cli">
GeminiCli
Gemini CLI
</SelectItem>
<SelectItem value="kiro">
Kiro

1
frontend/src/utils/oauth-icons.ts
View File

@@ -3,6 +3,7 @@ export const OAUTH_ICONS: Record<string, string> = {
linuxdo: `<svg viewBox="0 0 120 120" xmlns="http://www.w3.org/2000/svg"><clipPath id="ld"><circle cx="60" cy="60" r="47"/></clipPath><circle fill="#f0f0f0" cx="60" cy="60" r="50"/><rect fill="#1c1c1e" clip-path="url(#ld)" x="10" y="10" width="100" height="30"/><rect fill="#f0f0f0" clip-path="url(#ld)" x="10" y="40" width="100" height="40"/><rect fill="#ffb003" clip-path="url(#ld)" x="10" y="80" width="100" height="30"/></svg>`,
github: `<svg viewBox="0 0 24 24" fill="currentColor"><path d="M12 0c-6.626 0-12 5.373-12 12 0 5.302 3.438 9.8 8.207 11.387.599.111.793-.261.793-.577v-2.234c-3.338.726-4.033-1.416-4.033-1.416-.546-1.387-1.333-1.756-1.333-1.756-1.089-.745.083-.729.083-.729 1.205.084 1.839 1.237 1.839 1.237 1.07 1.834 2.807 1.304 3.492.997.107-.775.418-1.305.762-1.604-2.665-.305-5.467-1.334-5.467-5.931 0-1.311.469-2.381 1.236-3.221-.124-.303-.535-1.524.117-3.176 0 0 1.008-.322 3.301 1.23.957-.266 1.983-.399 3.003-.404 1.02.005 2.047.138 3.006.404 2.291-1.552 3.297-1.23 3.297-1.23.653 1.653.242 2.874.118 3.176.77.84 1.235 1.911 1.235 3.221 0 4.609-2.807 5.624-5.479 5.921.43.372.823 1.102.823 2.222v3.293c0 .319.192.694.801.576 4.765-1.589 8.199-6.086 8.199-11.386 0-6.627-5.373-12-12-12z"/></svg>`,
google: `<svg viewBox="0 0 24 24"><path fill="#4285F4" d="M22.56 12.25c0-.78-.07-1.53-.2-2.25H12v4.26h5.92c-.26 1.37-1.04 2.53-2.21 3.31v2.77h3.57c2.08-1.92 3.28-4.74 3.28-8.09z"/><path fill="#34A853" d="M12 23c2.97 0 5.46-.98 7.28-2.66l-3.57-2.77c-.98.66-2.23 1.06-3.71 1.06-2.86 0-5.29-1.93-6.16-4.53H2.18v2.84C3.99 20.53 7.7 23 12 23z"/><path fill="#FBBC05" d="M5.84 14.09c-.22-.66-.35-1.36-.35-2.09s.13-1.43.35-2.09V7.07H2.18C1.43 8.55 1 10.22 1 12s.43 3.45 1.18 4.93l2.85-2.22.81-.62z"/><path fill="#EA4335" d="M12 5.38c1.62 0 3.06.56 4.21 1.64l3.15-3.15C17.45 2.09 14.97 1 12 1 7.7 1 3.99 3.47 2.18 7.07l3.66 2.84c.87-2.6 3.3-4.53 6.16-4.53z"/></svg>`,
gemini_cli: `<svg viewBox="0 0 24 24"><path fill="#4285F4" d="M22.56 12.25c0-.78-.07-1.53-.2-2.25H12v4.26h5.92c-.26 1.37-1.04 2.53-2.21 3.31v2.77h3.57c2.08-1.92 3.28-4.74 3.28-8.09z"/><path fill="#34A853" d="M12 23c2.97 0 5.46-.98 7.28-2.66l-3.57-2.77c-.98.66-2.23 1.06-3.71 1.06-2.86 0-5.29-1.93-6.16-4.53H2.18v2.84C3.99 20.53 7.7 23 12 23z"/><path fill="#FBBC05" d="M5.84 14.09c-.22-.66-.35-1.36-.35-2.09s.13-1.43.35-2.09V7.07H2.18C1.43 8.55 1 10.22 1 12s.43 3.45 1.18 4.93l2.85-2.22.81-.62z"/><path fill="#EA4335" d="M12 5.38c1.62 0 3.06.56 4.21 1.64l3.15-3.15C17.45 2.09 14.97 1 12 1 7.7 1 3.99 3.47 2.18 7.07l3.66 2.84c.87-2.6 3.3-4.53 6.16-4.53z"/></svg>`,
}
// Default icon when provider type is not found

1
frontend/src/views/admin/PoolManagement.vue
View File

@@ -1353,6 +1353,7 @@ const poolHeaderMetaText = computed(() => {
const showAccountQuotaColumn = computed(() => {
return selectedProviderType.value === 'codex'
|| selectedProviderType.value === 'gemini_cli'
|| selectedProviderType.value === 'kiro'
|| selectedProviderType.value === 'antigravity'
})

156
src/api/admin/provider_query.py
View File

@@ -255,6 +255,97 @@ class TestModelFailoverResponse(BaseModel):
error: str | None = None
# ============ Internal helpers ============
def _test_check_response_has_error(resp: dict[str, Any]) -> bool:
"""快速判断 check_endpoint 结果是否失败。"""
if resp.get("error"):
return True
if int(resp.get("status_code", 0) or 0) != 200:
return True
resp_data = resp.get("response", {})
if isinstance(resp_data, dict) and resp_data.get("error"):
return True
resp_body = resp_data.get("response_body", {}) if isinstance(resp_data, dict) else {}
parsed = resp_body
if isinstance(resp_body, str):
try:
parsed = json.loads(resp_body)
except (json.JSONDecodeError, ValueError):
pass
return isinstance(parsed, dict) and bool(parsed.get("error"))
def _test_check_extract_error_message(resp: dict[str, Any]) -> str:
"""从 check_endpoint 结果中提取错误消息。"""
resp_data = resp.get("response", {}) if isinstance(resp, dict) else {}
# 直接检查 response dict 中的 error（check_endpoint 返回的结构）
if isinstance(resp_data, dict):
err = resp_data.get("error")
if isinstance(err, dict):
msg = err.get("message")
if isinstance(msg, str):
return msg
if isinstance(err, str):
return err
# 兼容 response_body 包装层
body = resp_data.get("response_body", {}) if isinstance(resp_data, dict) else {}
parsed = body
if isinstance(body, str):
try:
parsed = json.loads(body)
except (json.JSONDecodeError, ValueError):
parsed = body
if isinstance(parsed, dict):
err = parsed.get("error")
if isinstance(err, dict):
msg = err.get("message")
if isinstance(msg, str):
return msg
if isinstance(err, str):
return err
err_raw = resp.get("error")
if isinstance(err_raw, str):
return err_raw
if isinstance(err_raw, dict):
msg = err_raw.get("message")
if isinstance(msg, str):
return msg
return ""
def _test_check_should_fallback_to_non_stream(resp: dict[str, Any]) -> bool:
"""仅在流式特有失败时回退到非流式，避免无意义重试。"""
status = int(resp.get("status_code") or 0)
if status in {404, 405, 415, 501}:
return True
message = _test_check_extract_error_message(resp).lower()
if not message:
return False
stream_markers = ("stream", "sse", "streamgeneratecontent")
unsupported_markers = ("not support", "unsupported", "invalid argument")
if any(token in message for token in stream_markers) and any(
token in message for token in unsupported_markers
):
return True
decode_markers = (
"decompress",
"incorrect header check",
"decoding error",
"content decoding",
"contentdecodingerror",
)
return any(token in message for token in decode_markers)
# ============ API Endpoints ============
@@ -1507,31 +1598,46 @@ async def _execute_test_check(
if not adapter_class:
raise ValueError(f"Unknown API format: {endpoint.api_format}")
response = await adapter_class.check_endpoint(
None,
endpoint.base_url,
api_key_value,
{
**request_payload,
"model": effective_model,
},
extra_headers if extra_headers else None,
body_rules=getattr(endpoint, "body_rules", None),
header_rules=getattr(endpoint, "header_rules", None),
db=db,
user=user,
provider_name=provider_obj.name,
provider_id=str(provider_obj.id),
api_key_id=str(key.id),
model_name=effective_model,
auth_type=auth_type,
provider_type=provider_type if provider_type else None,
decrypted_auth_config=auth_config if auth_config else None,
provider_endpoint=endpoint,
provider_api_key=key,
proxy_config=effective_proxy,
timeout_seconds=request_timeout,
)
async def _run_check(stream: bool) -> dict[str, Any]:
return await adapter_class.check_endpoint(
None,
endpoint.base_url,
api_key_value,
{
**request_payload,
"model": effective_model,
"stream": stream,
},
extra_headers if extra_headers else None,
body_rules=getattr(endpoint, "body_rules", None),
header_rules=getattr(endpoint, "header_rules", None),
db=db,
user=user,
provider_name=provider_obj.name,
provider_id=str(provider_obj.id),
api_key_id=str(key.id),
model_name=effective_model,
auth_type=auth_type,
provider_type=provider_type if provider_type else None,
decrypted_auth_config=auth_config if auth_config else None,
provider_endpoint=endpoint,
provider_api_key=key,
proxy_config=effective_proxy,
timeout_seconds=request_timeout,
)
use_stream = bool(request_payload.get("stream", False))
response = await _run_check(use_stream)
if use_stream and _test_check_response_has_error(response):
if _test_check_should_fallback_to_non_stream(response):
logger.info(
"[test-model-failover] Stream check failed for provider={} model={}, fallback to non-stream",
provider_obj.name,
effective_model,
)
response = await _run_check(False)
return response, auth_type

31
src/api/handlers/base/endpoint_checker.py
View File

@@ -593,6 +593,16 @@ class HttpRequestExecutor:
def __init__(self, timeout: float = 30.0):
self.timeout = timeout
@staticmethod
def _unwrap_gemini_cli_response_wrapper(data: Any) -> Any:
"""对齐 gcli2api：Gemini CLI v1internal 响应可能多一层 `response` 包装。"""
if not isinstance(data, dict):
return data
response_obj = data.get("response")
if "candidates" not in data and isinstance(response_obj, dict):
return response_obj
return data
async def execute(self, request: EndpointCheckRequest) -> EndpointCheckResult:
"""执行HTTP请求（支持流式和非流式响应）"""
start_time = time.time()
@@ -661,6 +671,10 @@ class HttpRequestExecutor:
if response.status_code == 200:
try:
response_data = response.json()
if "/v1internal:" in (request.url or ""):
response_data = self._unwrap_gemini_cli_response_wrapper(
response_data
)
logger.debug(
f"[{request.api_format}] check_endpoint | response | json={_truncate_repr(response_data)}"
)
@@ -690,6 +704,17 @@ class HttpRequestExecutor:
return await ErrorHandler.handle_error(http_error, request)
except Exception as e:
logger.warning(
"[{}] endpoint check exception | provider={} provider_id={} api_key_id={} model={} stream={} local_protocol_error={} error={}",
request.api_format,
request.provider_name,
request.provider_id,
request.api_key_id,
request.model_name,
is_stream,
isinstance(e, httpx.LocalProtocolError) or "LocalProtocolError" in type(e).__name__,
e,
)
return await ErrorHandler.handle_error(e, request)
async def _execute_stream_request(
@@ -733,6 +758,8 @@ class HttpRequestExecutor:
try:
event = json.loads(data_str)
if "/v1internal:" in (request.url or ""):
event = self._unwrap_gemini_cli_response_wrapper(event)
event_type = event.get("type", "")
# OpenAI Responses API 事件
@@ -1476,7 +1503,9 @@ class EndpointCheckOrchestrator:
"""执行端点检查的完整流程"""
logger.info(
f"[{request.api_format}] Starting endpoint check | "
f"provider={request.provider_name}, model={request.model_name}"
f"provider={request.provider_name}, model={request.model_name}, "
f"api_key_id={request.api_key_id}, provider_id={request.provider_id}, "
f"stream={request.is_stream if request.is_stream is not None else request.json_body.get('stream', False)}"
)
# 1. 执行HTTP请求

23
src/api/handlers/base/handler_adapter_base.py
View File

@@ -360,6 +360,7 @@ class HandlerAdapterBase(ApiAdapter):
from src.core.provider_types import ProviderType
is_antigravity = provider_type == ProviderType.ANTIGRAVITY
is_gemini_cli = provider_type == ProviderType.GEMINI_CLI
is_vertex = provider_type == ProviderType.VERTEX_AI
is_kiro = provider_type == ProviderType.KIRO
is_oauth = auth_type == "oauth"
@@ -391,6 +392,12 @@ class HandlerAdapterBase(ApiAdapter):
effective_base_url = ordered_urls[0] if ordered_urls else base_url
path = V1INTERNAL_PATH_TEMPLATE.format(action="generateContent")
url = f"{str(effective_base_url).rstrip('/')}{path}"
elif is_gemini_cli:
from src.services.provider.adapters.gemini_cli.constants import V1INTERNAL_PATH_TEMPLATE
effective_base_url = base_url
path = V1INTERNAL_PATH_TEMPLATE.format(action="generateContent")
url = f"{str(effective_base_url).rstrip('/')}{path}"
elif is_vertex and provider_endpoint is not None and provider_api_key is not None:
from src.services.provider.auth import get_provider_auth
from src.services.provider.transport import build_provider_url
@@ -425,6 +432,12 @@ class HandlerAdapterBase(ApiAdapter):
if is_antigravity:
merged_extra.update(get_v1internal_extra_headers())
elif is_gemini_cli:
from src.services.provider.adapters.gemini_cli.constants import (
get_v1internal_extra_headers,
)
merged_extra.update(get_v1internal_extra_headers())
if is_kiro:
from src.services.provider.adapters.kiro.headers import (
@@ -477,6 +490,16 @@ class HandlerAdapterBase(ApiAdapter):
wrap_v1internal_request,
)
project_id = (decrypted_auth_config or {}).get("project_id", "")
effective_model = model_name or request_data.get("model", "")
body = wrap_v1internal_request(
body,
project_id=project_id,
model=effective_model,
)
elif is_gemini_cli:
from src.services.provider.adapters.gemini_cli.envelope import wrap_v1internal_request
project_id = (decrypted_auth_config or {}).get("project_id", "")
effective_model = model_name or request_data.get("model", "")
body = wrap_v1internal_request(

21
src/api/handlers/gemini/adapter.py
View File

@@ -260,6 +260,7 @@ class GeminiChatAdapter(ChatAdapterBase):
}
is_antigravity = provider_type and provider_type.lower() == "antigravity"
is_gemini_cli = provider_type and provider_type.lower() == "gemini_cli"
is_vertex = provider_type and provider_type.lower() == "vertex_ai"
is_oauth = auth_type == "oauth"
vertex_auth_info: Any | None = None
@@ -276,6 +277,11 @@ class GeminiChatAdapter(ChatAdapterBase):
ag_base = ordered_urls[0] if ordered_urls else base_url
path = V1INTERNAL_PATH_TEMPLATE.format(action="generateContent")
url = f"{str(ag_base).rstrip('/')}{path}"
elif is_gemini_cli:
from src.services.provider.adapters.gemini_cli.constants import V1INTERNAL_PATH_TEMPLATE
path = V1INTERNAL_PATH_TEMPLATE.format(action="generateContent")
url = f"{str(base_url).rstrip('/')}{path}"
elif is_vertex and provider_endpoint is not None and provider_api_key is not None:
# Vertex AI: test-model 必须走统一 provider transport/auth，
# 否则会错误命中普通 Gemini URL（导致 404）。
@@ -306,6 +312,12 @@ class GeminiChatAdapter(ChatAdapterBase):
# Antigravity 需要特定的 User-Agent
merged_extra = dict(extra_headers) if extra_headers else {}
if is_antigravity:
merged_extra.update(get_v1internal_extra_headers())
elif is_gemini_cli:
from src.services.provider.adapters.gemini_cli.constants import (
get_v1internal_extra_headers,
)
merged_extra.update(get_v1internal_extra_headers())
if is_vertex and provider_endpoint is not None and provider_api_key is not None:
headers = dict(merged_extra)
@@ -345,6 +357,15 @@ class GeminiChatAdapter(ChatAdapterBase):
model=effective_model_name,
request_type="endpoint_test",
)
elif is_gemini_cli:
from src.services.provider.adapters.gemini_cli.envelope import wrap_v1internal_request
project_id = (decrypted_auth_config or {}).get("project_id", "")
body = wrap_v1internal_request(
body,
project_id=project_id,
model=effective_model_name,
)
# 应用请求头规则（在请求头构建后应用）
if header_rules:

5
src/config/settings.py
View File

@@ -269,7 +269,10 @@ class Config:
"CLAUDE_CLI_USER_AGENT", "claude-code/1.0.1"
)
self.internal_user_agent_openai_cli = os.getenv("OPENAI_CLI_USER_AGENT", "openai-codex/1.0")
self.internal_user_agent_gemini_cli = os.getenv("GEMINI_CLI_USER_AGENT", "gemini-cli/0.1.0")
self.internal_user_agent_gemini_cli = os.getenv(
"GEMINI_CLI_USER_AGENT",
"GeminiCLI/0.1.5 (Windows; AMD64)",
)
# 邮箱验证配置
# VERIFICATION_CODE_EXPIRE_MINUTES: 验证码有效期（分钟）

4
src/core/provider_oauth_utils.py
View File

@@ -481,11 +481,9 @@ async def _enrich_gemini_cli(
def _bootstrap_auth_enrichers() -> None:
# 简单的内置 enrichers 直接注册
# 简单的内置 enrichers 直接注册（兜底，会被 plugin.register_all() 覆盖）
register_auth_enricher("claude_code", _enrich_claude_code)
register_auth_enricher("gemini_cli", _enrich_gemini_cli)
# Provider-specific enrichers 通过 plugin.register_all() 注册
# (called from envelope.py bootstrap)
_bootstrap_auth_enrichers()

65
src/services/orchestration/error_handler.py
View File

@@ -208,6 +208,13 @@ class ErrorHandlerService:
exception=converted_error,
request_id=request_id,
)
self._sync_gemini_cli_quota_state(
key=key,
provider=provider,
model_name=global_model_id,
error_text=error_response_text,
request_id=request_id,
)
# 所有非客户端错误均失效缓存
if can_invalidate:
@@ -455,3 +462,61 @@ class ErrorHandlerService:
logger.debug("auto cleanup side effect failed for key {}: {}", key_id[:8], exc)
task.add_done_callback(_log_async_error)
def _sync_gemini_cli_quota_state(
self,
*,
key: ProviderAPIKey | None,
provider: Provider | None,
model_name: str | None,
error_text: str | None,
request_id: str | None,
) -> None:
if key is None or provider is None:
return
from src.core.provider_types import ProviderType, normalize_provider_type
provider_type = normalize_provider_type(getattr(provider, "provider_type", None))
if provider_type != ProviderType.GEMINI_CLI:
return
normalized_model = str(model_name or "").strip()
if not normalized_model:
return
try:
from src.services.model.upstream_fetcher import merge_upstream_metadata
from src.services.provider.adapters.gemini_cli.quota import (
build_quota_exhausted_metadata,
extract_error_model_name,
)
resolved_model = extract_error_model_name(error_text, fallback=normalized_model)
if not resolved_model:
return
current_metadata = (
key.upstream_metadata if isinstance(key.upstream_metadata, dict) else {}
)
current_namespace = current_metadata.get("gemini_cli")
namespace_dict = current_namespace if isinstance(current_namespace, dict) else None
updates = build_quota_exhausted_metadata(
model_name=resolved_model,
error_text=error_text,
current_namespace=namespace_dict,
)
if not updates:
return
key.upstream_metadata = merge_upstream_metadata(current_metadata, updates)
self.db.add(key)
self.db.commit()
logger.info(
" [{}] Gemini CLI key {} 记录模型冷却: {}",
request_id,
str(getattr(key, "id", "") or "")[:8],
resolved_model,
)
except Exception as exc:
logger.debug(" [{}] Gemini CLI 冷却元数据写入失败: {}", request_id, exc)

5
src/services/provider/adapters/gemini_cli/__init__.py Normal file
View File

@@ -0,0 +1,5 @@
"""GeminiCLI provider adapter package."""
from .plugin import register_all
__all__ = ["register_all"]

192
src/services/provider/adapters/gemini_cli/client.py Normal file
View File

@@ -0,0 +1,192 @@
"""GeminiCLI upstream client helpers."""
from __future__ import annotations
import asyncio
from typing import Any
from src.clients.http_client import HTTPClientPool
from src.core.logger import logger
from src.services.provider.adapters.gemini_cli.constants import (
PROD_BASE_URL,
get_v1internal_extra_headers,
)
_CODE_ASSIST_METADATA = {
"ideType": "ANTIGRAVITY",
"platform": "PLATFORM_UNSPECIFIED",
"pluginType": "GEMINI",
}
def _extract_tier_raw(tier_obj: Any) -> str:
"""Extract raw tier string from loadCodeAssist response objects."""
if isinstance(tier_obj, str) and tier_obj.strip():
return tier_obj.strip()
if isinstance(tier_obj, dict):
for key in ("id", "tierType"):
value = tier_obj.get(key)
if isinstance(value, str) and value.strip():
return value.strip()
return ""
def extract_plan_type(data: dict[str, Any]) -> str | None:
"""Best-effort normalized plan type for GeminiCLI OAuth accounts."""
from src.core.oauth_plan import normalize_oauth_plan_type
for key in ("paidTier", "currentTier"):
raw = _extract_tier_raw(data.get(key))
normalized = normalize_oauth_plan_type(raw)
if normalized:
return normalized
return None
def extract_project_id(data: dict[str, Any]) -> str:
"""Extract project_id from loadCodeAssist/onboardUser responses."""
raw = data.get("cloudaicompanionProject")
if isinstance(raw, str) and raw.strip():
return raw.strip()
if isinstance(raw, dict):
project_id = raw.get("id", "")
if isinstance(project_id, str) and project_id.strip():
return project_id.strip()
return ""
def extract_tier_id(data: dict[str, Any]) -> str:
"""Choose a tier ID for onboarding when the account is not activated."""
allowed_tiers = data.get("allowedTiers")
if not isinstance(allowed_tiers, list):
return ""
for tier in allowed_tiers:
if isinstance(tier, dict) and tier.get("isDefault") is True:
tier_id = tier.get("id", "")
if isinstance(tier_id, str) and tier_id.strip():
return tier_id.strip()
for tier in allowed_tiers:
if isinstance(tier, dict):
tier_id = tier.get("id", "")
if isinstance(tier_id, str) and tier_id.strip():
return tier_id.strip()
return ""
async def load_code_assist(
access_token: str,
proxy_config: dict[str, Any] | None = None,
*,
timeout_seconds: float = 10.0,
) -> dict[str, Any]:
"""Load GeminiCLI account metadata from v1internal:loadCodeAssist."""
if not access_token:
raise ValueError("missing access_token")
client = await HTTPClientPool.get_proxy_client(proxy_config)
headers = {
"Authorization": f"Bearer {access_token}",
"Content-Type": "application/json",
**get_v1internal_extra_headers(),
}
resp = await client.post(
f"{PROD_BASE_URL.rstrip('/')}/v1internal:loadCodeAssist",
json={"metadata": _CODE_ASSIST_METADATA},
headers=headers,
timeout=timeout_seconds,
)
if 200 <= resp.status_code < 300:
data = resp.json()
return data if isinstance(data, dict) else {}
raise RuntimeError(
f"loadCodeAssist failed: status={resp.status_code} body={resp.text[:200] if resp.text else ''}"
)
async def onboard_user(
access_token: str,
*,
tier_id: str,
proxy_config: dict[str, Any] | None = None,
timeout_seconds: float = 30.0,
max_attempts: int = 5,
poll_interval: float = 2.0,
) -> str:
"""Activate GeminiCLI user and fetch project_id via v1internal:onboardUser."""
if not access_token:
raise ValueError("missing access_token")
if not tier_id:
raise ValueError("missing tier_id")
client = await HTTPClientPool.get_proxy_client(proxy_config)
headers = {
"Authorization": f"Bearer {access_token}",
"Content-Type": "application/json",
**get_v1internal_extra_headers(),
}
body = {
"tierId": tier_id,
"metadata": _CODE_ASSIST_METADATA,
}
url = f"{PROD_BASE_URL.rstrip('/')}/v1internal:onboardUser"
for attempt in range(1, max_attempts + 1):
resp = await client.post(url, json=body, headers=headers, timeout=timeout_seconds)
if not (200 <= resp.status_code < 300):
raise RuntimeError(
f"onboardUser failed: status={resp.status_code} body={resp.text[:200] if resp.text else ''}"
)
data = resp.json()
if not isinstance(data, dict):
raise RuntimeError(f"onboardUser: unexpected response type: {type(data)}")
if data.get("done") is True:
response_data = data.get("response")
if isinstance(response_data, dict):
return extract_project_id(response_data)
return ""
if attempt < max_attempts:
await asyncio.sleep(poll_interval)
raise RuntimeError(f"onboardUser: timeout after {max_attempts} attempts")
async def enrich_project_id(
access_token: str,
proxy_config: dict[str, Any] | None = None,
) -> str | None:
"""Best-effort project_id resolution for GeminiCLI OAuth keys."""
code_assist = await load_code_assist(access_token, proxy_config=proxy_config)
project_id = extract_project_id(code_assist)
if project_id:
return project_id
tier_id = extract_tier_id(code_assist)
if tier_id:
try:
project_id = await onboard_user(
access_token,
tier_id=tier_id,
proxy_config=proxy_config,
)
if project_id:
return project_id
except Exception as exc:
logger.warning("GeminiCLI onboardUser failed: {}", exc)
return None
__all__ = [
"enrich_project_id",
"extract_plan_type",
"extract_project_id",
"extract_tier_id",
"load_code_assist",
"onboard_user",
]

25
src/services/provider/adapters/gemini_cli/constants.py Normal file
View File

@@ -0,0 +1,25 @@
"""GeminiCLI provider constants."""
from __future__ import annotations
from src.config.settings import config
from src.core.provider_templates.fixed_providers import FIXED_PROVIDERS
from src.core.provider_templates.types import ProviderType
PROD_BASE_URL = FIXED_PROVIDERS[ProviderType.GEMINI_CLI].api_base_url
V1INTERNAL_PATH_TEMPLATE = "/v1internal:{action}"
def get_v1internal_extra_headers() -> dict[str, str]:
"""Headers required by GeminiCLI upstream requests."""
return {
"Accept-Encoding": "identity",
"User-Agent": config.internal_user_agent_gemini_cli,
}
__all__ = [
"PROD_BASE_URL",
"V1INTERNAL_PATH_TEMPLATE",
"get_v1internal_extra_headers",
]

96
src/services/provider/adapters/gemini_cli/envelope.py Normal file
View File

@@ -0,0 +1,96 @@
"""GeminiCLI v1internal request envelope."""
from __future__ import annotations
from typing import Any
from src.services.provider.adapters.gemini_cli.constants import get_v1internal_extra_headers
from src.services.provider.request_context import get_selected_base_url
def wrap_v1internal_request(
gemini_request: dict[str, Any],
*,
project_id: str,
model: str,
) -> dict[str, Any]:
"""Wrap a Gemini request into GeminiCLI v1internal format."""
inner_request = dict(gemini_request)
inner_request.pop("model", None)
inner_request.pop("stream", None)
return {
"model": model,
"project": project_id,
"request": inner_request,
}
class GeminiCliV1InternalEnvelope:
name = "gemini_cli:v1internal"
def extra_headers(self) -> dict[str, str] | None:
return get_v1internal_extra_headers()
def wrap_request(
self,
request_body: dict[str, Any],
*,
model: str,
url_model: str | None,
decrypted_auth_config: dict[str, Any] | None,
) -> tuple[dict[str, Any], str | None]:
_ = url_model
project_id = (decrypted_auth_config or {}).get("project_id")
if not isinstance(project_id, str) or not project_id:
from src.core.exceptions import ProviderNotAvailableException
raise ProviderNotAvailableException(
"GeminiCLI OAuth 配置缺少 project_id，请重新授权",
provider_name="gemini_cli",
upstream_response="missing auth_config.project_id",
)
wrapped = wrap_v1internal_request(
request_body,
project_id=project_id,
model=model,
)
return wrapped, None
def unwrap_response(self, data: Any) -> Any:
if not isinstance(data, dict):
return data
response_obj = data.get("response")
if "candidates" not in data and isinstance(response_obj, dict):
return response_obj
return data
def postprocess_unwrapped_response(self, *, model: str, data: Any) -> None:
_ = model, data
return
def capture_selected_base_url(self) -> str | None:
return get_selected_base_url()
def on_http_status(self, *, base_url: str | None, status_code: int) -> None:
_ = base_url, status_code
return
def on_connection_error(self, *, base_url: str | None, exc: Exception) -> None:
_ = base_url, exc
return
def force_stream_rewrite(self) -> bool:
return False
gemini_cli_v1internal_envelope = GeminiCliV1InternalEnvelope()
__all__ = [
"GeminiCliV1InternalEnvelope",
"gemini_cli_v1internal_envelope",
"wrap_v1internal_request",
]

158
src/services/provider/adapters/gemini_cli/plugin.py Normal file
View File

@@ -0,0 +1,158 @@
"""GeminiCLI provider plugin — unified registration entry."""
from __future__ import annotations
import time
from typing import Any
from urllib.parse import urlencode
from src.core.logger import logger
from src.services.provider.adapters.gemini_cli.constants import V1INTERNAL_PATH_TEMPLATE
from src.services.provider.preset_models import get_preset_models
from src.services.provider.request_context import set_selected_base_url
async def fetch_models_gemini_cli(
ctx: Any,
timeout_seconds: float,
) -> tuple[list[dict], list[str], bool, dict[str, Any] | None]:
"""GeminiCLI model fetcher.
Upstream does not expose a stable public models endpoint for OAuth CLI access,
so we currently return a curated preset model catalog and enrich account metadata
from loadCodeAssist when possible.
"""
from src.services.provider.adapters.gemini_cli.client import (
extract_plan_type,
load_code_assist,
)
models = get_preset_models("gemini_cli")
upstream_metadata: dict[str, Any] | None = None
access_token = str(getattr(ctx, "api_key_value", "") or "").strip()
if access_token:
try:
code_assist = await load_code_assist(
access_token,
proxy_config=getattr(ctx, "proxy_config", None),
timeout_seconds=timeout_seconds,
)
provider_meta: dict[str, Any] = {"updated_at": int(time.time())}
plan_type = extract_plan_type(code_assist)
if plan_type:
provider_meta["plan_type"] = plan_type
project_id = (getattr(ctx, "auth_config", None) or {}).get("project_id")
if isinstance(project_id, str) and project_id:
provider_meta["project_id"] = project_id
upstream_metadata = {"gemini_cli": provider_meta}
except Exception as exc:
logger.debug("GeminiCLI model metadata fetch failed: {}", exc)
return models, [], True, upstream_metadata
def build_gemini_cli_url(
endpoint: Any,
*,
is_stream: bool,
effective_query_params: dict[str, Any],
**_kwargs: Any,
) -> str:
"""Build GeminiCLI v1internal URL."""
base_url = str(getattr(endpoint, "base_url", "") or "").rstrip("/")
set_selected_base_url(base_url)
action = "streamGenerateContent" if is_stream else "generateContent"
path = V1INTERNAL_PATH_TEMPLATE.format(action=action)
url = f"{base_url}{path}"
if is_stream:
effective_query_params.setdefault("alt", "sse")
if effective_query_params:
query_string = urlencode(effective_query_params, doseq=True)
if query_string:
url = f"{url}?{query_string}"
return url
async def enrich_gemini_cli(
auth_config: dict[str, Any],
token_response: dict[str, Any],
access_token: str,
proxy_config: dict[str, Any] | None,
) -> dict[str, Any]:
"""GeminiCLI auth_config enrichment: email + project_id."""
from src.core.provider_oauth_utils import fetch_google_email
from src.services.provider.adapters.gemini_cli.client import (
enrich_project_id,
extract_plan_type,
load_code_assist,
)
if not auth_config.get("email"):
email = await fetch_google_email(
access_token,
proxy_config=proxy_config,
timeout_seconds=10.0,
)
if email:
auth_config["email"] = email
try:
code_assist = await load_code_assist(access_token, proxy_config=proxy_config)
except Exception as exc:
code_assist = None
logger.warning("[enrich] GeminiCLI loadCodeAssist failed: {}", exc)
if code_assist and not auth_config.get("tier"):
plan_type = extract_plan_type(code_assist)
if plan_type:
auth_config["tier"] = plan_type
if not auth_config.get("project_id"):
try:
project_id = (code_assist and code_assist.get("cloudaicompanionProject")) or None
if isinstance(project_id, dict):
project_id = project_id.get("id")
if isinstance(project_id, str) and project_id.strip():
auth_config["project_id"] = project_id.strip()
else:
project_id = await enrich_project_id(access_token, proxy_config=proxy_config)
if project_id:
auth_config["project_id"] = project_id
logger.info("[enrich] GeminiCLI project_id: {}", project_id[:8] + "...")
except Exception as exc:
logger.warning("[enrich] GeminiCLI project_id enrichment failed: {}", exc)
return auth_config
def register_all() -> None:
"""Register all GeminiCLI hooks into shared registries."""
from src.core.provider_oauth_utils import register_auth_enricher
from src.services.model.upstream_fetcher import UpstreamModelsFetcherRegistry
from src.services.provider.adapters.gemini_cli.envelope import gemini_cli_v1internal_envelope
from src.services.provider.envelope import register_envelope
from src.services.provider.transport import register_transport_hook
register_envelope("gemini_cli", "gemini:cli", gemini_cli_v1internal_envelope)
register_envelope("gemini_cli", "gemini:chat", gemini_cli_v1internal_envelope)
register_envelope("gemini_cli", "", gemini_cli_v1internal_envelope)
register_transport_hook("gemini_cli", "gemini:cli", build_gemini_cli_url)
register_transport_hook("gemini_cli", "gemini:chat", build_gemini_cli_url)
register_auth_enricher("gemini_cli", enrich_gemini_cli)
UpstreamModelsFetcherRegistry.register(
provider_types=["gemini_cli"],
fetcher=fetch_models_gemini_cli,
)
__all__ = [
"build_gemini_cli_url",
"enrich_gemini_cli",
"fetch_models_gemini_cli",
"register_all",
]

275
src/services/provider/adapters/gemini_cli/quota.py Normal file
View File

@@ -0,0 +1,275 @@
"""Gemini CLI quota / RESOURCE_EXHAUSTED helpers."""
from __future__ import annotations
import json
import re
import time
from datetime import datetime, timezone
from typing import Any
from uuid import UUID
_DURATION_TOKEN_RE = re.compile(r"(\d+(?:\.\d+)?)([dhms])")
_RESET_AFTER_RE = re.compile(r"reset after\s+([^.,;]+)", re.IGNORECASE)
def _parse_json(error_text: str | None) -> dict[str, Any] | None:
if not isinstance(error_text, str) or not error_text.strip():
return None
try:
data = json.loads(error_text)
except Exception:
return None
return data if isinstance(data, dict) else None
def _parse_duration_seconds(raw: Any) -> int | None:
if isinstance(raw, (int, float)):
return max(1, int(raw))
if not isinstance(raw, str):
return None
text = raw.strip().lower()
if not text:
return None
total_seconds = 0.0
matched = False
for amount_text, unit in _DURATION_TOKEN_RE.findall(text):
matched = True
amount = float(amount_text)
if unit == "d":
total_seconds += amount * 86400
elif unit == "h":
total_seconds += amount * 3600
elif unit == "m":
total_seconds += amount * 60
elif unit == "s":
total_seconds += amount
if not matched:
return None
return max(1, int(total_seconds))
def _iter_error_details(payload: dict[str, Any]) -> list[dict[str, Any]]:
error_obj = payload.get("error")
if not isinstance(error_obj, dict):
return []
details = error_obj.get("details")
if not isinstance(details, list):
return []
return [item for item in details if isinstance(item, dict)]
def _error_status(payload: dict[str, Any]) -> str:
error_obj = payload.get("error")
if not isinstance(error_obj, dict):
return ""
status = error_obj.get("status")
return status.strip() if isinstance(status, str) else ""
def _error_message(payload: dict[str, Any]) -> str:
error_obj = payload.get("error")
if not isinstance(error_obj, dict):
return ""
message = error_obj.get("message")
return message.strip() if isinstance(message, str) else ""
def _error_reason(payload: dict[str, Any]) -> str:
for detail in _iter_error_details(payload):
reason = detail.get("reason")
if isinstance(reason, str) and reason.strip():
return reason.strip()
return ""
def _looks_like_uuid(value: str | None) -> bool:
if not isinstance(value, str):
return False
text = value.strip()
if not text:
return False
try:
UUID(text)
except Exception:
return False
return True
def is_resource_exhausted_error(error_text: str | None) -> bool:
payload = _parse_json(error_text)
if payload is None:
return False
status = _error_status(payload).upper()
reason = _error_reason(payload).upper()
if status == "RESOURCE_EXHAUSTED" or reason == "QUOTA_EXHAUSTED":
return True
message = _error_message(payload).lower()
return ("exhausted your capacity" in message) or ("quota" in message and "exhaust" in message)
def parse_quota_reset_timestamp(error_text: str | None) -> int | None:
payload = _parse_json(error_text)
if payload is None:
return None
for detail in _iter_error_details(payload):
metadata = detail.get("metadata")
if not isinstance(metadata, dict):
continue
raw = metadata.get("quotaResetTimeStamp") or metadata.get("quotaResetTimestamp")
if not isinstance(raw, str) or not raw.strip():
continue
text = raw.strip()
try:
if text.endswith("Z"):
text = text[:-1] + "+00:00"
parsed = datetime.fromisoformat(text)
if parsed.tzinfo is None:
parsed = parsed.replace(tzinfo=timezone.utc)
return int(parsed.timestamp())
except Exception:
continue
return None
def parse_quota_reset_delay_seconds(error_text: str | None) -> int | None:
payload = _parse_json(error_text)
if payload is None:
return None
for detail in _iter_error_details(payload):
metadata = detail.get("metadata")
if not isinstance(metadata, dict):
continue
delay = metadata.get("quotaResetDelay")
parsed = _parse_duration_seconds(delay)
if parsed is not None:
return parsed
return None
def parse_quota_reset_message_seconds(error_text: str | None) -> int | None:
payload = _parse_json(error_text)
if payload is None:
return None
message = _error_message(payload)
if not message:
return None
matched = _RESET_AFTER_RE.search(message)
if not matched:
return None
return _parse_duration_seconds(matched.group(1))
def extract_error_model_name(error_text: str | None, *, fallback: str | None = None) -> str | None:
payload = _parse_json(error_text)
if payload is not None:
for detail in _iter_error_details(payload):
metadata = detail.get("metadata")
if not isinstance(metadata, dict):
continue
model = metadata.get("model")
if isinstance(model, str) and model.strip():
return model.strip()
fallback_text = str(fallback or "").strip()
if fallback_text and not _looks_like_uuid(fallback_text):
return fallback_text
return None
def extract_quota_cooldown_seconds(
error_text: str | None, *, now_ts: int | None = None
) -> int | None:
now = int(now_ts or time.time())
reset_at = parse_quota_reset_timestamp(error_text)
if reset_at is not None:
return max(1, reset_at - now)
delay = parse_quota_reset_delay_seconds(error_text)
if delay is not None:
return max(1, delay)
message_delay = parse_quota_reset_message_seconds(error_text)
if message_delay is not None:
return max(1, message_delay)
return None
def build_quota_exhausted_metadata(
*,
model_name: str,
error_text: str | None,
current_namespace: dict[str, Any] | None = None,
now_ts: int | None = None,
) -> dict[str, Any] | None:
normalized_model = str(model_name or "").strip()
if not normalized_model:
return None
if not is_resource_exhausted_error(error_text):
return None
now = int(now_ts or time.time())
reset_at = parse_quota_reset_timestamp(error_text)
if reset_at is None:
delay = parse_quota_reset_delay_seconds(error_text)
if delay is not None:
reset_at = now + delay
if reset_at is None:
message_delay = parse_quota_reset_message_seconds(error_text)
if message_delay is not None:
reset_at = now + message_delay
if reset_at is None:
return None
payload = _parse_json(error_text) or {}
namespace = dict(current_namespace) if isinstance(current_namespace, dict) else {}
quota_by_model_raw = namespace.get("quota_by_model")
quota_by_model = dict(quota_by_model_raw) if isinstance(quota_by_model_raw, dict) else {}
model_entry_raw = quota_by_model.get(normalized_model)
model_entry = dict(model_entry_raw) if isinstance(model_entry_raw, dict) else {}
model_entry["is_exhausted"] = True
model_entry["remaining_fraction"] = 0.0
model_entry["used_percent"] = 100.0
model_entry["updated_at"] = now
model_entry["reset_at"] = reset_at
model_entry["reset_time"] = datetime.fromtimestamp(reset_at, timezone.utc).isoformat()
model_entry["reset_seconds"] = max(0, reset_at - now)
reason = _error_reason(payload) or _error_status(payload) or _error_message(payload)
if reason:
model_entry["reason"] = reason
quota_by_model[normalized_model] = model_entry
namespace["quota_by_model"] = quota_by_model
namespace["updated_at"] = now
status = _error_status(payload)
if status:
namespace["last_error_status"] = status
if reason:
namespace["last_error_reason"] = reason
return {"gemini_cli": namespace}
__all__ = [
"extract_error_model_name",
"build_quota_exhausted_metadata",
"extract_quota_cooldown_seconds",
"is_resource_exhausted_error",
"parse_quota_reset_message_seconds",
"parse_quota_reset_delay_seconds",
"parse_quota_reset_timestamp",
]

2
src/services/provider/envelope.py
View File

@@ -157,12 +157,14 @@ def ensure_providers_bootstrapped() -> None:
register_all as _reg_claude_code,
)
from src.services.provider.adapters.codex.plugin import register_all as _reg_codex
from src.services.provider.adapters.gemini_cli.plugin import register_all as _reg_gemini_cli
from src.services.provider.adapters.kiro.plugin import register_all as _reg_kiro
from src.services.provider.adapters.vertex_ai.plugin import register_all as _reg_vertex_ai
_reg_antigravity()
_reg_claude_code()
_reg_codex()
_reg_gemini_cli()
_reg_kiro()
_reg_vertex_ai()

19
src/services/provider/pool/health_policy.py
View File

@@ -107,6 +107,23 @@ def _resolve_transient_cooldown_ttl(
return config.overload_cooldown_seconds
def _parse_google_quota_cooldown(error_body: str | None) -> int | None:
"""Parse Google-specific quota cooldown from error body.
Safe to call for any provider: returns None unless the error body
contains Google-specific fields (quotaResetTimeStamp / quotaResetDelay /
"reset after" message pattern).
"""
if not error_body:
return None
try:
from src.services.provider.adapters.gemini_cli.quota import extract_quota_cooldown_seconds
return extract_quota_cooldown_seconds(error_body)
except Exception:
return None
async def apply_health_policy(
*,
provider_id: str,
@@ -218,6 +235,8 @@ async def _apply(
# --- 429 Rate Limited ----------------------------------------------------
if status_code == 429:
retry_after = _parse_retry_after(response_headers)
if retry_after is None:
retry_after = _parse_google_quota_cooldown(error_body)
ttl = _resolve_transient_cooldown_ttl(
status_code=status_code,
retry_after_seconds=retry_after,

38
src/services/provider/preset_models.py
View File

@@ -27,6 +27,44 @@ ModelsFetcherFunc = Callable[
# 格式: provider_type -> list of model dicts
PRESET_MODELS: dict[str, list[dict[str, Any]]] = {
# GeminiCLI (Google Cloud Code / Gemini CLI OAuth upstream)
"gemini_cli": [
{
"id": "gemini-2.5-pro",
"object": "model",
"owned_by": "google",
"display_name": "Gemini 2.5 Pro",
"api_format": "gemini:cli",
},
{
"id": "gemini-2.5-flash",
"object": "model",
"owned_by": "google",
"display_name": "Gemini 2.5 Flash",
"api_format": "gemini:cli",
},
{
"id": "gemini-3-pro-preview",
"object": "model",
"owned_by": "google",
"display_name": "Gemini 3 Pro Preview",
"api_format": "gemini:cli",
},
{
"id": "gemini-3-flash-preview",
"object": "model",
"owned_by": "google",
"display_name": "Gemini 3 Flash Preview",
"api_format": "gemini:cli",
},
{
"id": "gemini-3.1-pro-preview",
"object": "model",
"owned_by": "google",
"display_name": "Gemini 3.1 Pro Preview",
"api_format": "gemini:cli",
},
],
# Kiro (Claude CLI 反代)
"kiro": [
{

109
src/services/provider_keys/quota_reader.py
View File

@@ -402,8 +402,116 @@ class AntigravityQuotaReader(PoolQuotaReader):
return f"最低剩余 {_format_percent(min_remaining)} ({len(remaining_list)} 模型)"
class GeminiCliQuotaReader(PoolQuotaReader):
namespace = "gemini_cli"
def _quota_by_model(self) -> dict[str, Any]:
quota_by_model = self._data.get("quota_by_model")
if not isinstance(quota_by_model, dict):
return {}
return quota_by_model
def _reset_at(self, model_info: dict[str, Any]) -> int | None:
reset_at = _to_float(model_info.get("reset_at"))
if reset_at is None or reset_at <= 0:
return None
if reset_at > 1_000_000_000_000:
reset_at /= 1000
return int(reset_at)
def _is_model_exhausted(self, model_info: dict[str, Any]) -> bool:
if _is_truthy_flag(model_info.get("is_exhausted")):
return True
remaining_fraction = _to_float(model_info.get("remaining_fraction"))
if remaining_fraction is not None and remaining_fraction <= 0.0:
return True
return _pct_is_exhausted(model_info.get("used_percent"))
def _active_exhausted_models(self) -> list[tuple[str, dict[str, Any], int | None]]:
now = int(time.time())
active: list[tuple[str, dict[str, Any], int | None]] = []
for model_name, raw_info in self._quota_by_model().items():
if not isinstance(raw_info, dict):
continue
if not self._is_model_exhausted(raw_info):
continue
reset_at = self._reset_at(raw_info)
if reset_at is not None and reset_at <= now:
continue
active.append((str(model_name), raw_info, reset_at))
return active
def is_exhausted(self, model_name: str | None = None) -> QuotaExhaustedResult:
if not model_name:
return QuotaExhaustedResult(False)
model_quota = self._quota_by_model().get(model_name)
if not isinstance(model_quota, dict) or not self._is_model_exhausted(model_quota):
return QuotaExhaustedResult(False)
reset_at = self._reset_at(model_quota)
if reset_at is not None:
now = int(time.time())
if reset_at <= now:
return QuotaExhaustedResult(False)
reset_text = _format_reset_after(reset_at - now)
if reset_text:
return QuotaExhaustedResult(
True, f"Gemini CLI 模型 {model_name} 冷却中（{reset_text}）"
)
return QuotaExhaustedResult(True, f"Gemini CLI 模型 {model_name} 配额已耗尽")
def usage_ratio(self) -> float | None:
active = self._active_exhausted_models()
if not active:
return None
return 1.0
def plan_type(self) -> str | None:
return _normalize_plan(self._data.get("plan_type")) or _normalize_plan(
self._data.get("tier")
)
def reset_seconds(self) -> float | None:
now = int(time.time())
reset_values = [
reset_at - now
for _, _, reset_at in self._active_exhausted_models()
if reset_at is not None and reset_at > now
]
if not reset_values:
return None
return float(min(reset_values))
def account_block(self) -> AccountBlockResult:
return AccountBlockResult(blocked=False)
def display_summary(self) -> str | None:
active = self._active_exhausted_models()
if not active:
return None
active_sorted = sorted(
active,
key=lambda item: item[2] if item[2] is not None else 2**31 - 1,
)
first_model, _, first_reset_at = active_sorted[0]
if len(active_sorted) == 1:
if first_reset_at is not None:
reset_text = _format_reset_after(first_reset_at - int(time.time()))
if reset_text:
return f"{first_model} 冷却中 ({reset_text})"
return f"{first_model} 冷却中"
if first_reset_at is not None:
reset_text = _format_reset_after(first_reset_at - int(time.time()))
if reset_text:
return f"{len(active_sorted)} 个模型冷却中（最早 {reset_text}）"
return f"{len(active_sorted)} 个模型冷却中"
_READER_CLASSES: dict[str, type[PoolQuotaReader]] = {
ProviderType.CODEX: CodexQuotaReader,
ProviderType.GEMINI_CLI: GeminiCliQuotaReader,
ProviderType.KIRO: KiroQuotaReader,
ProviderType.ANTIGRAVITY: AntigravityQuotaReader,
}
@@ -431,6 +539,7 @@ __all__ = [
"AccountBlockResult",
"AntigravityQuotaReader",
"CodexQuotaReader",
"GeminiCliQuotaReader",
"KiroQuotaReader",
"NullQuotaReader",
"PoolQuotaReader",

117
tests/services/test_gemini_cli_quota.py Normal file
View File

@@ -0,0 +1,117 @@
from __future__ import annotations
from types import SimpleNamespace
from typing import Any, cast
from src.services.orchestration.error_handler import ErrorHandlerService
from src.services.provider.adapters.gemini_cli.quota import (
build_quota_exhausted_metadata,
extract_quota_cooldown_seconds,
)
class _FakeDB:
def __init__(self) -> None:
self.added: list[object] = []
self.commit_count = 0
def add(self, obj: object) -> None:
self.added.append(obj)
def commit(self) -> None:
self.commit_count += 1
def test_extract_quota_cooldown_seconds_from_message_reset_after() -> None:
error_text = """
{
"error": {
"code": 429,
"message": "You have exhausted your capacity on this model. Your quota will reset after 27s.",
"status": "RESOURCE_EXHAUSTED",
"details": [
{
"@type": "type.googleapis.com/google.rpc.ErrorInfo",
"reason": "RATE_LIMIT_EXCEEDED",
"domain": "cloudcode-pa.googleapis.com",
"metadata": {
"model": "gemini-3-pro-preview"
}
}
]
}
}
"""
assert extract_quota_cooldown_seconds(error_text, now_ts=1_700_000_000) == 27
def test_build_quota_exhausted_metadata_skips_unknown_reset_window() -> None:
error_text = """
{
"error": {
"code": 429,
"message": "No capacity available for model gemini-3.1-pro-preview on the server",
"status": "RESOURCE_EXHAUSTED",
"details": [
{
"@type": "type.googleapis.com/google.rpc.ErrorInfo",
"reason": "MODEL_CAPACITY_EXHAUSTED",
"domain": "cloudcode-pa.googleapis.com",
"metadata": {
"model": "gemini-3.1-pro-preview"
}
}
]
}
}
"""
assert (
build_quota_exhausted_metadata(
model_name="gemini-3.1-pro-preview",
error_text=error_text,
current_namespace=None,
now_ts=1_700_000_000,
)
is None
)
def test_sync_gemini_cli_quota_state_uses_error_model_not_global_model_id() -> None:
db = _FakeDB()
service = ErrorHandlerService(db=cast(Any, db))
key = SimpleNamespace(id="k1", upstream_metadata={})
provider = SimpleNamespace(provider_type="gemini_cli")
error_text = """
{
"error": {
"code": 429,
"message": "You have exhausted your capacity on this model. Your quota will reset after 27s.",
"status": "RESOURCE_EXHAUSTED",
"details": [
{
"@type": "type.googleapis.com/google.rpc.ErrorInfo",
"reason": "RATE_LIMIT_EXCEEDED",
"domain": "cloudcode-pa.googleapis.com",
"metadata": {
"model": "gemini-3-pro-preview"
}
}
]
}
}
"""
service._sync_gemini_cli_quota_state(
key=cast(Any, key),
provider=cast(Any, provider),
model_name="17855939-7164-44b3-9354-f87851b25ec6",
error_text=error_text,
request_id="req-1",
)
quota_by_model = key.upstream_metadata["gemini_cli"]["quota_by_model"]
assert "gemini-3-pro-preview" in quota_by_model
assert "17855939-7164-44b3-9354-f87851b25ec6" not in quota_by_model
assert db.commit_count == 1