mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
refactor: 移除 Python 后端源码，全面迁移至 Rust gateway 架构
- 删除全部 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)
This commit is contained in:
fawney19
253
_deprecated_py_src/services/provider/adapters/kiro/usage.py
Normal file
253
_deprecated_py_src/services/provider/adapters/kiro/usage.py
Normal file
@@ -0,0 +1,253 @@
|
||||
"""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",
|
||||
]
|
||||
Reference in New Issue
Block a user