mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
feat: 新增 Kiro 适配器、OAuth 改进与多项功能增强
- 新增 Kiro provider 适配器（EventStream 协议解析、令牌管理、用量追踪） - 重构 OAuth 账户管理与统一配额机制 - 重构 Handler 基类（CLI adapter/handler、请求构建器、流处理器） - 增强缓存监控后端 API 与前端可视化 - 改进 Gemini 格式标准化器与请求头处理 - Antigravity/Codex 适配器更新，移除旧 metadata_collector - 新增数据库迁移：proxy provider API keys - 前端 UI 多项优化 Co-Authored-By: AAEE86 <ppk0227@hotmail.com>
This commit is contained in:
fawney19
@@ -78,6 +78,11 @@ async def fetch_models_for_key(
|
||||
timeout_seconds: float = 30.0,
|
||||
) -> tuple[list[dict], list[str], bool, dict[str, Any] | None]:
|
||||
"""统一入口：按 provider_type 选择策略获取模型列表（可附带 upstream_metadata）。"""
|
||||
# Ensure provider plugins (including custom model fetchers) are registered.
|
||||
from src.services.provider.envelope import ensure_providers_bootstrapped
|
||||
|
||||
ensure_providers_bootstrapped()
|
||||
|
||||
fetcher = UpstreamModelsFetcherRegistry.get(ctx.provider_type) or _fetch_models_default
|
||||
return await fetcher(ctx, timeout_seconds)
|
||||
|
||||
|
||||
@@ -105,22 +105,10 @@ ANTIGRAVITY_SYSTEM_INSTRUCTION = (
|
||||
"**Proactiveness**"
|
||||
)
|
||||
|
||||
# ============== JSON Schema 禁止字段 ==============
|
||||
FORBIDDEN_SCHEMA_FIELDS = frozenset(
|
||||
{
|
||||
"multipleOf",
|
||||
"exclusiveMinimum",
|
||||
"exclusiveMaximum",
|
||||
"contentEncoding",
|
||||
"contentMediaType",
|
||||
}
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"ANTIGRAVITY_SYSTEM_INSTRUCTION",
|
||||
"DAILY_BASE_URL",
|
||||
"DUMMY_THOUGHT_SIGNATURE",
|
||||
"FORBIDDEN_SCHEMA_FIELDS",
|
||||
"HTTP_USER_AGENT",
|
||||
"MIN_SIGNATURE_LENGTH",
|
||||
"PROD_BASE_URL",
|
||||
|
||||
@@ -22,7 +22,6 @@ from typing import Any
|
||||
|
||||
from src.services.provider.adapters.antigravity.constants import (
|
||||
ANTIGRAVITY_SYSTEM_INSTRUCTION,
|
||||
FORBIDDEN_SCHEMA_FIELDS,
|
||||
)
|
||||
from src.services.provider.adapters.antigravity.constants import (
|
||||
REQUEST_USER_AGENT as ANTIGRAVITY_REQUEST_USER_AGENT,
|
||||
@@ -199,33 +198,27 @@ def _clean_tool_declarations(inner_request: dict[str, Any]) -> None:
|
||||
|
||||
def _clean_json_schema(schema: dict[str, Any]) -> None:
|
||||
"""递归移除 Gemini 不支持的 JSON Schema 字段。"""
|
||||
for field in FORBIDDEN_SCHEMA_FIELDS:
|
||||
schema.pop(field, None)
|
||||
from src.core.api_format.schema_utils import clean_gemini_schema
|
||||
|
||||
# Recurse into properties
|
||||
props = schema.get("properties")
|
||||
if isinstance(props, dict):
|
||||
for prop_schema in props.values():
|
||||
if isinstance(prop_schema, dict):
|
||||
_clean_json_schema(prop_schema)
|
||||
clean_gemini_schema(schema)
|
||||
|
||||
# Recurse into items
|
||||
items = schema.get("items")
|
||||
if isinstance(items, dict):
|
||||
_clean_json_schema(items)
|
||||
|
||||
# Recurse into additionalProperties
|
||||
addl = schema.get("additionalProperties")
|
||||
if isinstance(addl, dict):
|
||||
_clean_json_schema(addl)
|
||||
def _compact_contents(inner_request: dict[str, Any]) -> None:
|
||||
"""Strip invalid parts, drop empty contents, merge consecutive same-role.
|
||||
|
||||
# Recurse into anyOf / oneOf / allOf
|
||||
for combo_key in ("anyOf", "oneOf", "allOf"):
|
||||
combo = schema.get(combo_key)
|
||||
if isinstance(combo, list):
|
||||
for sub_schema in combo:
|
||||
if isinstance(sub_schema, dict):
|
||||
_clean_json_schema(sub_schema)
|
||||
Delegates to :func:`compact_gemini_contents` from the Gemini normalizer to
|
||||
avoid duplicating the validation logic.
|
||||
|
||||
This function modifies *inner_request* in-place.
|
||||
"""
|
||||
from src.core.api_format.conversion.normalizers.gemini import compact_gemini_contents
|
||||
|
||||
contents = inner_request.get("contents")
|
||||
if not isinstance(contents, list):
|
||||
return
|
||||
|
||||
result = compact_gemini_contents(contents)
|
||||
inner_request["contents"] = result
|
||||
|
||||
|
||||
def _inject_system_instruction(inner_request: dict[str, Any]) -> None:
|
||||
@@ -349,8 +342,9 @@ def wrap_v1internal_request(
|
||||
5. Thinking budget 处理
|
||||
6. 工具声明清洗（图像生成模型跳过）
|
||||
7. System Instruction 注入（图像生成模型跳过）
|
||||
8. 注入 sessionId（对齐 CLIProxyAPI）
|
||||
9. 构建 v1internal 信封
|
||||
8. 清理空 parts 的 contents 并合并连续同角色条目
|
||||
9. 注入 sessionId（对齐 CLIProxyAPI）
|
||||
10. 构建 v1internal 信封
|
||||
"""
|
||||
from src.api.handlers.gemini.image_gen import is_image_gen_model
|
||||
|
||||
@@ -385,7 +379,12 @@ def wrap_v1internal_request(
|
||||
inner_request.pop("system_instruction", None)
|
||||
request_type = "image_gen"
|
||||
|
||||
# 6. 注入 sessionId（对齐 CLIProxyAPI/sub2api）
|
||||
# 6. 清理空 parts 的 contents 并合并连续同角色条目
|
||||
# 跨格式转换（如 Responses API reasoning 块）可能产生空 parts 的 content，
|
||||
# Gemini API 要求每个 content 至少有一个有效 part，并且严格交替 user/model 角色。
|
||||
_compact_contents(inner_request)
|
||||
|
||||
# 7. 注入 sessionId（对齐 CLIProxyAPI/sub2api）
|
||||
if "sessionId" not in inner_request:
|
||||
inner_request["sessionId"] = _generate_stable_session_id(inner_request)
|
||||
|
||||
|
||||
@@ -309,6 +309,31 @@ async def fetch_models_antigravity(
|
||||
return models, [], True, upstream_metadata
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Export builder
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_AG_SKIP_KEYS = frozenset(
|
||||
{
|
||||
"access_token",
|
||||
"expires_at",
|
||||
"updated_at",
|
||||
"token_type",
|
||||
"scope",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def antigravity_export_builder(
|
||||
auth_config: dict[str, Any],
|
||||
upstream_metadata: dict[str, Any] | None,
|
||||
) -> dict[str, Any]:
|
||||
"""Antigravity 导出：保留 refresh_token / email / project_id / tier。"""
|
||||
return {
|
||||
k: v for k, v in auth_config.items() if k not in _AG_SKIP_KEYS and v is not None and v != ""
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Unified Registration
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -321,6 +346,7 @@ def register_all() -> None:
|
||||
from src.services.provider.adapters.antigravity.envelope import antigravity_v1internal_envelope
|
||||
from src.services.provider.behavior import register_behavior_variant
|
||||
from src.services.provider.envelope import register_envelope
|
||||
from src.services.provider.export import register_export_builder
|
||||
from src.services.provider.transport import register_transport_hook
|
||||
|
||||
# Envelope
|
||||
@@ -337,6 +363,9 @@ def register_all() -> None:
|
||||
# Auth
|
||||
register_auth_enricher("antigravity", enrich_antigravity)
|
||||
|
||||
# Export
|
||||
register_export_builder("antigravity", antigravity_export_builder)
|
||||
|
||||
# Model Fetcher
|
||||
UpstreamModelsFetcherRegistry.register(
|
||||
provider_types=["antigravity"],
|
||||
|
||||
@@ -1,97 +0,0 @@
|
||||
"""
|
||||
Codex Provider 元数据采集器
|
||||
|
||||
从响应头解析 Codex 额度/限流信息：
|
||||
- x-codex-plan-type: 套餐类型
|
||||
- x-codex-primary-*: 主限额窗口（通常 7 天）
|
||||
- x-codex-secondary-*: 次级限额窗口（通常 5 小时）
|
||||
- x-codex-credits-*: 积分信息
|
||||
"""
|
||||
|
||||
from typing import Any, ClassVar
|
||||
|
||||
from src.services.provider.metadata_collectors import ( # noqa: E501 — core registry
|
||||
MetadataCollector,
|
||||
)
|
||||
|
||||
|
||||
def _safe_float(value: str | None) -> float | None:
|
||||
if value is None:
|
||||
return None
|
||||
try:
|
||||
return float(value)
|
||||
except (ValueError, TypeError):
|
||||
return None
|
||||
|
||||
|
||||
def _safe_int(value: str | None) -> int | None:
|
||||
if value is None:
|
||||
return None
|
||||
try:
|
||||
return int(float(value))
|
||||
except (ValueError, TypeError):
|
||||
return None
|
||||
|
||||
|
||||
def _safe_bool(value: str | None) -> bool | None:
|
||||
if value is None:
|
||||
return None
|
||||
return value.lower() in ("true", "1", "yes")
|
||||
|
||||
|
||||
class CodexMetadataCollector(MetadataCollector):
|
||||
"""Codex 额度/限流元数据采集器"""
|
||||
|
||||
# 支持 codex 类型，通过响应头判断是否是 Codex
|
||||
PROVIDER_TYPES: ClassVar[list[str]] = ["codex"]
|
||||
|
||||
def parse_headers(self, headers: dict[str, str]) -> dict[str, Any] | None:
|
||||
# 大小写不敏感查找
|
||||
lower_headers = {k.lower(): v for k, v in headers.items()}
|
||||
|
||||
plan_type = lower_headers.get("x-codex-plan-type")
|
||||
if plan_type is None:
|
||||
# 没有 Codex 特征头，跳过
|
||||
return None
|
||||
|
||||
result: dict[str, Any] = {"plan_type": plan_type}
|
||||
|
||||
# 主限额窗口（7 天）
|
||||
primary_used = _safe_float(lower_headers.get("x-codex-primary-used-percent"))
|
||||
if primary_used is not None:
|
||||
result["primary_used_percent"] = primary_used
|
||||
primary_reset_seconds = _safe_int(lower_headers.get("x-codex-primary-reset-after-seconds"))
|
||||
if primary_reset_seconds is not None:
|
||||
result["primary_reset_seconds"] = primary_reset_seconds
|
||||
primary_reset_at = _safe_int(lower_headers.get("x-codex-primary-reset-at"))
|
||||
if primary_reset_at is not None:
|
||||
result["primary_reset_at"] = primary_reset_at
|
||||
primary_window = _safe_int(lower_headers.get("x-codex-primary-window-minutes"))
|
||||
if primary_window is not None:
|
||||
result["primary_window_minutes"] = primary_window
|
||||
|
||||
# 次级限额窗口（5 小时）
|
||||
secondary_used = _safe_float(lower_headers.get("x-codex-secondary-used-percent"))
|
||||
if secondary_used is not None:
|
||||
result["secondary_used_percent"] = secondary_used
|
||||
secondary_reset_seconds = _safe_int(
|
||||
lower_headers.get("x-codex-secondary-reset-after-seconds")
|
||||
)
|
||||
if secondary_reset_seconds is not None:
|
||||
result["secondary_reset_seconds"] = secondary_reset_seconds
|
||||
secondary_reset_at = _safe_int(lower_headers.get("x-codex-secondary-reset-at"))
|
||||
if secondary_reset_at is not None:
|
||||
result["secondary_reset_at"] = secondary_reset_at
|
||||
secondary_window = _safe_int(lower_headers.get("x-codex-secondary-window-minutes"))
|
||||
if secondary_window is not None:
|
||||
result["secondary_window_minutes"] = secondary_window
|
||||
|
||||
# 积分信息
|
||||
has_credits = _safe_bool(lower_headers.get("x-codex-credits-has-credits"))
|
||||
if has_credits is not None:
|
||||
result["has_credits"] = has_credits
|
||||
credits_balance = _safe_float(lower_headers.get("x-codex-credits-balance"))
|
||||
if credits_balance is not None:
|
||||
result["credits_balance"] = credits_balance
|
||||
|
||||
return result
|
||||
@@ -5,6 +5,7 @@
|
||||
- Transport Hook (URL 构建)
|
||||
- Auth Enricher (OAuth enrichment)
|
||||
- Behavior Variants (格式变体)
|
||||
- Model Fetcher (fixed catalog — Codex has no /v1/models endpoint)
|
||||
|
||||
新增 provider 时参照此文件创建对应的 plugin.py 即可。
|
||||
"""
|
||||
@@ -16,6 +17,40 @@ from urllib.parse import urlencode
|
||||
|
||||
from src.core.logger import logger
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fixed model catalog
|
||||
# ---------------------------------------------------------------------------
|
||||
# Codex upstream (chatgpt.com/backend-api/codex) has no /v1/models endpoint.
|
||||
# Return a static list of known models.
|
||||
_CODEX_MODELS: list[dict[str, Any]] = [
|
||||
{
|
||||
"id": "gpt-5.2",
|
||||
"object": "model",
|
||||
"owned_by": "openai",
|
||||
"display_name": "gpt-5.2",
|
||||
},
|
||||
{
|
||||
"id": "gpt-5.2-codex",
|
||||
"object": "model",
|
||||
"owned_by": "openai",
|
||||
"display_name": "gpt-5.2-codex",
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
async def fetch_models_codex(
|
||||
ctx: Any,
|
||||
timeout_seconds: float, # noqa: ARG001
|
||||
) -> tuple[list[dict], list[str], bool, dict[str, Any] | None]:
|
||||
"""Return a fixed model catalog for Codex.
|
||||
|
||||
Codex upstream does not expose a ``/v1/models`` endpoint, so we skip the
|
||||
HTTP call entirely and return a hardcoded list.
|
||||
"""
|
||||
_ = ctx
|
||||
return list(_CODEX_MODELS), [], True, None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Transport Hook
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -87,6 +122,7 @@ async def enrich_codex(
|
||||
def register_all() -> None:
|
||||
"""一次性注册 Codex 的所有 hooks 到各通用 registry。"""
|
||||
from src.core.provider_oauth_utils import register_auth_enricher
|
||||
from src.services.model.upstream_fetcher import UpstreamModelsFetcherRegistry
|
||||
from src.services.provider.adapters.codex.envelope import codex_oauth_envelope
|
||||
from src.services.provider.behavior import register_behavior_variant
|
||||
from src.services.provider.envelope import register_envelope
|
||||
@@ -104,3 +140,12 @@ def register_all() -> None:
|
||||
|
||||
# Behavior
|
||||
register_behavior_variant("codex", same_format=True, cross_format=True)
|
||||
|
||||
# Export: Codex uses the default export builder (strip null + temp fields)
|
||||
# No need to register a custom one — the default in export.py suffices.
|
||||
|
||||
# Model Fetcher
|
||||
UpstreamModelsFetcherRegistry.register(
|
||||
provider_types=["codex"],
|
||||
fetcher=fetch_models_codex,
|
||||
)
|
||||
|
||||
3
src/services/provider/adapters/kiro/__init__.py
Normal file
3
src/services/provider/adapters/kiro/__init__.py
Normal file
@@ -0,0 +1,3 @@
|
||||
"""Kiro provider adapter."""
|
||||
|
||||
__all__ = []
|
||||
95
src/services/provider/adapters/kiro/constants.py
Normal file
95
src/services/provider/adapters/kiro/constants.py
Normal file
@@ -0,0 +1,95 @@
|
||||
"""Kiro adapter constants.
|
||||
|
||||
Kiro upstream uses AWS Event Stream (binary frames) for streaming responses.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import platform
|
||||
|
||||
AWS_EVENTSTREAM_CONTENT_TYPE = "application/vnd.amazon.eventstream"
|
||||
|
||||
# Kiro API endpoints
|
||||
KIRO_GENERATE_ASSISTANT_PATH = "/generateAssistantResponse"
|
||||
KIRO_USAGE_LIMITS_PATH = "/getUsageLimits"
|
||||
|
||||
# Default AWS region when not specified in credentials
|
||||
DEFAULT_REGION = "us-east-1"
|
||||
|
||||
# Default client fingerprints used in headers (best-effort)
|
||||
DEFAULT_KIRO_VERSION = "0.8.0"
|
||||
DEFAULT_NODE_VERSION = "22.21.1"
|
||||
|
||||
|
||||
def _detect_system_version() -> str:
|
||||
system = platform.system().lower() or "other"
|
||||
release = platform.release() or "unknown"
|
||||
# Match KiroIDE style: darwin#24.6.0, windows#10.0.22631, linux#6.8.0-...
|
||||
return f"{system}#{release}"
|
||||
|
||||
|
||||
DEFAULT_SYSTEM_VERSION = _detect_system_version()
|
||||
|
||||
# Header constants
|
||||
KIRO_AGENT_MODE = "vibe"
|
||||
CODEWHISPERER_OPTOUT = "true"
|
||||
|
||||
# aws-sdk-js versions observed in kiro.rs
|
||||
AWS_SDK_JS_MAIN_VERSION = "1.0.27"
|
||||
AWS_SDK_JS_USAGE_VERSION = "1.0.0"
|
||||
|
||||
# Claude model context window used by kiro.rs to convert contextUsage percentage -> tokens
|
||||
CONTEXT_WINDOW_TOKENS = 200_000
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Chunked-write policy injected into tool descriptions and system prompt
|
||||
# ---------------------------------------------------------------------------
|
||||
# Kiro upstream has lower per-message size limits than standard Claude.
|
||||
# We inject instructions for Write/Edit tools and a system-level policy
|
||||
# so the model splits large writes into smaller chunks automatically.
|
||||
|
||||
WRITE_TOOL_DESCRIPTION_SUFFIX = (
|
||||
"- IMPORTANT: If the content to write exceeds 150 lines, you MUST only write "
|
||||
"the first 50 lines using this tool, then use `Edit` tool to append the "
|
||||
"remaining content in chunks of no more than 50 lines each. If needed, leave "
|
||||
"a unique placeholder to help append content. Do NOT attempt to write all "
|
||||
"content at once."
|
||||
)
|
||||
|
||||
EDIT_TOOL_DESCRIPTION_SUFFIX = (
|
||||
"- IMPORTANT: If the `new_string` content exceeds 50 lines, you MUST split "
|
||||
"it into multiple Edit calls, each replacing no more than 50 lines at a time. "
|
||||
"If used to append content, leave a unique placeholder to help append content. "
|
||||
"On the final chunk, do NOT include the placeholder."
|
||||
)
|
||||
|
||||
TOOL_DESCRIPTION_SUFFIXES: dict[str, str] = {
|
||||
"Write": WRITE_TOOL_DESCRIPTION_SUFFIX,
|
||||
"Edit": EDIT_TOOL_DESCRIPTION_SUFFIX,
|
||||
}
|
||||
|
||||
SYSTEM_CHUNKED_POLICY = (
|
||||
"When the Write or Edit tool has content size limits, always comply silently. "
|
||||
"Never suggest bypassing these limits via alternative tools. "
|
||||
"Never ask the user whether to switch approaches. "
|
||||
"Complete all chunked operations without commentary."
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"AWS_EVENTSTREAM_CONTENT_TYPE",
|
||||
"AWS_SDK_JS_MAIN_VERSION",
|
||||
"AWS_SDK_JS_USAGE_VERSION",
|
||||
"CODEWHISPERER_OPTOUT",
|
||||
"CONTEXT_WINDOW_TOKENS",
|
||||
"DEFAULT_KIRO_VERSION",
|
||||
"DEFAULT_NODE_VERSION",
|
||||
"DEFAULT_REGION",
|
||||
"DEFAULT_SYSTEM_VERSION",
|
||||
"EDIT_TOOL_DESCRIPTION_SUFFIX",
|
||||
"KIRO_AGENT_MODE",
|
||||
"KIRO_GENERATE_ASSISTANT_PATH",
|
||||
"KIRO_USAGE_LIMITS_PATH",
|
||||
"SYSTEM_CHUNKED_POLICY",
|
||||
"TOOL_DESCRIPTION_SUFFIXES",
|
||||
"WRITE_TOOL_DESCRIPTION_SUFFIX",
|
||||
]
|
||||
42
src/services/provider/adapters/kiro/context.py
Normal file
42
src/services/provider/adapters/kiro/context.py
Normal file
@@ -0,0 +1,42 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import contextvars
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class KiroRequestContext:
|
||||
"""Per-request context for the Kiro adapter.
|
||||
|
||||
This bridges data from `KiroEnvelope.wrap_request()` (which receives the
|
||||
decrypted auth_config + original request body) to other layers that only
|
||||
expose parameterless hooks (extra_headers) or transport hooks.
|
||||
"""
|
||||
|
||||
region: str
|
||||
machine_id: str
|
||||
kiro_version: str | None = None
|
||||
system_version: str | None = None
|
||||
node_version: str | None = None
|
||||
thinking_enabled: bool = False
|
||||
|
||||
|
||||
_kiro_request_context: contextvars.ContextVar[KiroRequestContext | None] = contextvars.ContextVar(
|
||||
"kiro_request_context",
|
||||
default=None,
|
||||
)
|
||||
|
||||
|
||||
def set_kiro_request_context(ctx: KiroRequestContext | None) -> None:
|
||||
_kiro_request_context.set(ctx)
|
||||
|
||||
|
||||
def get_kiro_request_context() -> KiroRequestContext | None:
|
||||
return _kiro_request_context.get()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"KiroRequestContext",
|
||||
"get_kiro_request_context",
|
||||
"set_kiro_request_context",
|
||||
]
|
||||
620
src/services/provider/adapters/kiro/converter.py
Normal file
620
src/services/provider/adapters/kiro/converter.py
Normal file
@@ -0,0 +1,620 @@
|
||||
"""Claude Messages -> Kiro ConversationState converter (best-effort).
|
||||
|
||||
This mirrors `kiro.rs/src/anthropic/converter.rs` but focuses on the fields
|
||||
needed by generateAssistantResponse.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from typing import Any
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.services.provider.adapters.kiro.constants import (
|
||||
SYSTEM_CHUNKED_POLICY as _SYSTEM_CHUNKED_POLICY,
|
||||
)
|
||||
from src.services.provider.adapters.kiro.constants import (
|
||||
TOOL_DESCRIPTION_SUFFIXES as _TOOL_DESCRIPTION_SUFFIXES,
|
||||
)
|
||||
|
||||
|
||||
def map_model(model: str) -> str | None:
|
||||
"""Map an Anthropic model name to a Kiro-compatible model ID.
|
||||
|
||||
Kiro upstream expects specific model IDs (e.g. ``claude-sonnet-4.5``).
|
||||
If *model* already matches a Kiro ID it is returned as-is; otherwise we
|
||||
attempt a best-effort fuzzy mapping. Returns ``None`` when the model is
|
||||
unrecognised.
|
||||
"""
|
||||
raw = str(model or "").strip()
|
||||
if not raw:
|
||||
return None
|
||||
|
||||
model_lower = raw.lower()
|
||||
|
||||
# Already a Kiro-native ID?
|
||||
_KIRO_NATIVE_IDS = {
|
||||
"claude-sonnet-4.5",
|
||||
"claude-opus-4.5",
|
||||
"claude-opus-4.6",
|
||||
"claude-haiku-4.5",
|
||||
}
|
||||
if model_lower in _KIRO_NATIVE_IDS:
|
||||
return model_lower
|
||||
|
||||
# Fuzzy mapping from Anthropic-style model names
|
||||
if "sonnet" in model_lower:
|
||||
return "claude-sonnet-4.5"
|
||||
if "opus" in model_lower:
|
||||
if "4-5" in model_lower or "4.5" in model_lower:
|
||||
return "claude-opus-4.5"
|
||||
return "claude-opus-4.6"
|
||||
if "haiku" in model_lower:
|
||||
return "claude-haiku-4.5"
|
||||
|
||||
# Unrecognised — pass through as-is (let the upstream decide)
|
||||
return raw
|
||||
|
||||
|
||||
def _extract_session_id(user_id: str) -> str | None:
|
||||
text = str(user_id or "")
|
||||
pos = text.find("session_")
|
||||
if pos < 0:
|
||||
return None
|
||||
session_part = text[pos + 8 :]
|
||||
if len(session_part) < 36:
|
||||
return None
|
||||
candidate = session_part[:36]
|
||||
if candidate.count("-") != 4:
|
||||
return None
|
||||
return candidate
|
||||
|
||||
|
||||
def _generate_thinking_prefix(request_body: dict[str, Any]) -> str | None:
|
||||
thinking = request_body.get("thinking")
|
||||
if not isinstance(thinking, dict):
|
||||
return None
|
||||
|
||||
thinking_type = str(thinking.get("type") or "").strip()
|
||||
if thinking_type == "enabled":
|
||||
budget = thinking.get("budget_tokens")
|
||||
try:
|
||||
budget_i = int(budget) if budget is not None else 0
|
||||
except Exception:
|
||||
budget_i = 0
|
||||
return (
|
||||
f"<thinking_mode>enabled</thinking_mode>"
|
||||
f"<max_thinking_length>{budget_i}</max_thinking_length>"
|
||||
)
|
||||
|
||||
if thinking_type == "adaptive":
|
||||
output_cfg = request_body.get("output_config")
|
||||
effort = "high"
|
||||
if isinstance(output_cfg, dict):
|
||||
eff = output_cfg.get("effort")
|
||||
if isinstance(eff, str) and eff.strip():
|
||||
effort = eff.strip()
|
||||
return (
|
||||
f"<thinking_mode>adaptive</thinking_mode>"
|
||||
f"<thinking_effort>{effort}</thinking_effort>"
|
||||
)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _has_thinking_tags(content: str) -> bool:
|
||||
return "<thinking_mode>" in content or "<max_thinking_length>" in content
|
||||
|
||||
|
||||
def _system_to_text(system: Any) -> str:
|
||||
if system is None:
|
||||
return ""
|
||||
if isinstance(system, str):
|
||||
return system
|
||||
if isinstance(system, list):
|
||||
parts: list[str] = []
|
||||
for item in system:
|
||||
if isinstance(item, dict):
|
||||
t = item.get("text")
|
||||
if isinstance(t, str) and t:
|
||||
parts.append(t)
|
||||
else:
|
||||
# best-effort
|
||||
try:
|
||||
parts.append(str(item))
|
||||
except Exception:
|
||||
pass
|
||||
return "\n".join([p for p in parts if p])
|
||||
return ""
|
||||
|
||||
|
||||
def _get_image_format(media_type: str | None) -> str | None:
|
||||
if not isinstance(media_type, str) or "/" not in media_type:
|
||||
return None
|
||||
prefix, suffix = media_type.split("/", 1)
|
||||
if prefix != "image":
|
||||
return None
|
||||
suffix = suffix.strip().lower()
|
||||
if suffix in {"jpeg", "png", "gif", "webp"}:
|
||||
return suffix
|
||||
if suffix == "jpg":
|
||||
return "jpeg"
|
||||
return None
|
||||
|
||||
|
||||
def _process_message_content(
|
||||
content: Any,
|
||||
) -> tuple[str, list[dict[str, Any]], list[dict[str, Any]]]:
|
||||
"""Extract text/images/tool_results from a Claude content field."""
|
||||
text_parts: list[str] = []
|
||||
images: list[dict[str, Any]] = []
|
||||
tool_results: list[dict[str, Any]] = []
|
||||
|
||||
if isinstance(content, str):
|
||||
if content:
|
||||
text_parts.append(content)
|
||||
return "".join(text_parts), images, tool_results
|
||||
|
||||
if not isinstance(content, list):
|
||||
return "".join(text_parts), images, tool_results
|
||||
|
||||
for block in content:
|
||||
if not isinstance(block, dict):
|
||||
continue
|
||||
|
||||
btype = str(block.get("type") or "").strip()
|
||||
|
||||
if btype == "text":
|
||||
text = block.get("text")
|
||||
if isinstance(text, str) and text:
|
||||
text_parts.append(text)
|
||||
continue
|
||||
|
||||
if btype == "image":
|
||||
source = block.get("source")
|
||||
if not isinstance(source, dict):
|
||||
continue
|
||||
media_type = source.get("media_type") or source.get("mediaType")
|
||||
fmt = _get_image_format(media_type if isinstance(media_type, str) else None)
|
||||
data = source.get("data")
|
||||
if fmt and isinstance(data, str) and data:
|
||||
images.append({"format": fmt, "source": {"bytes": data}})
|
||||
continue
|
||||
|
||||
if btype == "tool_result":
|
||||
tool_use_id = block.get("tool_use_id") or block.get("toolUseId")
|
||||
if not isinstance(tool_use_id, str) or not tool_use_id.strip():
|
||||
continue
|
||||
|
||||
raw_content = block.get("content")
|
||||
if isinstance(raw_content, str):
|
||||
text = raw_content
|
||||
elif isinstance(raw_content, list):
|
||||
# Claude tool_result content blocks; keep only text parts.
|
||||
parts: list[str] = []
|
||||
for item in raw_content:
|
||||
if isinstance(item, dict) and item.get("type") == "text":
|
||||
t = item.get("text")
|
||||
if isinstance(t, str) and t:
|
||||
parts.append(t)
|
||||
text = "\n".join(parts)
|
||||
else:
|
||||
try:
|
||||
text = json.dumps(raw_content, ensure_ascii=False)
|
||||
except Exception:
|
||||
text = str(raw_content)
|
||||
|
||||
is_error = bool(block.get("is_error") or block.get("isError") or False)
|
||||
status = "error" if is_error else "success"
|
||||
|
||||
tool_results.append(
|
||||
{
|
||||
"toolUseId": tool_use_id.strip(),
|
||||
"content": [{"text": text or ""}],
|
||||
"status": status,
|
||||
"isError": bool(is_error),
|
||||
}
|
||||
)
|
||||
continue
|
||||
|
||||
return "".join(text_parts), images, tool_results
|
||||
|
||||
|
||||
def _convert_tools(tools: Any) -> list[dict[str, Any]]:
|
||||
if not isinstance(tools, list):
|
||||
return []
|
||||
|
||||
out: list[dict[str, Any]] = []
|
||||
for t in tools:
|
||||
if not isinstance(t, dict):
|
||||
continue
|
||||
name = t.get("name")
|
||||
if not isinstance(name, str) or not name.strip():
|
||||
continue
|
||||
|
||||
description = t.get("description")
|
||||
description_str = description if isinstance(description, str) else ""
|
||||
|
||||
# Inject chunked-write instructions for Write/Edit tools.
|
||||
suffix = _TOOL_DESCRIPTION_SUFFIXES.get(name.strip())
|
||||
if suffix:
|
||||
description_str = f"{description_str}\n{suffix}" if description_str else suffix
|
||||
|
||||
if len(description_str) > 10000:
|
||||
description_str = description_str[:10000]
|
||||
|
||||
input_schema = t.get("input_schema") or t.get("inputSchema") or {}
|
||||
if not isinstance(input_schema, dict):
|
||||
input_schema = {}
|
||||
|
||||
out.append(
|
||||
{
|
||||
"toolSpecification": {
|
||||
"name": name.strip(),
|
||||
"description": description_str,
|
||||
"inputSchema": {"json": input_schema},
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def _create_placeholder_tool(name: str) -> dict[str, Any]:
|
||||
return {
|
||||
"toolSpecification": {
|
||||
"name": name,
|
||||
"description": "Tool used in conversation history",
|
||||
"inputSchema": {
|
||||
"json": {
|
||||
"$schema": "http://json-schema.org/draft-07/schema#",
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
"required": [],
|
||||
"additionalProperties": True,
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def _convert_assistant_message(message: dict[str, Any]) -> dict[str, Any] | None:
|
||||
content = message.get("content")
|
||||
|
||||
tool_uses: list[dict[str, Any]] = []
|
||||
thinking_parts: list[str] = []
|
||||
text_parts: list[str] = []
|
||||
|
||||
if isinstance(content, str):
|
||||
if content:
|
||||
text_parts.append(content)
|
||||
elif isinstance(content, list):
|
||||
for block in content:
|
||||
if not isinstance(block, dict):
|
||||
continue
|
||||
btype = str(block.get("type") or "")
|
||||
if btype == "thinking":
|
||||
# Preserve thinking content so multi-turn context is not lost.
|
||||
t = block.get("thinking")
|
||||
if isinstance(t, str) and t:
|
||||
thinking_parts.append(t)
|
||||
elif btype == "text":
|
||||
t = block.get("text")
|
||||
if isinstance(t, str) and t:
|
||||
text_parts.append(t)
|
||||
elif btype == "tool_use":
|
||||
tool_use_id = block.get("id")
|
||||
name = block.get("name")
|
||||
if not isinstance(tool_use_id, str) or not tool_use_id.strip():
|
||||
continue
|
||||
if not isinstance(name, str) or not name.strip():
|
||||
continue
|
||||
inp = block.get("input")
|
||||
if not isinstance(inp, dict):
|
||||
inp = {}
|
||||
tool_uses.append(
|
||||
{
|
||||
"toolUseId": tool_use_id.strip(),
|
||||
"name": name.strip(),
|
||||
"input": inp,
|
||||
}
|
||||
)
|
||||
|
||||
# Combine thinking + text into final content.
|
||||
# Format: <thinking>...</thinking>\n\ntext
|
||||
thinking_str = "".join(thinking_parts)
|
||||
text_str = "".join(text_parts)
|
||||
|
||||
if thinking_str:
|
||||
if text_str:
|
||||
content_str = f"<thinking>{thinking_str}</thinking>\n\n{text_str}"
|
||||
else:
|
||||
content_str = f"<thinking>{thinking_str}</thinking>"
|
||||
else:
|
||||
content_str = text_str
|
||||
|
||||
if not content_str and tool_uses:
|
||||
content_str = " " # Kiro API requires non-empty content.
|
||||
|
||||
if not content_str and not tool_uses:
|
||||
return None
|
||||
|
||||
out: dict[str, Any] = {"content": content_str}
|
||||
if tool_uses:
|
||||
out["toolUses"] = tool_uses
|
||||
return out
|
||||
|
||||
|
||||
def convert_claude_messages_to_conversation_state(
|
||||
request_body: dict[str, Any],
|
||||
*,
|
||||
model: str,
|
||||
) -> dict[str, Any]:
|
||||
model_id = map_model(model)
|
||||
if not model_id:
|
||||
raise ValueError(f"kiro: model is required (got {model!r})")
|
||||
|
||||
messages = request_body.get("messages")
|
||||
if not isinstance(messages, list) or not messages:
|
||||
raise ValueError("kiro: empty messages")
|
||||
|
||||
conversation_id = None
|
||||
metadata = request_body.get("metadata")
|
||||
if isinstance(metadata, dict):
|
||||
user_id = metadata.get("user_id") or metadata.get("userId")
|
||||
if isinstance(user_id, str) and user_id:
|
||||
conversation_id = _extract_session_id(user_id)
|
||||
|
||||
if not conversation_id:
|
||||
conversation_id = str(uuid.uuid4())
|
||||
|
||||
agent_continuation_id = str(uuid.uuid4())
|
||||
|
||||
thinking_prefix = _generate_thinking_prefix(request_body)
|
||||
|
||||
history: list[dict[str, Any]] = []
|
||||
|
||||
# System injection: add as (user, assistant) pair.
|
||||
system_text = _system_to_text(request_body.get("system"))
|
||||
if system_text:
|
||||
# Append chunked-write policy so the model silently obeys tool limits.
|
||||
final_system = f"{system_text}\n{_SYSTEM_CHUNKED_POLICY}"
|
||||
if thinking_prefix and not _has_thinking_tags(system_text):
|
||||
final_system = f"{thinking_prefix}\n{final_system}"
|
||||
history.append(
|
||||
{
|
||||
"userInputMessage": {
|
||||
"content": final_system,
|
||||
"modelId": model_id,
|
||||
"origin": "AI_EDITOR",
|
||||
}
|
||||
}
|
||||
)
|
||||
history.append(
|
||||
{"assistantResponseMessage": {"content": "I will follow these instructions."}}
|
||||
)
|
||||
elif thinking_prefix:
|
||||
history.append(
|
||||
{
|
||||
"userInputMessage": {
|
||||
"content": thinking_prefix,
|
||||
"modelId": model_id,
|
||||
"origin": "AI_EDITOR",
|
||||
}
|
||||
}
|
||||
)
|
||||
history.append(
|
||||
{"assistantResponseMessage": {"content": "I will follow these instructions."}}
|
||||
)
|
||||
|
||||
# Build history from messages.
|
||||
# If the last message is assistant, include it in history (Kiro currentMessage
|
||||
# must be user; we synthesise one). Otherwise the last user message becomes
|
||||
# currentMessage and everything before it goes into history.
|
||||
last_msg = messages[-1]
|
||||
last_is_assistant = (
|
||||
isinstance(last_msg, dict) and str(last_msg.get("role") or "") == "assistant"
|
||||
)
|
||||
|
||||
if last_is_assistant:
|
||||
# All messages go into history; we'll synthesise a currentMessage later.
|
||||
history_end_index = len(messages)
|
||||
else:
|
||||
history_end_index = max(len(messages) - 1, 0)
|
||||
|
||||
user_buffer: list[dict[str, Any]] = []
|
||||
|
||||
def _flush_user_buffer() -> dict[str, Any] | None:
|
||||
nonlocal user_buffer
|
||||
if not user_buffer:
|
||||
return None
|
||||
|
||||
parts: list[str] = []
|
||||
images: list[dict[str, Any]] = []
|
||||
tool_results: list[dict[str, Any]] = []
|
||||
|
||||
for msg in user_buffer:
|
||||
text, imgs, results = _process_message_content(msg.get("content"))
|
||||
if text:
|
||||
parts.append(text)
|
||||
images.extend(imgs)
|
||||
tool_results.extend(results)
|
||||
|
||||
user_buffer = []
|
||||
|
||||
payload: dict[str, Any] = {
|
||||
"content": "\n".join(parts),
|
||||
"modelId": model_id,
|
||||
"origin": "AI_EDITOR",
|
||||
}
|
||||
|
||||
if images:
|
||||
payload["images"] = images
|
||||
|
||||
if tool_results:
|
||||
payload["userInputMessageContext"] = {"toolResults": tool_results}
|
||||
|
||||
return {"userInputMessage": payload}
|
||||
|
||||
for i in range(history_end_index):
|
||||
msg = messages[i]
|
||||
if not isinstance(msg, dict):
|
||||
continue
|
||||
role = str(msg.get("role") or "")
|
||||
if role == "user":
|
||||
user_buffer.append(msg)
|
||||
continue
|
||||
if role == "assistant":
|
||||
user_item = _flush_user_buffer()
|
||||
if user_item is not None:
|
||||
history.append(user_item)
|
||||
assistant_item = _convert_assistant_message(msg)
|
||||
if assistant_item is not None:
|
||||
history.append({"assistantResponseMessage": assistant_item})
|
||||
continue
|
||||
|
||||
# trailing unpaired user messages in history
|
||||
tail_user = _flush_user_buffer()
|
||||
if tail_user is not None:
|
||||
history.append(tail_user)
|
||||
history.append({"assistantResponseMessage": {"content": "OK"}})
|
||||
|
||||
# Current message: last message as user input.
|
||||
if last_is_assistant:
|
||||
# Synthesise a minimal user continuation message.
|
||||
text_content = "Continue."
|
||||
images: list[dict[str, Any]] = []
|
||||
tool_results: list[dict[str, Any]] = []
|
||||
else:
|
||||
last = messages[-1]
|
||||
if not isinstance(last, dict) or str(last.get("role") or "") != "user":
|
||||
raise ValueError("kiro: last message must be user")
|
||||
text_content, images, tool_results = _process_message_content(last.get("content"))
|
||||
|
||||
tools = _convert_tools(request_body.get("tools"))
|
||||
|
||||
# Ensure tools referenced in history assistant toolUses are defined.
|
||||
# Also collect ids for tool_use / tool_result pairing validation.
|
||||
history_tool_names: set[str] = set()
|
||||
history_tool_results_ids: set[str] = set()
|
||||
history_tool_use_ids: set[str] = set()
|
||||
|
||||
for item in history:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
u = item.get("userInputMessage")
|
||||
if isinstance(u, dict):
|
||||
ctx = u.get("userInputMessageContext")
|
||||
if isinstance(ctx, dict):
|
||||
results = ctx.get("toolResults")
|
||||
if isinstance(results, list):
|
||||
for r in results:
|
||||
if isinstance(r, dict):
|
||||
tid = r.get("toolUseId")
|
||||
if isinstance(tid, str) and tid:
|
||||
history_tool_results_ids.add(tid)
|
||||
a = item.get("assistantResponseMessage")
|
||||
if isinstance(a, dict):
|
||||
uses = a.get("toolUses")
|
||||
if isinstance(uses, list):
|
||||
for tu in uses:
|
||||
if not isinstance(tu, dict):
|
||||
continue
|
||||
nm = tu.get("name")
|
||||
if isinstance(nm, str) and nm:
|
||||
history_tool_names.add(nm)
|
||||
tid = tu.get("toolUseId")
|
||||
if isinstance(tid, str) and tid:
|
||||
history_tool_use_ids.add(tid)
|
||||
|
||||
existing_tool_names = {
|
||||
str(t.get("toolSpecification", {}).get("name", "")).lower() for t in tools
|
||||
}
|
||||
|
||||
for tool_name in sorted(history_tool_names):
|
||||
if tool_name.lower() not in existing_tool_names:
|
||||
tools.append(_create_placeholder_tool(tool_name))
|
||||
|
||||
# Filter tool_results: only keep those with matching tool_use in history, and not duplicated.
|
||||
validated_tool_results: list[dict[str, Any]] = []
|
||||
current_tool_result_ids: set[str] = set()
|
||||
for r in tool_results:
|
||||
if not isinstance(r, dict):
|
||||
continue
|
||||
tid = r.get("toolUseId")
|
||||
if not isinstance(tid, str) or not tid:
|
||||
continue
|
||||
if tid not in history_tool_use_ids:
|
||||
continue
|
||||
if tid in history_tool_results_ids:
|
||||
continue
|
||||
validated_tool_results.append(r)
|
||||
current_tool_result_ids.add(tid)
|
||||
|
||||
# Remove orphaned tool_uses from history.
|
||||
# Kiro API requires every tool_use to have a matching tool_result; otherwise
|
||||
# it returns 400 Bad Request.
|
||||
orphaned_tool_use_ids = (
|
||||
history_tool_use_ids - history_tool_results_ids - current_tool_result_ids
|
||||
)
|
||||
if orphaned_tool_use_ids:
|
||||
logger.warning(
|
||||
"kiro: removing {} orphaned tool_use(s) from history: {}",
|
||||
len(orphaned_tool_use_ids),
|
||||
orphaned_tool_use_ids,
|
||||
)
|
||||
for item in history:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
a = item.get("assistantResponseMessage")
|
||||
if not isinstance(a, dict):
|
||||
continue
|
||||
uses = a.get("toolUses")
|
||||
if not isinstance(uses, list):
|
||||
continue
|
||||
filtered = [
|
||||
u
|
||||
for u in uses
|
||||
if not (
|
||||
isinstance(u, dict)
|
||||
and isinstance(u.get("toolUseId"), str)
|
||||
and u["toolUseId"] in orphaned_tool_use_ids
|
||||
)
|
||||
]
|
||||
if not filtered:
|
||||
a.pop("toolUses", None)
|
||||
elif len(filtered) != len(uses):
|
||||
a["toolUses"] = filtered
|
||||
|
||||
user_ctx: dict[str, Any] = {}
|
||||
if tools:
|
||||
user_ctx["tools"] = tools
|
||||
if validated_tool_results:
|
||||
user_ctx["toolResults"] = validated_tool_results
|
||||
|
||||
user_input: dict[str, Any] = {
|
||||
"userInputMessageContext": user_ctx,
|
||||
"content": text_content,
|
||||
"modelId": model_id,
|
||||
"origin": "AI_EDITOR",
|
||||
}
|
||||
if images:
|
||||
user_input["images"] = images
|
||||
|
||||
conversation_state = {
|
||||
"agentContinuationId": agent_continuation_id,
|
||||
"agentTaskType": "vibe",
|
||||
"chatTriggerType": "MANUAL",
|
||||
"currentMessage": {"userInputMessage": user_input},
|
||||
"conversationId": conversation_id,
|
||||
"history": history,
|
||||
}
|
||||
|
||||
return conversation_state
|
||||
|
||||
|
||||
__all__ = [
|
||||
"convert_claude_messages_to_conversation_state",
|
||||
"map_model",
|
||||
]
|
||||
121
src/services/provider/adapters/kiro/envelope.py
Normal file
121
src/services/provider/adapters/kiro/envelope.py
Normal file
@@ -0,0 +1,121 @@
|
||||
"""Kiro provider envelope.
|
||||
|
||||
Kiro upstream is not Claude wire-compatible:
|
||||
- Request: wrap Claude Messages body into Kiro `conversationState` request.
|
||||
- Stream response: handled by StreamProcessor via binary EventStream rewrite.
|
||||
|
||||
We use contextvars to pass request-scoped values (region, machine_id, thinking)
|
||||
from wrap_request() to extra_headers() and transport hook.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from src.services.provider.adapters.kiro.context import KiroRequestContext, set_kiro_request_context
|
||||
from src.services.provider.adapters.kiro.converter import (
|
||||
convert_claude_messages_to_conversation_state,
|
||||
)
|
||||
from src.services.provider.adapters.kiro.headers import build_generate_assistant_headers
|
||||
from src.services.provider.adapters.kiro.models.credentials import KiroAuthConfig
|
||||
from src.services.provider.adapters.kiro.token_manager import generate_machine_id
|
||||
|
||||
|
||||
def _resolve_region(cfg: KiroAuthConfig) -> str:
|
||||
from src.services.provider.adapters.kiro.constants import DEFAULT_REGION
|
||||
|
||||
region = str(cfg.region or "").strip()
|
||||
return region or DEFAULT_REGION
|
||||
|
||||
|
||||
def _is_thinking_enabled(request_body: dict[str, Any]) -> bool:
|
||||
thinking = request_body.get("thinking")
|
||||
if not isinstance(thinking, dict):
|
||||
return False
|
||||
ttype = str(thinking.get("type") or "").strip().lower()
|
||||
return ttype in {"enabled", "adaptive"}
|
||||
|
||||
|
||||
class KiroEnvelope:
|
||||
name = "kiro:generateAssistantResponse"
|
||||
|
||||
def extra_headers(self) -> dict[str, str] | None:
|
||||
# Called after wrap_request(); relies on KiroRequestContext.
|
||||
from src.services.provider.adapters.kiro.context import get_kiro_request_context
|
||||
|
||||
ctx = get_kiro_request_context()
|
||||
if ctx is None:
|
||||
return None
|
||||
|
||||
host = f"q.{ctx.region}.amazonaws.com"
|
||||
return build_generate_assistant_headers(
|
||||
host=host,
|
||||
machine_id=ctx.machine_id,
|
||||
kiro_version=ctx.kiro_version,
|
||||
system_version=ctx.system_version,
|
||||
node_version=ctx.node_version,
|
||||
)
|
||||
|
||||
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]:
|
||||
cfg = KiroAuthConfig.from_dict(decrypted_auth_config or {})
|
||||
|
||||
region = _resolve_region(cfg)
|
||||
machine_id = generate_machine_id(cfg)
|
||||
|
||||
thinking_enabled = _is_thinking_enabled(request_body)
|
||||
|
||||
set_kiro_request_context(
|
||||
KiroRequestContext(
|
||||
region=region,
|
||||
machine_id=machine_id,
|
||||
kiro_version=cfg.kiro_version,
|
||||
system_version=cfg.system_version,
|
||||
node_version=cfg.node_version,
|
||||
thinking_enabled=thinking_enabled,
|
||||
)
|
||||
)
|
||||
|
||||
conversation_state = convert_claude_messages_to_conversation_state(
|
||||
request_body,
|
||||
model=model,
|
||||
)
|
||||
|
||||
wrapped: dict[str, Any] = {
|
||||
"conversationState": conversation_state,
|
||||
}
|
||||
if isinstance(cfg.profile_arn, str) and cfg.profile_arn.strip():
|
||||
wrapped["profileArn"] = cfg.profile_arn.strip()
|
||||
|
||||
return wrapped, url_model
|
||||
|
||||
def unwrap_response(self, data: Any) -> Any:
|
||||
return data
|
||||
|
||||
def postprocess_unwrapped_response(self, *, model: str, data: Any) -> None: # noqa: ARG002
|
||||
return
|
||||
|
||||
def capture_selected_base_url(self) -> str | None:
|
||||
return None
|
||||
|
||||
def on_http_status(self, *, base_url: str | None, status_code: int) -> None: # noqa: ARG002
|
||||
return
|
||||
|
||||
def on_connection_error(self, *, base_url: str | None, exc: Exception) -> None: # noqa: ARG002
|
||||
return
|
||||
|
||||
def force_stream_rewrite(self) -> bool:
|
||||
# Kiro streaming is binary AWS Event Stream and must be rewritten.
|
||||
return True
|
||||
|
||||
|
||||
kiro_envelope = KiroEnvelope()
|
||||
|
||||
|
||||
__all__ = ["KiroEnvelope", "kiro_envelope"]
|
||||
741
src/services/provider/adapters/kiro/eventstream_rewriter.py
Normal file
741
src/services/provider/adapters/kiro/eventstream_rewriter.py
Normal file
@@ -0,0 +1,741 @@
|
||||
"""AWS Event Stream -> Claude SSE rewriter for Kiro.
|
||||
|
||||
Kiro streaming responses are returned as `application/vnd.amazon.eventstream`
|
||||
(binary framed). This module decodes frames and emits Claude-style streaming
|
||||
SSE events (as UTF-8 bytes).
|
||||
|
||||
The output format uses ``event: {type}\\ndata: {...}\\n\\n`` for typed events and
|
||||
plain ``data: {...}\\n\\n`` for untyped events, matching how Aether parses Claude
|
||||
streams.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from collections.abc import AsyncGenerator
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.services.provider.adapters.kiro.constants import CONTEXT_WINDOW_TOKENS
|
||||
from src.services.provider.adapters.kiro.parser.decoder import EventStreamDecoder
|
||||
|
||||
# Safety limit for thinking_buffer to prevent memory exhaustion from
|
||||
# pathological upstream responses that never close the thinking tag.
|
||||
_MAX_THINKING_BUFFER = 1024 * 1024 # 1 MiB
|
||||
|
||||
_QUOTE_CHARS: frozenset[str] = frozenset("`\"'\\#!@$%^&*()-_=+[]{};:<>,.?/")
|
||||
|
||||
|
||||
def _is_quote_char(buffer: str, pos: int) -> bool:
|
||||
if pos < 0 or pos >= len(buffer):
|
||||
return False
|
||||
return buffer[pos] in _QUOTE_CHARS
|
||||
|
||||
|
||||
def _find_real_thinking_start_tag(buffer: str) -> int | None:
|
||||
tag = "<thinking>"
|
||||
search = 0
|
||||
while True:
|
||||
pos = buffer.find(tag, search)
|
||||
if pos < 0:
|
||||
return None
|
||||
has_before = pos > 0 and _is_quote_char(buffer, pos - 1)
|
||||
after_pos = pos + len(tag)
|
||||
has_after = _is_quote_char(buffer, after_pos)
|
||||
if not has_before and not has_after:
|
||||
return pos
|
||||
search = pos + 1
|
||||
|
||||
|
||||
def _find_real_thinking_end_tag(buffer: str) -> int | None:
|
||||
tag = "</thinking>"
|
||||
search = 0
|
||||
while True:
|
||||
pos = buffer.find(tag, search)
|
||||
if pos < 0:
|
||||
return None
|
||||
|
||||
has_before = pos > 0 and _is_quote_char(buffer, pos - 1)
|
||||
after_pos = pos + len(tag)
|
||||
has_after = _is_quote_char(buffer, after_pos)
|
||||
if has_before or has_after:
|
||||
search = pos + 1
|
||||
continue
|
||||
|
||||
after = buffer[after_pos:]
|
||||
if len(after) < 2:
|
||||
return None
|
||||
if after.startswith("\n\n"):
|
||||
return pos
|
||||
|
||||
search = pos + 1
|
||||
|
||||
|
||||
def _find_real_thinking_end_tag_at_buffer_end(buffer: str) -> int | None:
|
||||
tag = "</thinking>"
|
||||
search = 0
|
||||
while True:
|
||||
pos = buffer.find(tag, search)
|
||||
if pos < 0:
|
||||
return None
|
||||
|
||||
has_before = pos > 0 and _is_quote_char(buffer, pos - 1)
|
||||
after_pos = pos + len(tag)
|
||||
has_after = _is_quote_char(buffer, after_pos)
|
||||
if has_before or has_after:
|
||||
search = pos + 1
|
||||
continue
|
||||
|
||||
if buffer[after_pos:].strip() == "":
|
||||
return pos
|
||||
|
||||
search = pos + 1
|
||||
|
||||
|
||||
def _estimate_tokens(text: str) -> int:
|
||||
if not text:
|
||||
return 0
|
||||
chinese = 0
|
||||
other = 0
|
||||
for c in text:
|
||||
if "\u4e00" <= c <= "\u9fff":
|
||||
chinese += 1
|
||||
else:
|
||||
other += 1
|
||||
chinese_tokens = (chinese * 2 + 2) // 3
|
||||
other_tokens = (other + 3) // 4
|
||||
return max(chinese_tokens + other_tokens, 1)
|
||||
|
||||
|
||||
def _sse_data_bytes(obj: dict[str, Any]) -> bytes:
|
||||
data = json.dumps(obj, ensure_ascii=False)
|
||||
event_type = obj.get("type", "")
|
||||
if event_type:
|
||||
return f"event: {event_type}\ndata: {data}\n\n".encode("utf-8")
|
||||
return f"data: {data}\n\n".encode("utf-8")
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class _KiroStreamState:
|
||||
model: str
|
||||
thinking_enabled: bool
|
||||
estimated_input_tokens: int = 0
|
||||
|
||||
message_id: str = field(default_factory=lambda: f"msg_{uuid.uuid4().hex}")
|
||||
output_tokens: int = 0
|
||||
context_input_tokens: int | None = None
|
||||
|
||||
next_block_index: int = 0
|
||||
open_blocks: dict[int, str] = field(default_factory=dict)
|
||||
|
||||
text_block_index: int | None = None
|
||||
thinking_block_index: int | None = None
|
||||
tool_block_indices: dict[str, int] = field(default_factory=dict)
|
||||
|
||||
thinking_buffer: str = ""
|
||||
in_thinking_block: bool = False
|
||||
thinking_extracted: bool = False
|
||||
strip_thinking_leading_newline: bool = False
|
||||
|
||||
has_tool_use: bool = False
|
||||
stop_reason_override: str | None = None
|
||||
had_error: bool = False
|
||||
|
||||
def generate_initial_events(self) -> list[dict[str, Any]]:
|
||||
events: list[dict[str, Any]] = []
|
||||
|
||||
# message_start
|
||||
events.append(
|
||||
{
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"id": self.message_id,
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [],
|
||||
"model": self.model,
|
||||
"stop_reason": None,
|
||||
"stop_sequence": None,
|
||||
# Claude CLI clients expect usage to exist.
|
||||
"usage": {
|
||||
"input_tokens": int(self.estimated_input_tokens or 0),
|
||||
"output_tokens": 1,
|
||||
},
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
if not self.thinking_enabled:
|
||||
events.extend(self._ensure_text_block_open())
|
||||
|
||||
return events
|
||||
|
||||
def _ensure_text_block_open(self) -> list[dict[str, Any]]:
|
||||
if self.text_block_index is not None:
|
||||
if (
|
||||
self.text_block_index in self.open_blocks
|
||||
and self.open_blocks[self.text_block_index] == "text"
|
||||
):
|
||||
return []
|
||||
self.text_block_index = None
|
||||
|
||||
idx = self.next_block_index
|
||||
self.next_block_index += 1
|
||||
self.text_block_index = idx
|
||||
self.open_blocks[idx] = "text"
|
||||
|
||||
return [
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": idx,
|
||||
"content_block": {"type": "text", "text": ""},
|
||||
}
|
||||
]
|
||||
|
||||
def _close_block(self, idx: int) -> list[dict[str, Any]]:
|
||||
if idx not in self.open_blocks:
|
||||
return []
|
||||
self.open_blocks.pop(idx, None)
|
||||
return [{"type": "content_block_stop", "index": idx}]
|
||||
|
||||
def _ensure_thinking_block_open(self) -> list[dict[str, Any]]:
|
||||
if self.thinking_block_index is not None:
|
||||
if (
|
||||
self.thinking_block_index in self.open_blocks
|
||||
and self.open_blocks[self.thinking_block_index] == "thinking"
|
||||
):
|
||||
return []
|
||||
|
||||
idx = self.next_block_index
|
||||
self.next_block_index += 1
|
||||
self.thinking_block_index = idx
|
||||
self.open_blocks[idx] = "thinking"
|
||||
|
||||
return [
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": idx,
|
||||
"content_block": {"type": "thinking", "thinking": ""},
|
||||
}
|
||||
]
|
||||
|
||||
def _emit_text_delta(self, text: str) -> list[dict[str, Any]]:
|
||||
if not text:
|
||||
return []
|
||||
events: list[dict[str, Any]] = []
|
||||
events.extend(self._ensure_text_block_open())
|
||||
idx = int(self.text_block_index or 0)
|
||||
events.append(
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": idx,
|
||||
"delta": {"type": "text_delta", "text": text},
|
||||
}
|
||||
)
|
||||
return events
|
||||
|
||||
def _emit_thinking_delta(self, thinking: str) -> list[dict[str, Any]]:
|
||||
if not thinking:
|
||||
return []
|
||||
events: list[dict[str, Any]] = []
|
||||
events.extend(self._ensure_thinking_block_open())
|
||||
idx = int(self.thinking_block_index or 0)
|
||||
events.append(
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": idx,
|
||||
"delta": {"type": "thinking_delta", "thinking": thinking},
|
||||
}
|
||||
)
|
||||
return events
|
||||
|
||||
def _close_thinking_block(self) -> list[dict[str, Any]]:
|
||||
"""Send an empty thinking_delta sentinel and close the thinking block."""
|
||||
if self.thinking_block_index is None:
|
||||
return []
|
||||
idx = int(self.thinking_block_index)
|
||||
events: list[dict[str, Any]] = [
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": idx,
|
||||
"delta": {"type": "thinking_delta", "thinking": ""},
|
||||
}
|
||||
]
|
||||
events.extend(self._close_block(idx))
|
||||
return events
|
||||
|
||||
def process_context_usage(self, percentage: float) -> None:
|
||||
try:
|
||||
pct = float(percentage)
|
||||
except Exception:
|
||||
return
|
||||
# percentage * CONTEXT_WINDOW_TOKENS / 100
|
||||
self.context_input_tokens = int(pct * float(CONTEXT_WINDOW_TOKENS) / 100.0)
|
||||
|
||||
def process_exception(self, exception_type: str) -> None:
|
||||
if exception_type == "ContentLengthExceededException":
|
||||
# ContentLengthExceededException is a normal completion signal (output
|
||||
# exceeded size limit), not a fatal error. We record the stop_reason
|
||||
# but do NOT set had_error so that finalize() still emits message_delta
|
||||
# with stop_reason="max_tokens" and message_stop.
|
||||
self.stop_reason_override = "max_tokens"
|
||||
return
|
||||
|
||||
def process_assistant_response(self, content: str) -> list[dict[str, Any]]:
|
||||
if not content:
|
||||
return []
|
||||
|
||||
self.output_tokens += _estimate_tokens(content)
|
||||
|
||||
if not self.thinking_enabled:
|
||||
return self._emit_text_delta(content)
|
||||
|
||||
self.thinking_buffer += content
|
||||
|
||||
# Safety: flush as text if thinking_buffer grows too large without closing tag
|
||||
if len(self.thinking_buffer) > _MAX_THINKING_BUFFER:
|
||||
logger.warning(
|
||||
"kiro thinking_buffer exceeded {} bytes, force-flushing as text",
|
||||
_MAX_THINKING_BUFFER,
|
||||
)
|
||||
overflow = self.thinking_buffer
|
||||
self.thinking_buffer = ""
|
||||
if self.in_thinking_block:
|
||||
result = self._emit_thinking_delta(overflow)
|
||||
result.extend(self._close_thinking_block())
|
||||
self.in_thinking_block = False
|
||||
self.thinking_extracted = True
|
||||
return result
|
||||
return self._emit_text_delta(overflow)
|
||||
|
||||
events: list[dict[str, Any]] = []
|
||||
|
||||
while True:
|
||||
if not self.in_thinking_block and not self.thinking_extracted:
|
||||
start_pos = _find_real_thinking_start_tag(self.thinking_buffer)
|
||||
if start_pos is not None:
|
||||
before = self.thinking_buffer[:start_pos]
|
||||
if before and before.strip():
|
||||
events.extend(self._emit_text_delta(before))
|
||||
|
||||
self.in_thinking_block = True
|
||||
self.strip_thinking_leading_newline = True
|
||||
self.thinking_buffer = self.thinking_buffer[start_pos + len("<thinking>") :]
|
||||
events.extend(self._ensure_thinking_block_open())
|
||||
continue
|
||||
|
||||
# Keep a short suffix in buffer for partial tag detection.
|
||||
keep = len("<thinking>")
|
||||
if len(self.thinking_buffer) > keep:
|
||||
safe = self.thinking_buffer[:-keep]
|
||||
if safe and safe.strip():
|
||||
events.extend(self._emit_text_delta(safe))
|
||||
self.thinking_buffer = self.thinking_buffer[-keep:]
|
||||
break
|
||||
|
||||
if self.in_thinking_block:
|
||||
# Strip a single leading \n after <thinking> tag.
|
||||
# The model outputs `<thinking>\n` and the \n may arrive in the
|
||||
# same chunk or the next one; we drop it for cleaner output.
|
||||
if self.strip_thinking_leading_newline:
|
||||
if self.thinking_buffer.startswith("\n"):
|
||||
self.thinking_buffer = self.thinking_buffer[1:]
|
||||
self.strip_thinking_leading_newline = False
|
||||
elif self.thinking_buffer:
|
||||
# Buffer is non-empty but doesn't start with \n; stop waiting.
|
||||
self.strip_thinking_leading_newline = False
|
||||
# else: buffer is empty, keep the flag for the next chunk.
|
||||
|
||||
end_pos = _find_real_thinking_end_tag(self.thinking_buffer)
|
||||
if end_pos is not None:
|
||||
thinking_text = self.thinking_buffer[:end_pos]
|
||||
if thinking_text:
|
||||
events.extend(self._emit_thinking_delta(thinking_text))
|
||||
|
||||
events.extend(self._close_thinking_block())
|
||||
|
||||
self.in_thinking_block = False
|
||||
self.thinking_extracted = True
|
||||
self.thinking_buffer = self.thinking_buffer[end_pos + len("</thinking>") :]
|
||||
continue
|
||||
|
||||
keep = len("</thinking>")
|
||||
if len(self.thinking_buffer) > keep:
|
||||
safe = self.thinking_buffer[:-keep]
|
||||
if safe:
|
||||
events.extend(self._emit_thinking_delta(safe))
|
||||
self.thinking_buffer = self.thinking_buffer[-keep:]
|
||||
break
|
||||
|
||||
# thinking extracted: remaining buffer is text
|
||||
if self.thinking_buffer:
|
||||
remaining = self.thinking_buffer
|
||||
self.thinking_buffer = ""
|
||||
events.extend(self._emit_text_delta(remaining))
|
||||
break
|
||||
|
||||
return events
|
||||
|
||||
def process_tool_use(
|
||||
self,
|
||||
*,
|
||||
name: str,
|
||||
tool_use_id: str,
|
||||
input_json: str,
|
||||
stop: bool,
|
||||
) -> list[dict[str, Any]]:
|
||||
if not tool_use_id:
|
||||
return []
|
||||
|
||||
self.has_tool_use = True
|
||||
|
||||
events: list[dict[str, Any]] = []
|
||||
|
||||
# Boundary: close thinking block if needed, filtering a dangling </thinking>.
|
||||
if self.thinking_enabled and self.in_thinking_block and self.thinking_buffer:
|
||||
end_pos = _find_real_thinking_end_tag_at_buffer_end(self.thinking_buffer)
|
||||
if end_pos is not None:
|
||||
thinking_text = self.thinking_buffer[:end_pos]
|
||||
if thinking_text:
|
||||
events.extend(self._emit_thinking_delta(thinking_text))
|
||||
|
||||
events.extend(self._close_thinking_block())
|
||||
|
||||
after_pos = end_pos + len("</thinking>")
|
||||
remaining = self.thinking_buffer[after_pos:]
|
||||
self.thinking_buffer = ""
|
||||
self.in_thinking_block = False
|
||||
self.thinking_extracted = True
|
||||
if remaining:
|
||||
events.extend(self._emit_text_delta(remaining))
|
||||
else:
|
||||
# Best-effort flush all as thinking
|
||||
events.extend(self._emit_thinking_delta(self.thinking_buffer))
|
||||
events.extend(self._close_thinking_block())
|
||||
self.thinking_buffer = ""
|
||||
self.in_thinking_block = False
|
||||
self.thinking_extracted = True
|
||||
|
||||
# Flush any buffered pre-thinking tail so tool_use doesn't swallow it.
|
||||
if (
|
||||
self.thinking_enabled
|
||||
and not self.in_thinking_block
|
||||
and not self.thinking_extracted
|
||||
and self.thinking_buffer
|
||||
):
|
||||
buffered = self.thinking_buffer
|
||||
self.thinking_buffer = ""
|
||||
events.extend(self._emit_text_delta(buffered))
|
||||
|
||||
# Close current text block before tool_use.
|
||||
if self.text_block_index is not None:
|
||||
idx = int(self.text_block_index)
|
||||
events.extend(self._close_block(idx))
|
||||
|
||||
block_index = self.tool_block_indices.get(tool_use_id)
|
||||
if block_index is None:
|
||||
block_index = self.next_block_index
|
||||
self.next_block_index += 1
|
||||
self.tool_block_indices[tool_use_id] = block_index
|
||||
|
||||
# Start tool block if not open.
|
||||
if block_index not in self.open_blocks:
|
||||
self.open_blocks[block_index] = "tool_use"
|
||||
events.append(
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": block_index,
|
||||
"content_block": {
|
||||
"type": "tool_use",
|
||||
"id": tool_use_id,
|
||||
"name": name,
|
||||
"input": {},
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
if input_json:
|
||||
self.output_tokens += _estimate_tokens(input_json)
|
||||
events.append(
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": block_index,
|
||||
"delta": {"type": "input_json_delta", "partial_json": input_json},
|
||||
}
|
||||
)
|
||||
|
||||
if stop:
|
||||
events.extend(self._close_block(block_index))
|
||||
|
||||
return events
|
||||
|
||||
def finalize(self) -> list[dict[str, Any]]:
|
||||
events: list[dict[str, Any]] = []
|
||||
|
||||
# Flush remaining thinking/text buffer.
|
||||
if self.thinking_enabled and self.thinking_buffer:
|
||||
if self.in_thinking_block:
|
||||
end_pos = _find_real_thinking_end_tag_at_buffer_end(self.thinking_buffer)
|
||||
if end_pos is not None:
|
||||
thinking_text = self.thinking_buffer[:end_pos]
|
||||
if thinking_text:
|
||||
events.extend(self._emit_thinking_delta(thinking_text))
|
||||
|
||||
events.extend(self._close_thinking_block())
|
||||
|
||||
after_pos = end_pos + len("</thinking>")
|
||||
remaining = self.thinking_buffer[after_pos:]
|
||||
if remaining:
|
||||
events.extend(self._emit_text_delta(remaining))
|
||||
else:
|
||||
events.extend(self._emit_thinking_delta(self.thinking_buffer))
|
||||
events.extend(self._close_thinking_block())
|
||||
|
||||
else:
|
||||
events.extend(self._emit_text_delta(self.thinking_buffer))
|
||||
|
||||
self.thinking_buffer = ""
|
||||
self.in_thinking_block = False
|
||||
self.thinking_extracted = True
|
||||
|
||||
# Close any open blocks (best-effort).
|
||||
for idx in sorted(list(self.open_blocks.keys()), reverse=True):
|
||||
events.extend(self._close_block(idx))
|
||||
|
||||
stop_reason = self.stop_reason_override
|
||||
if not stop_reason:
|
||||
stop_reason = "tool_use" if self.has_tool_use else "end_turn"
|
||||
|
||||
input_tokens = (
|
||||
int(self.context_input_tokens)
|
||||
if self.context_input_tokens is not None
|
||||
else int(self.estimated_input_tokens or 0)
|
||||
)
|
||||
|
||||
events.append(
|
||||
{
|
||||
"type": "message_delta",
|
||||
"delta": {"stop_reason": stop_reason, "stop_sequence": None},
|
||||
"usage": {"input_tokens": input_tokens, "output_tokens": int(self.output_tokens)},
|
||||
}
|
||||
)
|
||||
events.append({"type": "message_stop"})
|
||||
|
||||
return events
|
||||
|
||||
|
||||
async def rewrite_eventstream_to_sse(
|
||||
byte_iterator: Any,
|
||||
*,
|
||||
model: str,
|
||||
thinking_enabled: bool,
|
||||
estimated_input_tokens: int = 0,
|
||||
) -> AsyncGenerator[bytes]:
|
||||
"""Rewrite Kiro AWS Event Stream bytes to Claude SSE bytes."""
|
||||
decoder = EventStreamDecoder()
|
||||
state = _KiroStreamState(
|
||||
model=str(model or ""),
|
||||
thinking_enabled=bool(thinking_enabled),
|
||||
estimated_input_tokens=int(estimated_input_tokens or 0),
|
||||
)
|
||||
|
||||
# 收集原始字节用于错误诊断
|
||||
raw_bytes_buffer = b""
|
||||
|
||||
# Initial events
|
||||
for evt in state.generate_initial_events():
|
||||
yield _sse_data_bytes(evt)
|
||||
|
||||
async for chunk in byte_iterator:
|
||||
if not chunk:
|
||||
continue
|
||||
|
||||
# 保留原始字节用于错误诊断（限制大小）
|
||||
if len(raw_bytes_buffer) < 4096:
|
||||
raw_bytes_buffer += chunk
|
||||
|
||||
try:
|
||||
decoder.feed(chunk)
|
||||
frames = decoder.decode_available()
|
||||
except Exception as e:
|
||||
logger.warning("kiro eventstream decode error: {}", e)
|
||||
# 尝试解析原始响应为 JSON 错误
|
||||
error_message = f"kiro eventstream decode failed: {type(e).__name__}"
|
||||
try:
|
||||
raw_text = raw_bytes_buffer.decode("utf-8", errors="replace")
|
||||
# 尝试解析为 JSON
|
||||
error_json = json.loads(raw_text)
|
||||
if isinstance(error_json, dict):
|
||||
# 提取上游错误信息
|
||||
upstream_msg = error_json.get("message") or error_json.get("error", {}).get(
|
||||
"message"
|
||||
)
|
||||
if upstream_msg:
|
||||
error_message = f"Kiro API error: {upstream_msg}"
|
||||
except Exception:
|
||||
pass
|
||||
yield _sse_data_bytes(
|
||||
{
|
||||
"type": "error",
|
||||
"error": {
|
||||
"type": "upstream_stream_error",
|
||||
"message": error_message,
|
||||
},
|
||||
}
|
||||
)
|
||||
break
|
||||
|
||||
for frame in frames:
|
||||
mtype = (frame.message_type() or "event").strip().lower()
|
||||
etype = (frame.event_type() or "").strip()
|
||||
payload_text = frame.payload_as_text()
|
||||
|
||||
if mtype == "event":
|
||||
try:
|
||||
payload = json.loads(payload_text) if payload_text else {}
|
||||
except Exception:
|
||||
payload = {}
|
||||
|
||||
if etype == "assistantResponseEvent":
|
||||
content = payload.get("content") if isinstance(payload, dict) else None
|
||||
if isinstance(content, str) and content:
|
||||
for evt in state.process_assistant_response(content):
|
||||
yield _sse_data_bytes(evt)
|
||||
continue
|
||||
|
||||
if etype == "toolUseEvent":
|
||||
if isinstance(payload, dict):
|
||||
name = str(payload.get("name") or "")
|
||||
tool_use_id = payload.get("toolUseId") or payload.get("tool_use_id")
|
||||
tool_use_id = str(tool_use_id or "")
|
||||
raw_input = payload.get("input")
|
||||
if raw_input is None:
|
||||
input_json = ""
|
||||
elif isinstance(raw_input, str):
|
||||
input_json = raw_input
|
||||
else:
|
||||
try:
|
||||
input_json = json.dumps(raw_input, ensure_ascii=False)
|
||||
except Exception:
|
||||
input_json = str(raw_input)
|
||||
stop = bool(payload.get("stop", False))
|
||||
for evt in state.process_tool_use(
|
||||
name=name,
|
||||
tool_use_id=tool_use_id,
|
||||
input_json=input_json,
|
||||
stop=stop,
|
||||
):
|
||||
yield _sse_data_bytes(evt)
|
||||
continue
|
||||
|
||||
if etype == "contextUsageEvent":
|
||||
if isinstance(payload, dict):
|
||||
pct = payload.get("contextUsagePercentage")
|
||||
if pct is not None:
|
||||
try:
|
||||
state.process_context_usage(float(pct))
|
||||
except (ValueError, TypeError):
|
||||
logger.debug(
|
||||
"kiro: failed to parse contextUsagePercentage: {!r}", pct
|
||||
)
|
||||
continue
|
||||
|
||||
# meteringEvent / unknown: ignore
|
||||
continue
|
||||
|
||||
if mtype == "exception":
|
||||
ex_type = frame.headers.exception_type() or "UnknownException"
|
||||
state.process_exception(ex_type)
|
||||
# ContentLengthExceededException is handled by process_exception
|
||||
# (sets stop_reason_override) and should NOT prevent finalize().
|
||||
if not state.stop_reason_override:
|
||||
state.had_error = True
|
||||
logger.debug("kiro upstream exception: {} | {}", ex_type, payload_text[:200])
|
||||
if state.had_error:
|
||||
yield _sse_data_bytes(
|
||||
{
|
||||
"type": "error",
|
||||
"error": {
|
||||
"type": "upstream_exception",
|
||||
"message": ex_type,
|
||||
},
|
||||
}
|
||||
)
|
||||
continue
|
||||
|
||||
if mtype == "error":
|
||||
err_code = frame.headers.error_code() or "UnknownError"
|
||||
state.had_error = True
|
||||
logger.debug("kiro upstream error: {} | {}", err_code, payload_text[:200])
|
||||
yield _sse_data_bytes(
|
||||
{
|
||||
"type": "error",
|
||||
"error": {
|
||||
"type": "upstream_error",
|
||||
"message": err_code,
|
||||
},
|
||||
}
|
||||
)
|
||||
continue
|
||||
|
||||
if not state.had_error:
|
||||
for evt in state.finalize():
|
||||
yield _sse_data_bytes(evt)
|
||||
|
||||
|
||||
def apply_kiro_stream_rewrite(
|
||||
byte_iter: Any,
|
||||
*,
|
||||
model: str = "",
|
||||
input_tokens: int = 0,
|
||||
prefetched_chunks: list[bytes] | None = None,
|
||||
) -> AsyncGenerator[bytes]:
|
||||
"""Apply Kiro EventStream->SSE rewrite if context is available.
|
||||
|
||||
Consolidates the repeated import-context-rewrite pattern used across
|
||||
``chat_handler_base``, ``cli_handler_base``, and ``stream_processor``.
|
||||
|
||||
Args:
|
||||
byte_iter: Upstream byte iterator (raw AWS Event Stream).
|
||||
model: Model name for SSE events.
|
||||
input_tokens: Estimated input token count.
|
||||
prefetched_chunks: Optional pre-fetched bytes to prepend.
|
||||
|
||||
Returns:
|
||||
An async generator of Claude-compatible SSE bytes.
|
||||
"""
|
||||
from src.services.provider.adapters.kiro.context import get_kiro_request_context
|
||||
|
||||
kiro_ctx = get_kiro_request_context()
|
||||
thinking_enabled = bool(getattr(kiro_ctx, "thinking_enabled", False)) if kiro_ctx else False
|
||||
|
||||
if prefetched_chunks:
|
||||
upstream = byte_iter
|
||||
prefix = list(prefetched_chunks)
|
||||
|
||||
async def _combined() -> AsyncGenerator[bytes, None]:
|
||||
for c in prefix:
|
||||
if c:
|
||||
yield c
|
||||
async for c in upstream:
|
||||
if c:
|
||||
yield c
|
||||
|
||||
source: Any = _combined()
|
||||
else:
|
||||
source = byte_iter
|
||||
|
||||
return rewrite_eventstream_to_sse(
|
||||
source,
|
||||
model=str(model or ""),
|
||||
thinking_enabled=thinking_enabled,
|
||||
estimated_input_tokens=int(input_tokens or 0),
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"apply_kiro_stream_rewrite",
|
||||
"rewrite_eventstream_to_sse",
|
||||
]
|
||||
102
src/services/provider/adapters/kiro/headers.py
Normal file
102
src/services/provider/adapters/kiro/headers.py
Normal file
@@ -0,0 +1,102 @@
|
||||
"""Kiro header builders."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
|
||||
from src.services.provider.adapters.kiro.constants import (
|
||||
AWS_EVENTSTREAM_CONTENT_TYPE,
|
||||
AWS_SDK_JS_MAIN_VERSION,
|
||||
AWS_SDK_JS_USAGE_VERSION,
|
||||
CODEWHISPERER_OPTOUT,
|
||||
DEFAULT_KIRO_VERSION,
|
||||
DEFAULT_NODE_VERSION,
|
||||
DEFAULT_SYSTEM_VERSION,
|
||||
KIRO_AGENT_MODE,
|
||||
)
|
||||
|
||||
|
||||
def build_kiro_ide_tag(*, kiro_version: str, machine_id: str) -> str:
|
||||
version = (kiro_version or DEFAULT_KIRO_VERSION).strip() or DEFAULT_KIRO_VERSION
|
||||
mid = (machine_id or "").strip()
|
||||
return f"KiroIDE-{version}-{mid}" if mid else f"KiroIDE-{version}"
|
||||
|
||||
|
||||
def build_x_amz_user_agent_main(*, kiro_version: str, machine_id: str) -> str:
|
||||
return f"aws-sdk-js/{AWS_SDK_JS_MAIN_VERSION} {build_kiro_ide_tag(kiro_version=kiro_version, machine_id=machine_id)}"
|
||||
|
||||
|
||||
def build_user_agent_main(
|
||||
*, system_version: str, node_version: str, kiro_version: str, machine_id: str
|
||||
) -> str:
|
||||
os_tag = (system_version or DEFAULT_SYSTEM_VERSION).strip() or DEFAULT_SYSTEM_VERSION
|
||||
node_tag = (node_version or DEFAULT_NODE_VERSION).strip() or DEFAULT_NODE_VERSION
|
||||
ide = build_kiro_ide_tag(kiro_version=kiro_version, machine_id=machine_id)
|
||||
return (
|
||||
f"aws-sdk-js/{AWS_SDK_JS_MAIN_VERSION} ua/2.1 os/{os_tag} lang/js "
|
||||
f"md/nodejs#{node_tag} api/codewhispererstreaming#{AWS_SDK_JS_MAIN_VERSION} m/E {ide}"
|
||||
)
|
||||
|
||||
|
||||
def build_x_amz_user_agent_usage(*, kiro_version: str, machine_id: str) -> str:
|
||||
ide = build_kiro_ide_tag(kiro_version=kiro_version, machine_id=machine_id)
|
||||
return f"aws-sdk-js/{AWS_SDK_JS_USAGE_VERSION} {ide}"
|
||||
|
||||
|
||||
def build_user_agent_usage(*, kiro_version: str, machine_id: str) -> str:
|
||||
ide = build_kiro_ide_tag(kiro_version=kiro_version, machine_id=machine_id)
|
||||
os_tag = DEFAULT_SYSTEM_VERSION
|
||||
node_tag = DEFAULT_NODE_VERSION
|
||||
return (
|
||||
f"aws-sdk-js/{AWS_SDK_JS_USAGE_VERSION} ua/2.1 os/{os_tag} lang/js "
|
||||
f"md/nodejs#{node_tag} api/codewhispererruntime#1.0.0 m/N,E {ide}"
|
||||
)
|
||||
|
||||
|
||||
def build_generate_assistant_headers(
|
||||
*,
|
||||
host: str,
|
||||
access_token: str | None = None,
|
||||
machine_id: str,
|
||||
kiro_version: str | None = None,
|
||||
system_version: str | None = None,
|
||||
node_version: str | None = None,
|
||||
) -> dict[str, str]:
|
||||
version = (kiro_version or DEFAULT_KIRO_VERSION).strip() or DEFAULT_KIRO_VERSION
|
||||
sys_ver = (system_version or DEFAULT_SYSTEM_VERSION).strip() or DEFAULT_SYSTEM_VERSION
|
||||
node_ver = (node_version or DEFAULT_NODE_VERSION).strip() or DEFAULT_NODE_VERSION
|
||||
|
||||
headers: dict[str, str] = {
|
||||
"Content-Type": "application/json",
|
||||
"Accept": AWS_EVENTSTREAM_CONTENT_TYPE,
|
||||
"host": host,
|
||||
"Connection": "close",
|
||||
"x-amzn-codewhisperer-optout": CODEWHISPERER_OPTOUT,
|
||||
"x-amzn-kiro-agent-mode": KIRO_AGENT_MODE,
|
||||
"x-amz-user-agent": build_x_amz_user_agent_main(
|
||||
kiro_version=version, machine_id=machine_id
|
||||
),
|
||||
"User-Agent": build_user_agent_main(
|
||||
system_version=sys_ver,
|
||||
node_version=node_ver,
|
||||
kiro_version=version,
|
||||
machine_id=machine_id,
|
||||
),
|
||||
"amz-sdk-invocation-id": str(uuid.uuid4()),
|
||||
"amz-sdk-request": "attempt=1; max=3",
|
||||
}
|
||||
|
||||
if access_token:
|
||||
headers["Authorization"] = f"Bearer {access_token}"
|
||||
|
||||
return headers
|
||||
|
||||
|
||||
__all__ = [
|
||||
"build_generate_assistant_headers",
|
||||
"build_kiro_ide_tag",
|
||||
"build_user_agent_main",
|
||||
"build_user_agent_usage",
|
||||
"build_x_amz_user_agent_main",
|
||||
"build_x_amz_user_agent_usage",
|
||||
]
|
||||
21
src/services/provider/adapters/kiro/models/__init__.py
Normal file
21
src/services/provider/adapters/kiro/models/__init__.py
Normal file
@@ -0,0 +1,21 @@
|
||||
from .credentials import KiroAuthConfig
|
||||
from .usage_limits import (
|
||||
Bonus,
|
||||
FreeTrialInfo,
|
||||
SubscriptionInfo,
|
||||
UsageBreakdown,
|
||||
UsageLimitsResponse,
|
||||
calculate_current_usage,
|
||||
calculate_total_usage_limit,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"Bonus",
|
||||
"FreeTrialInfo",
|
||||
"KiroAuthConfig",
|
||||
"SubscriptionInfo",
|
||||
"UsageBreakdown",
|
||||
"UsageLimitsResponse",
|
||||
"calculate_current_usage",
|
||||
"calculate_total_usage_limit",
|
||||
]
|
||||
178
src/services/provider/adapters/kiro/models/credentials.py
Normal file
178
src/services/provider/adapters/kiro/models/credentials.py
Normal file
@@ -0,0 +1,178 @@
|
||||
"""Internal Kiro credential schema (stored in ProviderAPIKey.auth_config)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
|
||||
def _parse_epoch_seconds(value: object) -> int | None:
|
||||
if value is None:
|
||||
return None
|
||||
try:
|
||||
if isinstance(value, (int, float)):
|
||||
return int(value)
|
||||
if isinstance(value, str) and value.strip().isdigit():
|
||||
return int(value.strip())
|
||||
except Exception:
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def _get_str(raw: dict[str, Any], *keys: str) -> str | None:
|
||||
"""Return the first non-empty stripped string for *keys*, or ``None``."""
|
||||
for k in keys:
|
||||
v = raw.get(k)
|
||||
if isinstance(v, str) and v.strip():
|
||||
return v.strip()
|
||||
return None
|
||||
|
||||
|
||||
def _parse_iso_to_epoch_seconds(value: object) -> int | None:
|
||||
if not isinstance(value, str) or not value.strip():
|
||||
return None
|
||||
text = value.strip()
|
||||
# Support RFC3339 with Z suffix.
|
||||
if text.endswith("Z"):
|
||||
text = text[:-1] + "+00:00"
|
||||
try:
|
||||
dt = datetime.fromisoformat(text)
|
||||
except Exception:
|
||||
return None
|
||||
if dt.tzinfo is None:
|
||||
dt = dt.replace(tzinfo=timezone.utc)
|
||||
return int(dt.timestamp())
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class KiroAuthConfig:
|
||||
provider_type: str = "kiro"
|
||||
|
||||
auth_method: str = "social" # social | idc
|
||||
refresh_token: str = ""
|
||||
expires_at: int = 0
|
||||
|
||||
profile_arn: str | None = None
|
||||
region: str | None = None
|
||||
|
||||
client_id: str | None = None
|
||||
client_secret: str | None = None
|
||||
|
||||
machine_id: str | None = None
|
||||
kiro_version: str | None = None
|
||||
system_version: str | None = None
|
||||
node_version: str | None = None
|
||||
|
||||
email: str | None = None # 账号邮箱
|
||||
|
||||
# 缓存的 access_token（可选，用于避免频繁刷新）
|
||||
access_token: str | None = None
|
||||
|
||||
@staticmethod
|
||||
def infer_auth_method(raw: dict[str, Any]) -> str:
|
||||
"""
|
||||
根据凭据字段自动推断认证类型。
|
||||
|
||||
规则：
|
||||
- 包含 clientId + clientSecret -> IdC
|
||||
- 仅含 refreshToken -> Social
|
||||
"""
|
||||
client_id = raw.get("client_id") or raw.get("clientId")
|
||||
client_secret = raw.get("client_secret") or raw.get("clientSecret")
|
||||
|
||||
if client_id and client_secret:
|
||||
return "idc"
|
||||
return "social"
|
||||
|
||||
@staticmethod
|
||||
def validate_required_fields(raw: dict[str, Any]) -> tuple[bool, str]:
|
||||
"""
|
||||
验证凭据是否包含必需字段。
|
||||
|
||||
返回: (is_valid, error_message)
|
||||
"""
|
||||
refresh_token = raw.get("refresh_token") or raw.get("refreshToken") or ""
|
||||
refresh_token = str(refresh_token).strip()
|
||||
|
||||
if not refresh_token:
|
||||
return False, "refreshToken 为必填字段"
|
||||
|
||||
# refreshToken 不能含有 ...（表示被截断）
|
||||
if "..." in refresh_token:
|
||||
return False, "refreshToken 不完整（含有 ...），请导出完整的 Token"
|
||||
|
||||
# IdC 类型需要 clientId 和 clientSecret
|
||||
auth_method = KiroAuthConfig.infer_auth_method(raw)
|
||||
if auth_method == "idc":
|
||||
client_id = raw.get("client_id") or raw.get("clientId")
|
||||
client_secret = raw.get("client_secret") or raw.get("clientSecret")
|
||||
if not client_id:
|
||||
return False, "IdC 类型需要 clientId"
|
||||
if not client_secret:
|
||||
return False, "IdC 类型需要 clientSecret"
|
||||
|
||||
return True, ""
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, raw: dict[str, Any]) -> "KiroAuthConfig":
|
||||
if not isinstance(raw, dict):
|
||||
raw = {}
|
||||
|
||||
provider_type = _get_str(raw, "provider_type", "providerType") or "kiro"
|
||||
|
||||
# 自动推断 auth_method（如果未显式指定）
|
||||
explicit_method = _get_str(raw, "auth_method", "authMethod")
|
||||
auth_method = explicit_method.lower() if explicit_method else cls.infer_auth_method(raw)
|
||||
|
||||
refresh_token = (_get_str(raw, "refresh_token", "refreshToken") or "").strip()
|
||||
|
||||
expires_at = _parse_epoch_seconds(raw.get("expires_at"))
|
||||
if expires_at is None:
|
||||
expires_at = _parse_iso_to_epoch_seconds(raw.get("expiresAt"))
|
||||
if expires_at is None:
|
||||
expires_at = 0
|
||||
|
||||
cfg = cls(
|
||||
provider_type=provider_type,
|
||||
auth_method=(auth_method or "social").lower(),
|
||||
refresh_token=refresh_token,
|
||||
expires_at=int(expires_at),
|
||||
profile_arn=_get_str(raw, "profile_arn", "profileArn"),
|
||||
region=_get_str(raw, "region"),
|
||||
client_id=_get_str(raw, "client_id", "clientId"),
|
||||
client_secret=_get_str(raw, "client_secret", "clientSecret"),
|
||||
machine_id=_get_str(raw, "machine_id", "machineId"),
|
||||
kiro_version=_get_str(raw, "kiro_version", "kiroVersion"),
|
||||
system_version=_get_str(raw, "system_version", "systemVersion"),
|
||||
node_version=_get_str(raw, "node_version", "nodeVersion"),
|
||||
email=_get_str(raw, "email"),
|
||||
access_token=_get_str(raw, "access_token", "accessToken"),
|
||||
)
|
||||
|
||||
# Normalize auth_method aliases.
|
||||
if cfg.auth_method in {"builder-id", "builder_id", "iam"}:
|
||||
cfg.auth_method = "idc"
|
||||
|
||||
return cfg
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"provider_type": self.provider_type,
|
||||
"auth_method": self.auth_method,
|
||||
"refresh_token": self.refresh_token,
|
||||
"expires_at": self.expires_at,
|
||||
"profile_arn": self.profile_arn,
|
||||
"region": self.region,
|
||||
"client_id": self.client_id,
|
||||
"client_secret": self.client_secret,
|
||||
"machine_id": self.machine_id,
|
||||
"kiro_version": self.kiro_version,
|
||||
"system_version": self.system_version,
|
||||
"node_version": self.node_version,
|
||||
"email": self.email,
|
||||
"access_token": self.access_token,
|
||||
}
|
||||
|
||||
|
||||
__all__ = ["KiroAuthConfig"]
|
||||
222
src/services/provider/adapters/kiro/models/usage_limits.py
Normal file
222
src/services/provider/adapters/kiro/models/usage_limits.py
Normal file
@@ -0,0 +1,222 @@
|
||||
"""Kiro getUsageLimits response models (best-effort).
|
||||
|
||||
The AWS API uses camelCase fields; we parse defensively.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class SubscriptionInfo:
|
||||
subscription_title: str | None = None
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, raw: Any) -> "SubscriptionInfo | None":
|
||||
if not isinstance(raw, dict):
|
||||
return None
|
||||
title = raw.get("subscriptionTitle")
|
||||
if isinstance(title, str) and title.strip():
|
||||
return cls(subscription_title=title.strip())
|
||||
return cls(subscription_title=None)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class Bonus:
|
||||
current_usage: float = 0.0
|
||||
usage_limit: float = 0.0
|
||||
status: str | None = None
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, raw: Any) -> "Bonus | None":
|
||||
if not isinstance(raw, dict):
|
||||
return None
|
||||
status = raw.get("status")
|
||||
status_val = status.strip() if isinstance(status, str) and status.strip() else None
|
||||
cu = raw.get("currentUsage")
|
||||
ul = raw.get("usageLimit")
|
||||
try:
|
||||
cu_f = float(cu) if cu is not None else 0.0
|
||||
except Exception:
|
||||
cu_f = 0.0
|
||||
try:
|
||||
ul_f = float(ul) if ul is not None else 0.0
|
||||
except Exception:
|
||||
ul_f = 0.0
|
||||
return cls(current_usage=cu_f, usage_limit=ul_f, status=status_val)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class FreeTrialInfo:
|
||||
current_usage: int = 0
|
||||
current_usage_with_precision: float = 0.0
|
||||
usage_limit: int = 0
|
||||
usage_limit_with_precision: float = 0.0
|
||||
free_trial_expiry: float | None = None
|
||||
free_trial_status: str | None = None
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, raw: Any) -> "FreeTrialInfo | None":
|
||||
if not isinstance(raw, dict):
|
||||
return None
|
||||
|
||||
def _int(v: Any) -> int:
|
||||
try:
|
||||
return int(v)
|
||||
except Exception:
|
||||
return 0
|
||||
|
||||
def _float(v: Any) -> float:
|
||||
try:
|
||||
return float(v)
|
||||
except Exception:
|
||||
return 0.0
|
||||
|
||||
expiry = raw.get("freeTrialExpiry")
|
||||
try:
|
||||
expiry_f = float(expiry) if expiry is not None else None
|
||||
except Exception:
|
||||
expiry_f = None
|
||||
|
||||
status = raw.get("freeTrialStatus")
|
||||
status_val = status.strip() if isinstance(status, str) and status.strip() else None
|
||||
|
||||
return cls(
|
||||
current_usage=_int(raw.get("currentUsage")),
|
||||
current_usage_with_precision=_float(raw.get("currentUsageWithPrecision")),
|
||||
usage_limit=_int(raw.get("usageLimit")),
|
||||
usage_limit_with_precision=_float(raw.get("usageLimitWithPrecision")),
|
||||
free_trial_expiry=expiry_f,
|
||||
free_trial_status=status_val,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class UsageBreakdown:
|
||||
current_usage: int = 0
|
||||
current_usage_with_precision: float = 0.0
|
||||
usage_limit: int = 0
|
||||
usage_limit_with_precision: float = 0.0
|
||||
next_date_reset: float | None = None
|
||||
bonuses: list[Bonus] = field(default_factory=list)
|
||||
free_trial_info: FreeTrialInfo | None = None
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, raw: Any) -> "UsageBreakdown | None":
|
||||
if not isinstance(raw, dict):
|
||||
return None
|
||||
|
||||
def _int(v: Any) -> int:
|
||||
try:
|
||||
return int(v)
|
||||
except Exception:
|
||||
return 0
|
||||
|
||||
def _float(v: Any) -> float:
|
||||
try:
|
||||
return float(v)
|
||||
except Exception:
|
||||
return 0.0
|
||||
|
||||
next_reset = raw.get("nextDateReset")
|
||||
try:
|
||||
next_reset_f = float(next_reset) if next_reset is not None else None
|
||||
except Exception:
|
||||
next_reset_f = None
|
||||
|
||||
bonuses_raw = raw.get("bonuses")
|
||||
bonuses: list[Bonus] = []
|
||||
if isinstance(bonuses_raw, list):
|
||||
for b in bonuses_raw:
|
||||
parsed = Bonus.from_dict(b)
|
||||
if parsed is not None:
|
||||
bonuses.append(parsed)
|
||||
|
||||
return cls(
|
||||
current_usage=_int(raw.get("currentUsage")),
|
||||
current_usage_with_precision=_float(raw.get("currentUsageWithPrecision")),
|
||||
usage_limit=_int(raw.get("usageLimit")),
|
||||
usage_limit_with_precision=_float(raw.get("usageLimitWithPrecision")),
|
||||
next_date_reset=next_reset_f,
|
||||
bonuses=bonuses,
|
||||
free_trial_info=FreeTrialInfo.from_dict(raw.get("freeTrialInfo")),
|
||||
)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class UsageLimitsResponse:
|
||||
next_date_reset: float | None = None
|
||||
subscription_info: SubscriptionInfo | None = None
|
||||
usage_breakdown_list: list[UsageBreakdown] = field(default_factory=list)
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, raw: Any) -> "UsageLimitsResponse":
|
||||
if not isinstance(raw, dict):
|
||||
raw = {}
|
||||
|
||||
next_reset = raw.get("nextDateReset")
|
||||
try:
|
||||
next_reset_f = float(next_reset) if next_reset is not None else None
|
||||
except Exception:
|
||||
next_reset_f = None
|
||||
|
||||
breakdown_raw = raw.get("usageBreakdownList")
|
||||
breakdowns: list[UsageBreakdown] = []
|
||||
if isinstance(breakdown_raw, list):
|
||||
for b in breakdown_raw:
|
||||
parsed = UsageBreakdown.from_dict(b)
|
||||
if parsed is not None:
|
||||
breakdowns.append(parsed)
|
||||
|
||||
return cls(
|
||||
next_date_reset=next_reset_f,
|
||||
subscription_info=SubscriptionInfo.from_dict(raw.get("subscriptionInfo")),
|
||||
usage_breakdown_list=breakdowns,
|
||||
)
|
||||
|
||||
|
||||
def calculate_total_usage_limit(response: UsageLimitsResponse) -> float:
|
||||
if not response.usage_breakdown_list:
|
||||
return 0.0
|
||||
|
||||
breakdown = response.usage_breakdown_list[0]
|
||||
total = breakdown.usage_limit_with_precision
|
||||
|
||||
if breakdown.free_trial_info and breakdown.free_trial_info.free_trial_status == "ACTIVE":
|
||||
total += breakdown.free_trial_info.usage_limit_with_precision
|
||||
|
||||
for bonus in breakdown.bonuses:
|
||||
if bonus.status == "ACTIVE":
|
||||
total += bonus.usage_limit
|
||||
|
||||
return total
|
||||
|
||||
|
||||
def calculate_current_usage(response: UsageLimitsResponse) -> float:
|
||||
if not response.usage_breakdown_list:
|
||||
return 0.0
|
||||
|
||||
breakdown = response.usage_breakdown_list[0]
|
||||
total = breakdown.current_usage_with_precision
|
||||
|
||||
if breakdown.free_trial_info and breakdown.free_trial_info.free_trial_status == "ACTIVE":
|
||||
total += breakdown.free_trial_info.current_usage_with_precision
|
||||
|
||||
for bonus in breakdown.bonuses:
|
||||
if bonus.status == "ACTIVE":
|
||||
total += bonus.current_usage
|
||||
|
||||
return total
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Bonus",
|
||||
"FreeTrialInfo",
|
||||
"SubscriptionInfo",
|
||||
"UsageBreakdown",
|
||||
"UsageLimitsResponse",
|
||||
"calculate_current_usage",
|
||||
"calculate_total_usage_limit",
|
||||
]
|
||||
6
src/services/provider/adapters/kiro/parser/__init__.py
Normal file
6
src/services/provider/adapters/kiro/parser/__init__.py
Normal file
@@ -0,0 +1,6 @@
|
||||
"""AWS Event Stream parser for Kiro."""
|
||||
|
||||
from .decoder import EventStreamDecoder
|
||||
from .frame import Frame
|
||||
|
||||
__all__ = ["EventStreamDecoder", "Frame"]
|
||||
13
src/services/provider/adapters/kiro/parser/crc.py
Normal file
13
src/services/provider/adapters/kiro/parser/crc.py
Normal file
@@ -0,0 +1,13 @@
|
||||
"""CRC helpers for AWS Event Stream frames."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import binascii
|
||||
|
||||
|
||||
def crc32(data: bytes) -> int:
|
||||
"""Compute unsigned CRC32 (IEEE)."""
|
||||
return binascii.crc32(data) & 0xFFFFFFFF
|
||||
|
||||
|
||||
__all__ = ["crc32"]
|
||||
91
src/services/provider/adapters/kiro/parser/decoder.py
Normal file
91
src/services/provider/adapters/kiro/parser/decoder.py
Normal file
@@ -0,0 +1,91 @@
|
||||
"""Incremental AWS Event Stream decoder."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from .error import BufferOverflowError, EventStreamParseError
|
||||
from .frame import MAX_MESSAGE_SIZE, Frame, parse_frame
|
||||
|
||||
DEFAULT_MAX_BUFFER_SIZE = MAX_MESSAGE_SIZE
|
||||
DEFAULT_MAX_ERRORS = 5
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class DecoderStats:
|
||||
frames_decoded: int = 0
|
||||
bytes_skipped: int = 0
|
||||
error_count: int = 0
|
||||
|
||||
|
||||
class EventStreamDecoder:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
max_buffer_size: int = DEFAULT_MAX_BUFFER_SIZE,
|
||||
max_errors: int = DEFAULT_MAX_ERRORS,
|
||||
) -> None:
|
||||
self._buffer = bytearray()
|
||||
self._max_buffer_size = int(max_buffer_size)
|
||||
self._max_errors = int(max_errors)
|
||||
self._stopped = False
|
||||
self.stats = DecoderStats()
|
||||
|
||||
@property
|
||||
def stopped(self) -> bool:
|
||||
return self._stopped
|
||||
|
||||
def feed(self, data: bytes) -> None:
|
||||
if self._stopped:
|
||||
return
|
||||
if not data:
|
||||
return
|
||||
new_size = len(self._buffer) + len(data)
|
||||
if new_size > self._max_buffer_size:
|
||||
self._stopped = True
|
||||
raise BufferOverflowError(size=new_size, max_size=self._max_buffer_size)
|
||||
self._buffer.extend(data)
|
||||
|
||||
def decode_available(self) -> list[Frame]:
|
||||
"""Decode all complete frames currently in buffer."""
|
||||
out: list[Frame] = []
|
||||
if self._stopped:
|
||||
return out
|
||||
|
||||
while True:
|
||||
try:
|
||||
# Use memoryview to avoid full buffer copy on each iteration
|
||||
parsed = parse_frame(memoryview(self._buffer))
|
||||
except EventStreamParseError:
|
||||
self.stats.error_count += 1
|
||||
if self.stats.error_count >= self._max_errors:
|
||||
self._stopped = True
|
||||
raise
|
||||
|
||||
# Recovery: skip a byte and keep scanning.
|
||||
if self._buffer:
|
||||
del self._buffer[0]
|
||||
self.stats.bytes_skipped += 1
|
||||
else:
|
||||
break
|
||||
continue
|
||||
|
||||
if parsed is None:
|
||||
break
|
||||
|
||||
frame, consumed = parsed
|
||||
if consumed <= 0:
|
||||
break
|
||||
|
||||
out.append(frame)
|
||||
del self._buffer[:consumed]
|
||||
self.stats.frames_decoded += 1
|
||||
self.stats.error_count = 0
|
||||
|
||||
return out
|
||||
|
||||
|
||||
__all__ = [
|
||||
"DecoderStats",
|
||||
"EventStreamDecoder",
|
||||
]
|
||||
72
src/services/provider/adapters/kiro/parser/error.py
Normal file
72
src/services/provider/adapters/kiro/parser/error.py
Normal file
@@ -0,0 +1,72 @@
|
||||
"""AWS Event Stream parsing errors."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
class EventStreamParseError(Exception):
|
||||
"""Base error for AWS Event Stream decoding."""
|
||||
|
||||
|
||||
class IncompleteFrameError(EventStreamParseError):
|
||||
def __init__(self, *, needed: int, available: int) -> None:
|
||||
super().__init__(f"incomplete frame: needed={needed} available={available}")
|
||||
self.needed = needed
|
||||
self.available = available
|
||||
|
||||
|
||||
class MessageTooSmallError(EventStreamParseError):
|
||||
def __init__(self, *, length: int, min_length: int) -> None:
|
||||
super().__init__(f"message too small: length={length} min={min_length}")
|
||||
self.length = length
|
||||
self.min_length = min_length
|
||||
|
||||
|
||||
class MessageTooLargeError(EventStreamParseError):
|
||||
def __init__(self, *, length: int, max_length: int) -> None:
|
||||
super().__init__(f"message too large: length={length} max={max_length}")
|
||||
self.length = length
|
||||
self.max_length = max_length
|
||||
|
||||
|
||||
class PreludeCrcMismatchError(EventStreamParseError):
|
||||
def __init__(self, *, expected: int, actual: int) -> None:
|
||||
super().__init__(f"prelude crc mismatch: expected={expected} actual={actual}")
|
||||
self.expected = expected
|
||||
self.actual = actual
|
||||
|
||||
|
||||
class MessageCrcMismatchError(EventStreamParseError):
|
||||
def __init__(self, *, expected: int, actual: int) -> None:
|
||||
super().__init__(f"message crc mismatch: expected={expected} actual={actual}")
|
||||
self.expected = expected
|
||||
self.actual = actual
|
||||
|
||||
|
||||
class InvalidHeaderTypeError(EventStreamParseError):
|
||||
def __init__(self, type_id: int) -> None:
|
||||
super().__init__(f"invalid header type: {type_id}")
|
||||
self.type_id = type_id
|
||||
|
||||
|
||||
class HeaderParseError(EventStreamParseError):
|
||||
pass
|
||||
|
||||
|
||||
class BufferOverflowError(EventStreamParseError):
|
||||
def __init__(self, *, size: int, max_size: int) -> None:
|
||||
super().__init__(f"buffer overflow: size={size} max={max_size}")
|
||||
self.size = size
|
||||
self.max_size = max_size
|
||||
|
||||
|
||||
__all__ = [
|
||||
"BufferOverflowError",
|
||||
"EventStreamParseError",
|
||||
"HeaderParseError",
|
||||
"IncompleteFrameError",
|
||||
"InvalidHeaderTypeError",
|
||||
"MessageCrcMismatchError",
|
||||
"MessageTooLargeError",
|
||||
"MessageTooSmallError",
|
||||
"PreludeCrcMismatchError",
|
||||
]
|
||||
95
src/services/provider/adapters/kiro/parser/frame.py
Normal file
95
src/services/provider/adapters/kiro/parser/frame.py
Normal file
@@ -0,0 +1,95 @@
|
||||
"""AWS Event Stream message frame parsing."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from .crc import crc32
|
||||
from .error import (
|
||||
HeaderParseError,
|
||||
IncompleteFrameError,
|
||||
MessageCrcMismatchError,
|
||||
MessageTooLargeError,
|
||||
MessageTooSmallError,
|
||||
PreludeCrcMismatchError,
|
||||
)
|
||||
from .header import Headers, parse_headers
|
||||
|
||||
PRELUDE_SIZE = 12
|
||||
MIN_MESSAGE_SIZE = PRELUDE_SIZE + 4
|
||||
MAX_MESSAGE_SIZE = 16 * 1024 * 1024
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class Frame:
|
||||
headers: Headers
|
||||
payload: bytes
|
||||
|
||||
def message_type(self) -> str | None:
|
||||
return self.headers.message_type()
|
||||
|
||||
def event_type(self) -> str | None:
|
||||
return self.headers.event_type()
|
||||
|
||||
def payload_as_text(self) -> str:
|
||||
return self.payload.decode("utf-8", errors="replace")
|
||||
|
||||
|
||||
def parse_frame(buffer: bytes | memoryview) -> tuple[Frame, int] | None:
|
||||
"""Parse a single frame from the front of buffer.
|
||||
|
||||
Returns:
|
||||
(frame, consumed_bytes) if a full frame is available, otherwise None.
|
||||
|
||||
Raises:
|
||||
EventStreamParseError subclasses on validation errors.
|
||||
"""
|
||||
if len(buffer) < PRELUDE_SIZE:
|
||||
return None
|
||||
|
||||
total_length = int.from_bytes(buffer[0:4], "big", signed=False)
|
||||
header_length = int.from_bytes(buffer[4:8], "big", signed=False)
|
||||
prelude_crc = int.from_bytes(buffer[8:12], "big", signed=False)
|
||||
|
||||
if total_length < MIN_MESSAGE_SIZE:
|
||||
raise MessageTooSmallError(length=total_length, min_length=MIN_MESSAGE_SIZE)
|
||||
if total_length > MAX_MESSAGE_SIZE:
|
||||
raise MessageTooLargeError(length=total_length, max_length=MAX_MESSAGE_SIZE)
|
||||
|
||||
if len(buffer) < total_length:
|
||||
return None
|
||||
|
||||
actual_prelude_crc = crc32(buffer[0:8])
|
||||
if actual_prelude_crc != prelude_crc:
|
||||
raise PreludeCrcMismatchError(expected=prelude_crc, actual=actual_prelude_crc)
|
||||
|
||||
message_crc = int.from_bytes(buffer[total_length - 4 : total_length], "big", signed=False)
|
||||
actual_message_crc = crc32(buffer[0 : total_length - 4])
|
||||
if actual_message_crc != message_crc:
|
||||
raise MessageCrcMismatchError(expected=message_crc, actual=actual_message_crc)
|
||||
|
||||
headers_start = PRELUDE_SIZE
|
||||
headers_end = headers_start + header_length
|
||||
|
||||
if headers_end > total_length - 4:
|
||||
raise HeaderParseError("header length exceeds frame boundary")
|
||||
|
||||
headers = parse_headers(bytes(buffer[headers_start:headers_end]), header_length)
|
||||
|
||||
payload_start = headers_end
|
||||
payload_end = total_length - 4
|
||||
if payload_end < payload_start:
|
||||
raise IncompleteFrameError(needed=payload_start, available=payload_end)
|
||||
|
||||
payload = bytes(buffer[payload_start:payload_end])
|
||||
|
||||
return Frame(headers=headers, payload=payload), total_length
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Frame",
|
||||
"MAX_MESSAGE_SIZE",
|
||||
"MIN_MESSAGE_SIZE",
|
||||
"PRELUDE_SIZE",
|
||||
"parse_frame",
|
||||
]
|
||||
144
src/services/provider/adapters/kiro/parser/header.py
Normal file
144
src/services/provider/adapters/kiro/parser/header.py
Normal file
@@ -0,0 +1,144 @@
|
||||
"""AWS Event Stream header parsing."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from enum import IntEnum
|
||||
|
||||
from .error import HeaderParseError, IncompleteFrameError, InvalidHeaderTypeError
|
||||
|
||||
|
||||
class HeaderValueType(IntEnum):
|
||||
BOOL_TRUE = 0
|
||||
BOOL_FALSE = 1
|
||||
BYTE = 2
|
||||
SHORT = 3
|
||||
INTEGER = 4
|
||||
LONG = 5
|
||||
BYTE_ARRAY = 6
|
||||
STRING = 7
|
||||
TIMESTAMP = 8
|
||||
UUID = 9
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class Headers:
|
||||
values: dict[str, object]
|
||||
|
||||
def get(self, name: str) -> object | None:
|
||||
return self.values.get(name)
|
||||
|
||||
def get_string(self, name: str) -> str | None:
|
||||
v = self.values.get(name)
|
||||
return v if isinstance(v, str) else None
|
||||
|
||||
def message_type(self) -> str | None:
|
||||
return self.get_string(":message-type")
|
||||
|
||||
def event_type(self) -> str | None:
|
||||
return self.get_string(":event-type")
|
||||
|
||||
def exception_type(self) -> str | None:
|
||||
return self.get_string(":exception-type")
|
||||
|
||||
def error_code(self) -> str | None:
|
||||
return self.get_string(":error-code")
|
||||
|
||||
|
||||
def _ensure_bytes(data: bytes, offset: int, needed: int) -> None:
|
||||
available = len(data) - offset
|
||||
if available < needed:
|
||||
raise IncompleteFrameError(needed=needed, available=available)
|
||||
|
||||
|
||||
def parse_headers(data: bytes, header_length: int) -> Headers:
|
||||
if len(data) < header_length:
|
||||
raise IncompleteFrameError(needed=header_length, available=len(data))
|
||||
|
||||
values: dict[str, object] = {}
|
||||
offset = 0
|
||||
|
||||
while offset < header_length:
|
||||
_ensure_bytes(data, offset, 1)
|
||||
name_len = data[offset]
|
||||
offset += 1
|
||||
if name_len == 0:
|
||||
raise HeaderParseError("header name length cannot be 0")
|
||||
|
||||
_ensure_bytes(data, offset, name_len)
|
||||
name = data[offset : offset + name_len].decode("utf-8", errors="replace")
|
||||
offset += name_len
|
||||
|
||||
_ensure_bytes(data, offset, 1)
|
||||
type_id = data[offset]
|
||||
offset += 1
|
||||
try:
|
||||
value_type = HeaderValueType(type_id)
|
||||
except ValueError as e:
|
||||
raise InvalidHeaderTypeError(type_id) from e
|
||||
|
||||
if value_type == HeaderValueType.BOOL_TRUE:
|
||||
values[name] = True
|
||||
continue
|
||||
if value_type == HeaderValueType.BOOL_FALSE:
|
||||
values[name] = False
|
||||
continue
|
||||
|
||||
if value_type == HeaderValueType.BYTE:
|
||||
_ensure_bytes(data, offset, 1)
|
||||
values[name] = int.from_bytes(data[offset : offset + 1], "big", signed=True)
|
||||
offset += 1
|
||||
continue
|
||||
|
||||
if value_type == HeaderValueType.SHORT:
|
||||
_ensure_bytes(data, offset, 2)
|
||||
values[name] = int.from_bytes(data[offset : offset + 2], "big", signed=True)
|
||||
offset += 2
|
||||
continue
|
||||
|
||||
if value_type == HeaderValueType.INTEGER:
|
||||
_ensure_bytes(data, offset, 4)
|
||||
values[name] = int.from_bytes(data[offset : offset + 4], "big", signed=True)
|
||||
offset += 4
|
||||
continue
|
||||
|
||||
if value_type in (HeaderValueType.LONG, HeaderValueType.TIMESTAMP):
|
||||
_ensure_bytes(data, offset, 8)
|
||||
values[name] = int.from_bytes(data[offset : offset + 8], "big", signed=True)
|
||||
offset += 8
|
||||
continue
|
||||
|
||||
if value_type == HeaderValueType.BYTE_ARRAY:
|
||||
_ensure_bytes(data, offset, 2)
|
||||
length = int.from_bytes(data[offset : offset + 2], "big", signed=False)
|
||||
offset += 2
|
||||
_ensure_bytes(data, offset, length)
|
||||
values[name] = data[offset : offset + length]
|
||||
offset += length
|
||||
continue
|
||||
|
||||
if value_type == HeaderValueType.STRING:
|
||||
_ensure_bytes(data, offset, 2)
|
||||
length = int.from_bytes(data[offset : offset + 2], "big", signed=False)
|
||||
offset += 2
|
||||
_ensure_bytes(data, offset, length)
|
||||
values[name] = data[offset : offset + length].decode("utf-8", errors="replace")
|
||||
offset += length
|
||||
continue
|
||||
|
||||
if value_type == HeaderValueType.UUID:
|
||||
_ensure_bytes(data, offset, 16)
|
||||
values[name] = bytes(data[offset : offset + 16])
|
||||
offset += 16
|
||||
continue
|
||||
|
||||
raise HeaderParseError(f"unhandled header type: {value_type}")
|
||||
|
||||
return Headers(values=values)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"HeaderValueType",
|
||||
"Headers",
|
||||
"parse_headers",
|
||||
]
|
||||
166
src/services/provider/adapters/kiro/plugin.py
Normal file
166
src/services/provider/adapters/kiro/plugin.py
Normal file
@@ -0,0 +1,166 @@
|
||||
"""Kiro provider plugin — unified registration entry.
|
||||
|
||||
Kiro upstream looks like Claude CLI (Bearer token) from the outside, but uses a
|
||||
custom wire protocol:
|
||||
- Request: Claude Messages API -> Kiro generateAssistantResponse envelope
|
||||
- Response (stream): AWS Event Stream (binary) -> Claude SSE events
|
||||
|
||||
This plugin registers:
|
||||
- Envelope
|
||||
- Transport hook (dynamic region base_url)
|
||||
- Model fetcher (fixed model catalog — Kiro has no /v1/models endpoint)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
from urllib.parse import urlencode
|
||||
|
||||
from src.services.provider.adapters.kiro.constants import (
|
||||
DEFAULT_REGION,
|
||||
KIRO_GENERATE_ASSISTANT_PATH,
|
||||
)
|
||||
from src.services.provider.adapters.kiro.context import get_kiro_request_context
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fixed model catalog
|
||||
# ---------------------------------------------------------------------------
|
||||
# Kiro upstream has no /v1/models endpoint. We return a static list matching
|
||||
# the models accepted by map_model() in converter.py.
|
||||
_KIRO_MODELS: list[dict[str, Any]] = [
|
||||
{
|
||||
"id": "claude-sonnet-4.5",
|
||||
"object": "model",
|
||||
"owned_by": "anthropic",
|
||||
"display_name": "Claude Sonnet 4.5",
|
||||
},
|
||||
{
|
||||
"id": "claude-opus-4.5",
|
||||
"object": "model",
|
||||
"owned_by": "anthropic",
|
||||
"display_name": "Claude Opus 4.5",
|
||||
},
|
||||
{
|
||||
"id": "claude-opus-4.6",
|
||||
"object": "model",
|
||||
"owned_by": "anthropic",
|
||||
"display_name": "Claude Opus 4.6",
|
||||
},
|
||||
{
|
||||
"id": "claude-haiku-4.5",
|
||||
"object": "model",
|
||||
"owned_by": "anthropic",
|
||||
"display_name": "Claude Haiku 4.5",
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
async def fetch_models_kiro(
|
||||
ctx: Any,
|
||||
timeout_seconds: float, # noqa: ARG001
|
||||
) -> tuple[list[dict], list[str], bool, dict[str, Any] | None]:
|
||||
"""Return a fixed model catalog for Kiro.
|
||||
|
||||
Kiro upstream does not expose a ``/v1/models`` endpoint, so we skip the
|
||||
HTTP call entirely and return a hardcoded list.
|
||||
"""
|
||||
_ = ctx # not needed — no upstream call
|
||||
return list(_KIRO_MODELS), [], True, None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Transport hook
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def build_kiro_url(
|
||||
endpoint: Any,
|
||||
*,
|
||||
is_stream: bool,
|
||||
effective_query_params: dict[str, Any],
|
||||
) -> str:
|
||||
"""Build Kiro generateAssistantResponse URL.
|
||||
|
||||
Endpoint base_url may contain a `{region}` placeholder. The actual region is
|
||||
resolved from per-request context (set by the envelope).
|
||||
"""
|
||||
_ = is_stream
|
||||
|
||||
base = str(getattr(endpoint, "base_url", "") or "").rstrip("/")
|
||||
|
||||
ctx = get_kiro_request_context()
|
||||
region = (ctx.region if ctx else "") or DEFAULT_REGION
|
||||
if "{region}" in base:
|
||||
base = base.replace("{region}", region)
|
||||
|
||||
path = KIRO_GENERATE_ASSISTANT_PATH
|
||||
url = base if base.endswith(path) else f"{base}{path}"
|
||||
|
||||
if effective_query_params:
|
||||
query_string = urlencode(effective_query_params, doseq=True)
|
||||
if query_string:
|
||||
url = f"{url}?{query_string}"
|
||||
|
||||
return url
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Export builder
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_KIRO_SKIP_KEYS = frozenset(
|
||||
{
|
||||
"access_token",
|
||||
"expires_at",
|
||||
"updated_at",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def kiro_export_builder(
|
||||
auth_config: dict[str, Any],
|
||||
upstream_metadata: dict[str, Any] | None,
|
||||
) -> dict[str, Any]:
|
||||
"""Kiro 导出：保留 auth_method / refresh_token / machine_id / profile_arn 等，
|
||||
IdC 模式额外保留 client_id / client_secret / region。"""
|
||||
data = {
|
||||
k: v
|
||||
for k, v in auth_config.items()
|
||||
if k not in _KIRO_SKIP_KEYS and v is not None and v != ""
|
||||
}
|
||||
# email 可能仅在 upstream_metadata.kiro 中
|
||||
if not data.get("email"):
|
||||
kiro_meta = (upstream_metadata or {}).get("kiro") or {}
|
||||
if kiro_meta.get("email"):
|
||||
data["email"] = kiro_meta["email"]
|
||||
return data
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Registration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def register_all() -> None:
|
||||
"""Register all Kiro hooks into shared registries."""
|
||||
|
||||
from src.services.model.upstream_fetcher import UpstreamModelsFetcherRegistry
|
||||
from src.services.provider.adapters.kiro.envelope import kiro_envelope
|
||||
from src.services.provider.envelope import register_envelope
|
||||
from src.services.provider.export import register_export_builder
|
||||
from src.services.provider.transport import register_transport_hook
|
||||
|
||||
register_envelope("kiro", "claude:cli", kiro_envelope)
|
||||
register_envelope("kiro", "", kiro_envelope)
|
||||
|
||||
register_transport_hook("kiro", "claude:cli", build_kiro_url)
|
||||
|
||||
register_export_builder("kiro", kiro_export_builder)
|
||||
|
||||
UpstreamModelsFetcherRegistry.register(
|
||||
provider_types=["kiro"],
|
||||
fetcher=fetch_models_kiro,
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["build_kiro_url", "fetch_models_kiro", "kiro_export_builder", "register_all"]
|
||||
312
src/services/provider/adapters/kiro/token_manager.py
Normal file
312
src/services/provider/adapters/kiro/token_manager.py
Normal file
@@ -0,0 +1,312 @@
|
||||
"""Kiro token refresh helpers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import json
|
||||
import re
|
||||
import time
|
||||
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_kiro_ide_tag
|
||||
from src.services.provider.adapters.kiro.models.credentials import KiroAuthConfig
|
||||
|
||||
IDC_AMZ_USER_AGENT = (
|
||||
"aws-sdk-js/3.738.0 ua/2.1 os/other lang/js md/browser#unknown_unknown "
|
||||
"api/sso-oidc#3.738.0 m/E KiroIDE"
|
||||
)
|
||||
|
||||
_REGION_RE = re.compile(r"^[a-z]{2}-[a-z0-9-]+-\d+$")
|
||||
_HEX64_RE = re.compile(r"^[0-9a-fA-F]{64}$")
|
||||
_UUID_RE = re.compile(
|
||||
r"^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$"
|
||||
)
|
||||
|
||||
|
||||
def validate_refresh_token(refresh_token: str) -> None:
|
||||
token = str(refresh_token or "").strip()
|
||||
if not token:
|
||||
raise ValueError("missing refresh_token")
|
||||
|
||||
# kiro.rs: length < 100 or contains "..." is considered truncated.
|
||||
if len(token) < 100 or token.endswith("...") or "..." in token:
|
||||
raise ValueError(
|
||||
"refresh_token appears truncated; please export the full token from Kiro IDE"
|
||||
)
|
||||
|
||||
|
||||
def normalize_machine_id(machine_id: str) -> str | None:
|
||||
raw = str(machine_id or "").strip()
|
||||
if not raw:
|
||||
return None
|
||||
|
||||
if _HEX64_RE.fullmatch(raw):
|
||||
return raw.lower()
|
||||
|
||||
if _UUID_RE.fullmatch(raw):
|
||||
without = raw.replace("-", "").lower()
|
||||
return without + without
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def generate_machine_id(cfg: KiroAuthConfig) -> str:
|
||||
normalized = normalize_machine_id(cfg.machine_id or "")
|
||||
if normalized:
|
||||
return normalized
|
||||
|
||||
validate_refresh_token(cfg.refresh_token)
|
||||
seed = f"KotlinNativeAPI/{cfg.refresh_token}".encode("utf-8")
|
||||
return hashlib.sha256(seed).hexdigest()
|
||||
|
||||
|
||||
def is_token_expired(expires_at: int | None, *, skew_seconds: int = 120) -> bool:
|
||||
try:
|
||||
ts = int(expires_at or 0)
|
||||
except Exception:
|
||||
ts = 0
|
||||
if ts <= 0:
|
||||
return True
|
||||
return int(time.time()) >= ts - int(skew_seconds)
|
||||
|
||||
|
||||
def _resolve_region(cfg: KiroAuthConfig) -> str:
|
||||
region = str(cfg.region or "").strip()
|
||||
if region and _REGION_RE.fullmatch(region):
|
||||
return region
|
||||
# Keep best-effort fallback; actual host parsing happens in transport hook.
|
||||
from src.services.provider.adapters.kiro.constants import DEFAULT_REGION
|
||||
|
||||
return region or DEFAULT_REGION
|
||||
|
||||
|
||||
def _try_extract_email_from_jwt(token: str) -> str | None:
|
||||
"""尝试从 JWT access_token 中提取 email。
|
||||
|
||||
Kiro Social / IdC 返回的 accessToken 可能是 JWT 格式，
|
||||
payload 中可能包含 email 字段。仅做 base64 解码，不验证签名。
|
||||
失败时静默返回 None。
|
||||
"""
|
||||
try:
|
||||
parts = token.split(".")
|
||||
if len(parts) != 3:
|
||||
return None
|
||||
# base64url decode the payload (second segment)
|
||||
payload_b64 = parts[1]
|
||||
# Add padding
|
||||
padding = 4 - len(payload_b64) % 4
|
||||
if padding != 4:
|
||||
payload_b64 += "=" * padding
|
||||
payload_bytes = base64.urlsafe_b64decode(payload_b64)
|
||||
claims = json.loads(payload_bytes)
|
||||
# Try common email claim keys
|
||||
for key in ("email", "Email", "mail", "upn"):
|
||||
val = claims.get(key)
|
||||
if isinstance(val, str) and "@" in val:
|
||||
return val.strip()
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
async def refresh_social_token(
|
||||
cfg: KiroAuthConfig,
|
||||
*,
|
||||
proxy_config: dict[str, Any] | None,
|
||||
timeout_seconds: float = 30.0,
|
||||
) -> tuple[str, KiroAuthConfig]:
|
||||
"""Refresh access token via Kiro Social refresh endpoint."""
|
||||
validate_refresh_token(cfg.refresh_token)
|
||||
|
||||
region = _resolve_region(cfg)
|
||||
url = f"https://prod.{region}.auth.desktop.kiro.dev/refreshToken"
|
||||
host = f"prod.{region}.auth.desktop.kiro.dev"
|
||||
|
||||
machine_id = generate_machine_id(cfg)
|
||||
kiro_version = (cfg.kiro_version or "").strip() or "0.8.0"
|
||||
ua = build_kiro_ide_tag(kiro_version=kiro_version, machine_id=machine_id)
|
||||
|
||||
body = {"refreshToken": cfg.refresh_token}
|
||||
headers = {
|
||||
"User-Agent": ua,
|
||||
"Host": host,
|
||||
"Accept": "application/json, text/plain, */*",
|
||||
"Content-Type": "application/json",
|
||||
"Connection": "close",
|
||||
"Accept-Encoding": "gzip, compress, deflate, br",
|
||||
}
|
||||
|
||||
client = await HTTPClientPool.get_proxy_client(proxy_config=proxy_config)
|
||||
resp = await client.post(
|
||||
url,
|
||||
headers=headers,
|
||||
json=body,
|
||||
timeout=httpx.Timeout(timeout_seconds),
|
||||
)
|
||||
|
||||
if resp.status_code < 200 or resp.status_code >= 300:
|
||||
logger.debug(
|
||||
"kiro social refresh error: HTTP {} | {}",
|
||||
resp.status_code,
|
||||
(resp.text or "").strip()[:200],
|
||||
)
|
||||
raise RuntimeError(f"kiro social refresh failed: HTTP {resp.status_code}")
|
||||
|
||||
data: dict[str, Any]
|
||||
try:
|
||||
data = resp.json()
|
||||
except Exception as e:
|
||||
raise RuntimeError("kiro social refresh: invalid json response") from e
|
||||
|
||||
access_token = str(data.get("accessToken") or "").strip()
|
||||
if not access_token:
|
||||
raise RuntimeError("kiro social refresh returned empty accessToken")
|
||||
|
||||
new_cfg = KiroAuthConfig.from_dict(cfg.to_dict())
|
||||
|
||||
# refreshToken/profileArn may rotate
|
||||
rt = data.get("refreshToken")
|
||||
if isinstance(rt, str) and rt.strip():
|
||||
new_cfg.refresh_token = rt.strip()
|
||||
|
||||
profile_arn = data.get("profileArn")
|
||||
if isinstance(profile_arn, str) and profile_arn.strip():
|
||||
new_cfg.profile_arn = profile_arn.strip()
|
||||
|
||||
expires_in = data.get("expiresIn")
|
||||
try:
|
||||
if expires_in is not None:
|
||||
new_cfg.expires_at = int(time.time()) + int(expires_in)
|
||||
except Exception:
|
||||
new_cfg.expires_at = int(time.time()) + 3600
|
||||
|
||||
# Persist computed machine_id if user didn't provide one.
|
||||
if not (cfg.machine_id or "").strip():
|
||||
new_cfg.machine_id = machine_id
|
||||
|
||||
# 尝试从 accessToken 中提取 email（如果尚未设置）
|
||||
if not (new_cfg.email or "").strip():
|
||||
extracted_email = _try_extract_email_from_jwt(access_token)
|
||||
if extracted_email:
|
||||
new_cfg.email = extracted_email
|
||||
logger.debug("kiro social: extracted email from accessToken: {}", extracted_email)
|
||||
|
||||
# 缓存 access_token
|
||||
new_cfg.access_token = access_token
|
||||
|
||||
return access_token, new_cfg
|
||||
|
||||
|
||||
async def refresh_idc_token(
|
||||
cfg: KiroAuthConfig,
|
||||
*,
|
||||
proxy_config: dict[str, Any] | None,
|
||||
timeout_seconds: float = 30.0,
|
||||
) -> tuple[str, KiroAuthConfig]:
|
||||
"""Refresh access token via AWS SSO OIDC endpoint (IdC)."""
|
||||
validate_refresh_token(cfg.refresh_token)
|
||||
|
||||
if not (cfg.client_id or "").strip() or not (cfg.client_secret or "").strip():
|
||||
raise ValueError("idc refresh requires client_id and client_secret")
|
||||
|
||||
region = _resolve_region(cfg)
|
||||
url = f"https://oidc.{region}.amazonaws.com/token"
|
||||
host = f"oidc.{region}.amazonaws.com"
|
||||
|
||||
body = {
|
||||
"clientId": cfg.client_id,
|
||||
"clientSecret": cfg.client_secret,
|
||||
"refreshToken": cfg.refresh_token,
|
||||
"grantType": "refresh_token",
|
||||
}
|
||||
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Host": host,
|
||||
"x-amz-user-agent": IDC_AMZ_USER_AGENT,
|
||||
"User-Agent": "node",
|
||||
"Accept": "*/*",
|
||||
}
|
||||
|
||||
client = await HTTPClientPool.get_proxy_client(proxy_config=proxy_config)
|
||||
resp = await client.post(
|
||||
url,
|
||||
headers=headers,
|
||||
json=body,
|
||||
timeout=httpx.Timeout(timeout_seconds),
|
||||
)
|
||||
|
||||
if resp.status_code < 200 or resp.status_code >= 300:
|
||||
logger.debug(
|
||||
"kiro idc refresh error: HTTP {} | {}",
|
||||
resp.status_code,
|
||||
(resp.text or "").strip()[:200],
|
||||
)
|
||||
raise RuntimeError(f"kiro idc refresh failed: HTTP {resp.status_code}")
|
||||
|
||||
data: dict[str, Any]
|
||||
try:
|
||||
data = resp.json()
|
||||
except Exception as e:
|
||||
raise RuntimeError("kiro idc refresh: invalid json response") from e
|
||||
|
||||
access_token = str(data.get("accessToken") or "").strip()
|
||||
if not access_token:
|
||||
raise RuntimeError("kiro idc refresh returned empty accessToken")
|
||||
|
||||
new_cfg = KiroAuthConfig.from_dict(cfg.to_dict())
|
||||
|
||||
rt = data.get("refreshToken")
|
||||
if isinstance(rt, str) and rt.strip():
|
||||
new_cfg.refresh_token = rt.strip()
|
||||
|
||||
expires_in = data.get("expiresIn")
|
||||
try:
|
||||
if expires_in is not None:
|
||||
new_cfg.expires_at = int(time.time()) + int(expires_in)
|
||||
except Exception:
|
||||
new_cfg.expires_at = int(time.time()) + 3600
|
||||
|
||||
# Persist computed machine_id if user didn't provide one.
|
||||
if not (cfg.machine_id or "").strip():
|
||||
new_cfg.machine_id = generate_machine_id(cfg)
|
||||
|
||||
# 尝试从 accessToken 中提取 email（如果尚未设置）
|
||||
if not (new_cfg.email or "").strip():
|
||||
extracted_email = _try_extract_email_from_jwt(access_token)
|
||||
if extracted_email:
|
||||
new_cfg.email = extracted_email
|
||||
logger.debug("kiro idc: extracted email from accessToken: {}", extracted_email)
|
||||
|
||||
# 缓存 access_token
|
||||
new_cfg.access_token = access_token
|
||||
|
||||
return access_token, new_cfg
|
||||
|
||||
|
||||
async def refresh_access_token(
|
||||
cfg: KiroAuthConfig,
|
||||
*,
|
||||
proxy_config: dict[str, Any] | None,
|
||||
) -> tuple[str, KiroAuthConfig]:
|
||||
method = (cfg.auth_method or "social").strip().lower()
|
||||
if method == "idc":
|
||||
return await refresh_idc_token(cfg, proxy_config=proxy_config)
|
||||
return await refresh_social_token(cfg, proxy_config=proxy_config)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"IDC_AMZ_USER_AGENT",
|
||||
"generate_machine_id",
|
||||
"is_token_expired",
|
||||
"normalize_machine_id",
|
||||
"refresh_access_token",
|
||||
"refresh_idc_token",
|
||||
"refresh_social_token",
|
||||
"validate_refresh_token",
|
||||
]
|
||||
183
src/services/provider/adapters/kiro/usage.py
Normal file
183
src/services/provider/adapters/kiro/usage.py
Normal file
@@ -0,0 +1,183 @@
|
||||
"""Kiro usage/quota fetching utilities."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
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.token_manager import (
|
||||
generate_machine_id,
|
||||
is_token_expired,
|
||||
refresh_access_token,
|
||||
)
|
||||
|
||||
|
||||
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")
|
||||
|
||||
# 构建请求
|
||||
from src.services.provider.adapters.kiro.constants import DEFAULT_REGION
|
||||
|
||||
region = (updated_cfg.region if updated_cfg else cfg.region) or DEFAULT_REGION
|
||||
host = f"q.{region}.amazonaws.com"
|
||||
machine_id = generate_machine_id(updated_cfg or cfg)
|
||||
kiro_version = (updated_cfg.kiro_version if updated_cfg else cfg.kiro_version) or "0.8.0"
|
||||
|
||||
# 构建 URL（添加 isEmailRequired=true 获取邮箱）
|
||||
url = f"https://{host}/getUsageLimits?origin=AI_EDITOR&resourceType=AGENTIC_REQUEST&isEmailRequired=true"
|
||||
|
||||
profile_arn = updated_cfg.profile_arn if updated_cfg else cfg.profile_arn
|
||||
if profile_arn:
|
||||
from urllib.parse import quote
|
||||
|
||||
url += f"&profileArn={quote(profile_arn, safe='')}"
|
||||
|
||||
# 构建 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",
|
||||
}
|
||||
|
||||
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:
|
||||
error_msg = {
|
||||
401: "认证失败，Token 无效或已过期",
|
||||
403: "权限不足，无法获取使用额度",
|
||||
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 or "").strip()[:200],
|
||||
)
|
||||
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__ = [
|
||||
"fetch_kiro_usage_limits",
|
||||
"parse_kiro_usage_response",
|
||||
]
|
||||
@@ -120,9 +120,11 @@ def ensure_providers_bootstrapped() -> None:
|
||||
register_all as _reg_antigravity,
|
||||
)
|
||||
from src.services.provider.adapters.codex.plugin import register_all as _reg_codex
|
||||
from src.services.provider.adapters.kiro.plugin import register_all as _reg_kiro
|
||||
|
||||
_reg_antigravity()
|
||||
_reg_codex()
|
||||
_reg_kiro()
|
||||
|
||||
|
||||
__all__ = [
|
||||
|
||||
52
src/services/provider/export.py
Normal file
52
src/services/provider/export.py
Normal file
@@ -0,0 +1,52 @@
|
||||
"""OAuth Key 导出：provider-specific export builders.
|
||||
|
||||
每个 Provider adapter 在 ``register_all()`` 时注册自己的 ``build_export_data``，
|
||||
导出端点通过 ``build_export_data()`` 分发。
|
||||
|
||||
未注册的 provider_type 使用默认实现（strip null + 临时字段）。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Callable
|
||||
|
||||
# 所有 provider 共用的临时/无用字段
|
||||
_DEFAULT_SKIP_KEYS = frozenset(
|
||||
{
|
||||
"access_token",
|
||||
"expires_at",
|
||||
"updated_at",
|
||||
"token_type",
|
||||
"scope",
|
||||
}
|
||||
)
|
||||
|
||||
ExportBuilder = Callable[[dict[str, Any], dict[str, Any] | None], dict[str, Any]]
|
||||
|
||||
_BUILDERS: dict[str, ExportBuilder] = {}
|
||||
|
||||
|
||||
def register_export_builder(provider_type: str, builder: ExportBuilder) -> None:
|
||||
_BUILDERS[provider_type] = builder
|
||||
|
||||
|
||||
def _default_builder(
|
||||
auth_config: dict[str, Any],
|
||||
upstream_metadata: dict[str, Any] | None,
|
||||
) -> dict[str, Any]:
|
||||
"""默认导出：去掉 null、空字符串、临时字段。"""
|
||||
return {
|
||||
k: v
|
||||
for k, v in auth_config.items()
|
||||
if k not in _DEFAULT_SKIP_KEYS and v is not None and v != ""
|
||||
}
|
||||
|
||||
|
||||
def build_export_data(
|
||||
provider_type: str,
|
||||
auth_config: dict[str, Any],
|
||||
upstream_metadata: dict[str, Any] | None,
|
||||
) -> dict[str, Any]:
|
||||
"""调用 provider-specific builder 构建导出数据。"""
|
||||
builder = _BUILDERS.get(provider_type, _default_builder)
|
||||
return builder(auth_config, upstream_metadata)
|
||||
@@ -1,156 +0,0 @@
|
||||
"""
|
||||
上游元数据采集器（MetadataCollector）
|
||||
|
||||
可扩展注册表模式：
|
||||
- 每个 Provider 类型可注册一个 MetadataCollector
|
||||
- 从响应头解析有价值的元数据（额度、限流等）
|
||||
- 解析结果存入 ProviderAPIKey.upstream_metadata
|
||||
|
||||
扩展方式：
|
||||
1. 创建新文件实现 MetadataCollector
|
||||
2. 在本文件底部注册
|
||||
"""
|
||||
|
||||
import time
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, ClassVar
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.logger import logger
|
||||
|
||||
# 节流：每个 key_id 至少间隔 _THROTTLE_SECONDS 秒才写入一次
|
||||
_THROTTLE_SECONDS = 30
|
||||
_last_write_ts: dict[str, float] = {}
|
||||
|
||||
|
||||
class MetadataCollector(ABC):
|
||||
"""元数据采集器基类"""
|
||||
|
||||
# 支持的 provider_type 列表（小写）
|
||||
PROVIDER_TYPES: ClassVar[list[str]] = []
|
||||
|
||||
@abstractmethod
|
||||
def parse_headers(self, headers: dict[str, str]) -> dict[str, Any] | None:
|
||||
"""解析响应头，返回结构化元数据。返回 None 表示无可用数据。"""
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class MetadataCollectorRegistry:
|
||||
"""元数据采集器注册表"""
|
||||
|
||||
_collectors: ClassVar[list[MetadataCollector]] = []
|
||||
_type_index: ClassVar[dict[str, MetadataCollector]] = {}
|
||||
|
||||
@classmethod
|
||||
def register(cls, collector: MetadataCollector) -> None:
|
||||
cls._collectors.append(collector)
|
||||
for pt in collector.PROVIDER_TYPES:
|
||||
cls._type_index[pt.lower()] = collector
|
||||
logger.debug(
|
||||
"[MetadataCollectorRegistry] 注册: {} -> {}",
|
||||
collector.__class__.__name__,
|
||||
collector.PROVIDER_TYPES,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def collect(cls, provider_type: str, headers: dict[str, str]) -> dict[str, Any] | None:
|
||||
"""根据 provider_type 查找采集器并解析响应头"""
|
||||
collector = cls._type_index.get(provider_type.lower())
|
||||
if collector is None:
|
||||
return None
|
||||
try:
|
||||
return collector.parse_headers(headers)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"[MetadataCollectorRegistry] {} 解析失败", collector.__class__.__name__
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
_initialized = False
|
||||
|
||||
|
||||
def _ensure_collectors_registered() -> None:
|
||||
"""惰性注册所有采集器（首次调用时执行，避免循环导入）"""
|
||||
global _initialized
|
||||
if _initialized:
|
||||
return
|
||||
_initialized = True
|
||||
|
||||
# 延迟导入，避免模块加载时的循环依赖
|
||||
from src.services.provider.adapters.codex.metadata_collector import CodexMetadataCollector
|
||||
|
||||
MetadataCollectorRegistry.register(CodexMetadataCollector())
|
||||
|
||||
|
||||
def ensure_collectors_registered() -> None:
|
||||
"""Ensure metadata collectors are registered (idempotent)."""
|
||||
_ensure_collectors_registered()
|
||||
|
||||
|
||||
def collect_and_save_upstream_metadata(
|
||||
db: Session,
|
||||
*,
|
||||
provider_type: str,
|
||||
key_id: str,
|
||||
response_headers: dict[str, str],
|
||||
request_id: str,
|
||||
) -> None:
|
||||
"""采集上游元数据并更新 ProviderAPIKey.upstream_metadata（带节流）
|
||||
|
||||
每个 key_id 至少间隔 _THROTTLE_SECONDS 秒才执行一次数据库写入，
|
||||
避免高并发时频繁更新同一行。
|
||||
|
||||
Args:
|
||||
db: 数据库 Session
|
||||
provider_type: Provider 类型（如 "codex"）
|
||||
key_id: ProviderAPIKey.id
|
||||
response_headers: 上游响应头
|
||||
request_id: 请求 ID（用于日志）
|
||||
"""
|
||||
if not provider_type or not key_id or not response_headers:
|
||||
return
|
||||
|
||||
# 确保采集器已注册
|
||||
_ensure_collectors_registered()
|
||||
|
||||
# 节流检查
|
||||
now = time.monotonic()
|
||||
last_ts = _last_write_ts.get(key_id, 0.0)
|
||||
if now - last_ts < _THROTTLE_SECONDS:
|
||||
return
|
||||
|
||||
try:
|
||||
metadata = MetadataCollectorRegistry.collect(provider_type, response_headers)
|
||||
if metadata is None:
|
||||
return
|
||||
|
||||
from src.models.database import ProviderAPIKey
|
||||
|
||||
key = db.query(ProviderAPIKey).filter(ProviderAPIKey.id == key_id).first()
|
||||
if key is None:
|
||||
return
|
||||
|
||||
key.upstream_metadata = metadata
|
||||
db.commit()
|
||||
_last_write_ts[key_id] = now
|
||||
logger.debug(
|
||||
"[{}] 已更新 ProviderAPIKey({}) upstream_metadata",
|
||||
request_id,
|
||||
key_id,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("[{}] 采集上游元数据失败", request_id)
|
||||
try:
|
||||
db.rollback()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MetadataCollector",
|
||||
"MetadataCollectorRegistry",
|
||||
"collect_and_save_upstream_metadata",
|
||||
"ensure_collectors_registered",
|
||||
]
|
||||
@@ -78,12 +78,18 @@ def get_upstream_stream_policy(
|
||||
and parsed == UpstreamStreamPolicy.FORCE_NON_STREAM
|
||||
):
|
||||
return UpstreamStreamPolicy.FORCE_STREAM
|
||||
if pt == ProviderType.KIRO and parsed == UpstreamStreamPolicy.FORCE_NON_STREAM:
|
||||
return UpstreamStreamPolicy.FORCE_STREAM
|
||||
return parsed
|
||||
|
||||
# Safe-by-default: Codex Responses OAuth behaves like SSE-only.
|
||||
if pt == ProviderType.CODEX and sig == "openai:cli":
|
||||
return UpstreamStreamPolicy.FORCE_STREAM
|
||||
|
||||
# Kiro upstream streams binary AWS Event Stream; treat as stream-only.
|
||||
if pt == ProviderType.KIRO:
|
||||
return UpstreamStreamPolicy.FORCE_STREAM
|
||||
|
||||
return UpstreamStreamPolicy.AUTO
|
||||
|
||||
|
||||
|
||||
@@ -101,8 +101,18 @@ def get_antigravity_base_url() -> str | None:
|
||||
return get_selected_base_url()
|
||||
|
||||
|
||||
def _get_provider_type(endpoint: Any, key: "ProviderAPIKey" | None = None) -> str | None:
|
||||
"""尽力获取 Provider.provider_type（用于 Antigravity 等 Provider 特判）。"""
|
||||
def _get_provider_type(
|
||||
endpoint: Any,
|
||||
key: "ProviderAPIKey" | None = None,
|
||||
decrypted_auth_config: dict[str, Any] | None = None,
|
||||
) -> str | None:
|
||||
"""尽力获取 Provider.provider_type（用于 Antigravity 等 Provider 特判）。
|
||||
|
||||
优先级:
|
||||
1. endpoint.provider.provider_type
|
||||
2. key.provider.provider_type
|
||||
3. decrypted_auth_config["provider_type"]（OAuth 导入的凭证）
|
||||
"""
|
||||
try:
|
||||
provider = getattr(endpoint, "provider", None)
|
||||
if provider is not None:
|
||||
@@ -122,6 +132,12 @@ def _get_provider_type(endpoint: Any, key: "ProviderAPIKey" | None = None) -> st
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Fallback: OAuth 导入的凭证可能包含 provider_type（如 Kiro）
|
||||
if decrypted_auth_config:
|
||||
pt = decrypted_auth_config.get("provider_type")
|
||||
if isinstance(pt, str) and pt.strip():
|
||||
return pt.strip().lower()
|
||||
|
||||
return None
|
||||
|
||||
|
||||
@@ -179,7 +195,7 @@ def build_provider_url(
|
||||
# endpoint_sig 为空时保持为空（更安全：默认路径回退到 "/"，避免误判为 claude:chat）
|
||||
endpoint_sig = normalize_endpoint_signature(endpoint_sig) if endpoint_sig else ""
|
||||
|
||||
provider_type = _get_provider_type(endpoint, key)
|
||||
provider_type = _get_provider_type(endpoint, key, decrypted_auth_config)
|
||||
|
||||
# 合并查询参数（部分逻辑需要先拿到 query_params）
|
||||
effective_query_params = dict(query_params) if query_params else {}
|
||||
|
||||
@@ -268,6 +268,30 @@ def resolve_ops_proxy(
|
||||
return None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Key 级别代理优先解析
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def resolve_effective_proxy(
|
||||
provider_proxy: dict[str, Any] | None,
|
||||
key_proxy: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any] | None:
|
||||
"""
|
||||
解析有效的代理配置，Key 级别代理优先于 Provider 级别代理。
|
||||
|
||||
Args:
|
||||
provider_proxy: Provider 级别代理配置
|
||||
key_proxy: Key 级别代理配置（可选），非 None 且 enabled 时覆盖 Provider 级别
|
||||
|
||||
Returns:
|
||||
有效的代理配置字典，或 None（无代理）
|
||||
"""
|
||||
if key_proxy and key_proxy.get("enabled", True):
|
||||
return key_proxy
|
||||
return provider_proxy
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 代理 URL 构建
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
Reference in New Issue
Block a user