mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
1d9c77522a
- 删除全部 Python 源码 (src/) 及 Alembic 迁移脚本，归档至 _deprecated_py_src/ - 重构 Rust gateway ai_pipeline: 拆分 planner/finalize 模块，新增 contracts/adaptation 层 - 重组 handlers 模块为 admin/public/proxy/internal/shared 子模块结构 - 新增 executor 模块，引入 Rust 原生数据库迁移 (aether-data/migrations) - 简化 CI/Docker 构建流程，移除 base image 二级构建，统一为单一 app image - 移除 Python 相关基础设施文件 (entrypoint.sh, gunicorn_conf.py, Dockerfile.base)
254 lines
8.8 KiB
Python
254 lines
8.8 KiB
Python
"""Kiro usage/quota fetching utilities."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import re
|
||
import time
|
||
import uuid
|
||
from typing import Any
|
||
|
||
import httpx
|
||
|
||
from src.clients.http_client import HTTPClientPool
|
||
from src.core.logger import logger
|
||
from src.services.provider.adapters.kiro.headers import (
|
||
build_user_agent_usage,
|
||
build_x_amz_user_agent_usage,
|
||
)
|
||
from src.services.provider.adapters.kiro.models.credentials import KiroAuthConfig
|
||
from src.services.provider.adapters.kiro.models.usage_limits import (
|
||
UsageLimitsResponse,
|
||
calculate_current_usage,
|
||
calculate_total_usage_limit,
|
||
)
|
||
from src.services.provider.adapters.kiro.request import get_profile_arn_for_payload
|
||
from src.services.provider.adapters.kiro.rust_http import execute_kiro_rust_http_request
|
||
from src.services.provider.adapters.kiro.token_manager import (
|
||
generate_machine_id,
|
||
is_token_expired,
|
||
refresh_access_token,
|
||
)
|
||
|
||
|
||
class KiroAccountBannedException(Exception):
|
||
"""Kiro 账户被封禁异常"""
|
||
|
||
def __init__(
|
||
self,
|
||
message: str = "账户已封禁",
|
||
status_code: int = 403,
|
||
reason: str | None = None,
|
||
):
|
||
super().__init__(message)
|
||
self.message = message
|
||
self.status_code = status_code
|
||
self.reason = reason
|
||
|
||
|
||
async def fetch_kiro_usage_limits(
|
||
auth_config: dict[str, Any],
|
||
proxy_config: dict[str, Any] | None = None,
|
||
) -> dict[str, Any]:
|
||
"""
|
||
调用 Kiro getUsageLimits API 获取使用额度信息
|
||
|
||
Args:
|
||
auth_config: 解密后的 KiroAuthConfig 数据
|
||
proxy_config: 代理配置（可选）
|
||
|
||
Returns:
|
||
包含 usage_data 和 updated_auth_config 的字典
|
||
|
||
Raises:
|
||
RuntimeError: 请求失败时抛出
|
||
"""
|
||
cfg = KiroAuthConfig.from_dict(auth_config)
|
||
|
||
# 检查是否有缓存的 access_token 且未过期
|
||
access_token: str | None = None
|
||
updated_cfg: KiroAuthConfig | None = None
|
||
|
||
if cfg.access_token and not is_token_expired(cfg.expires_at):
|
||
# 使用缓存的 token
|
||
access_token = cfg.access_token
|
||
updated_cfg = cfg
|
||
logger.debug("[KIRO_QUOTA] 使用缓存的 access_token")
|
||
else:
|
||
# token 过期或不存在，需要刷新
|
||
logger.debug("[KIRO_QUOTA] Token 已过期或不存在，正在刷新...")
|
||
access_token, updated_cfg = await refresh_access_token(cfg, proxy_config=proxy_config)
|
||
|
||
if not access_token:
|
||
raise RuntimeError("无法获取 Kiro access_token")
|
||
|
||
# 构建请求
|
||
effective_cfg = updated_cfg or cfg
|
||
region = effective_cfg.effective_api_region()
|
||
host = f"q.{region}.amazonaws.com"
|
||
machine_id = generate_machine_id(effective_cfg)
|
||
kiro_version = (effective_cfg.kiro_version or "0.8.0").strip() or "0.8.0"
|
||
|
||
# 构建 URL（添加 isEmailRequired=true 获取邮箱）
|
||
url = f"https://{host}/getUsageLimits?origin=AI_EDITOR&resourceType=AGENTIC_REQUEST&isEmailRequired=true"
|
||
|
||
profile_arn = get_profile_arn_for_payload(effective_cfg)
|
||
if profile_arn:
|
||
from urllib.parse import quote
|
||
|
||
url += f"&profileArn={quote(profile_arn, safe='')}"
|
||
|
||
logger.debug("[KIRO_QUOTA] 请求 URL: {}", url)
|
||
|
||
# 构建 headers
|
||
headers = {
|
||
"x-amz-user-agent": build_x_amz_user_agent_usage(
|
||
kiro_version=kiro_version, machine_id=machine_id
|
||
),
|
||
"User-Agent": build_user_agent_usage(kiro_version=kiro_version, machine_id=machine_id),
|
||
"host": host,
|
||
"amz-sdk-invocation-id": str(uuid.uuid4()),
|
||
"amz-sdk-request": "attempt=1; max=1",
|
||
"Authorization": f"Bearer {access_token}",
|
||
"Connection": "close",
|
||
}
|
||
|
||
response = await execute_kiro_rust_http_request(
|
||
method="GET",
|
||
url=url,
|
||
headers=headers,
|
||
body=None,
|
||
proxy_config=proxy_config,
|
||
request_id=f"kiro-usage:{region}:{machine_id}",
|
||
provider_api_format="kiro:usage",
|
||
)
|
||
if response is None:
|
||
client = await HTTPClientPool.get_proxy_client(proxy_config=proxy_config)
|
||
response = await client.get(url, headers=headers, timeout=httpx.Timeout(30.0))
|
||
|
||
if response.status_code != 200:
|
||
response_text = (response.text or "").strip()
|
||
|
||
# 检测账户封禁（403/423 均视为封禁）
|
||
# 对于 Kiro，403 表示账户权限被拒绝（无论是封禁、权限不足还是其他原因），
|
||
# 都应该标记为异常状态，以便管理员及时处理
|
||
is_banned = False
|
||
ban_reason = None
|
||
if response.status_code in (403, 423):
|
||
is_banned = True
|
||
# 检测封禁相关错误（用于提供更详细的原因说明）
|
||
banned_keywords = [
|
||
"AccountSuspendedException",
|
||
"account.*suspend",
|
||
"account.*banned",
|
||
"account.*disabled",
|
||
"account.*access.*denied",
|
||
]
|
||
for keyword in banned_keywords:
|
||
if re.search(keyword, response_text, re.IGNORECASE):
|
||
ban_reason = (
|
||
response_text[:200] if response_text else f"HTTP {response.status_code}"
|
||
)
|
||
break
|
||
# 如果没有匹配到特定关键词，使用通用原因
|
||
if not ban_reason:
|
||
if response.status_code == 423:
|
||
ban_reason = response_text[:200] if response_text else "HTTP 423 Locked"
|
||
else:
|
||
ban_reason = response_text[:200] if response_text else "HTTP 403 权限被拒绝"
|
||
|
||
error_msg = {
|
||
401: "认证失败，Token 无效或已过期",
|
||
403: "账户异常，权限被拒绝",
|
||
423: "账户已封禁",
|
||
429: "请求过于频繁，已被限流",
|
||
}.get(response.status_code, "获取使用额度失败")
|
||
if 500 <= response.status_code < 600:
|
||
error_msg = "服务器错误，AWS 服务暂时不可用"
|
||
logger.debug(
|
||
"kiro usage API error: HTTP {} | {}",
|
||
response.status_code,
|
||
response_text[:200],
|
||
)
|
||
|
||
# 如果检测到封禁，抛出带有封禁标记的异常
|
||
if is_banned:
|
||
raise KiroAccountBannedException(
|
||
message=error_msg,
|
||
status_code=response.status_code,
|
||
reason=ban_reason,
|
||
)
|
||
|
||
raise RuntimeError(f"{error_msg}: HTTP {response.status_code}")
|
||
|
||
try:
|
||
data = response.json()
|
||
except Exception as exc:
|
||
raise RuntimeError(f"获取使用额度成功但响应解析失败: HTTP {response.status_code}") from exc
|
||
|
||
# 返回刷新后的配置（用于更新 auth_config）
|
||
return {
|
||
"usage_data": data,
|
||
"updated_auth_config": updated_cfg.to_dict() if updated_cfg else None,
|
||
}
|
||
|
||
|
||
def parse_kiro_usage_response(data: dict) -> dict | None:
|
||
"""
|
||
解析 Kiro getUsageLimits API 响应，提取限额信息和用户邮箱
|
||
|
||
返回格式与 kiro.rs BalanceResponse 类似：
|
||
- subscription_title: 订阅类型（如 "KIRO PRO+"）
|
||
- current_usage: 当前使用量
|
||
- usage_limit: 使用限额
|
||
- remaining: 剩余额度
|
||
- usage_percentage: 使用百分比
|
||
- next_reset_at: 下次重置时间（Unix 时间戳）
|
||
- email: 用户邮箱（通过 isEmailRequired=true 获取）
|
||
"""
|
||
if not data:
|
||
return None
|
||
|
||
usage_resp = UsageLimitsResponse.from_dict(data)
|
||
|
||
current_usage = calculate_current_usage(usage_resp)
|
||
usage_limit = calculate_total_usage_limit(usage_resp)
|
||
remaining = max(usage_limit - current_usage, 0.0)
|
||
usage_percentage = (current_usage / usage_limit * 100.0) if usage_limit > 0 else 0.0
|
||
usage_percentage = min(usage_percentage, 100.0)
|
||
|
||
result: dict[str, Any] = {
|
||
"current_usage": current_usage,
|
||
"usage_limit": usage_limit,
|
||
"remaining": remaining,
|
||
"usage_percentage": usage_percentage,
|
||
}
|
||
|
||
# 订阅类型
|
||
if usage_resp.subscription_info and usage_resp.subscription_info.subscription_title:
|
||
result["subscription_title"] = usage_resp.subscription_info.subscription_title
|
||
|
||
# 下次重置时间
|
||
if usage_resp.next_date_reset is not None:
|
||
result["next_reset_at"] = usage_resp.next_date_reset
|
||
elif usage_resp.usage_breakdown_list and usage_resp.usage_breakdown_list[0].next_date_reset:
|
||
result["next_reset_at"] = usage_resp.usage_breakdown_list[0].next_date_reset
|
||
|
||
# 解析用户邮箱（从 desktopUserInfo 或 userInfo 中获取）
|
||
user_info = data.get("desktopUserInfo") or data.get("userInfo") or {}
|
||
if isinstance(user_info, dict):
|
||
email = user_info.get("email")
|
||
if isinstance(email, str) and email.strip():
|
||
result["email"] = email.strip()
|
||
|
||
# 添加更新时间戳
|
||
result["updated_at"] = int(time.time())
|
||
|
||
return result
|
||
|
||
|
||
__all__ = [
|
||
"KiroAccountBannedException",
|
||
"fetch_kiro_usage_limits",
|
||
"parse_kiro_usage_response",
|
||
]
|