from __future__ import annotations import asyncio import gzip import json import time import uuid from dataclasses import dataclass, field from typing import Any from fastapi import HTTPException, Request from sqlalchemy.orm import Session from starlette.requests import ClientDisconnect from src.config.settings import config from src.core.api_format.headers import get_header_value from src.core.http_compression import is_gzip_content_encoding, normalize_content_encoding from src.core.logger import logger from src.models.database import ApiKey, ManagementToken, User from src.utils.perf import PerfRecorder from src.utils.request_utils import ( get_request_identity_metadata, update_request_state, ) def _snapshot_optional_str(value: Any) -> str | None: return value.strip() if isinstance(value, str) and value.strip() else None def _snapshot_optional_float(value: Any) -> float | None: if value is None: return None if isinstance(value, bool): return None if isinstance(value, (int, float)): return float(value) if isinstance(value, str): raw = value.strip() if not raw: return None try: return float(raw) except ValueError: return None return None @dataclass(frozen=True) class _ContextBuildSnapshot: request_id: str start_time: float request_method: str request_path: str client_ip: str user_agent: str original_headers: dict[str, str] request_content_type: str | None query_params: dict[str, str] path_params: dict[str, Any] prefetched_balance_remaining: float | None gateway_execution_path: str | None rate_limit_scope: str | None tx_committed_by_route: bool client_content_encoding: str | None client_accept_encoding: str | None perf_metrics: dict[str, Any] | None @classmethod def from_request(cls, request: Request) -> _ContextBuildSnapshot: request_state = getattr(request, "state", None) original_headers = dict(request.headers) identity = get_request_identity_metadata(request) request_id = identity.request_id or str(uuid.uuid4())[:8] accept_encoding = get_header_value(original_headers, "accept-encoding") if isinstance(accept_encoding, str): accept_encoding = accept_encoding.strip() or None perf_metrics = getattr(request.state, "perf_metrics", None) perf_payload = perf_metrics if isinstance(perf_metrics, dict) and perf_metrics else None return cls( request_id=request_id, start_time=time.time(), request_method=request.method, request_path=request.url.path, client_ip=identity.client_ip, user_agent=identity.user_agent, original_headers=original_headers, request_content_type=get_header_value(original_headers, "content-type"), query_params=dict(request.query_params), path_params=dict(getattr(request, "path_params", {}) or {}), prefetched_balance_remaining=_snapshot_optional_float( getattr(request_state, "prefetched_balance_remaining", None) ), gateway_execution_path=_snapshot_optional_str( getattr(request_state, "gateway_execution_path", None) ), rate_limit_scope=_snapshot_optional_str( getattr(request_state, "rate_limit_scope", None) ), tx_committed_by_route=getattr(request_state, "tx_committed_by_route", False) is True, client_content_encoding=normalize_content_encoding( get_header_value(original_headers, "content-encoding") ), client_accept_encoding=accept_encoding, perf_metrics=perf_payload, ) def _apply_context_state_markers( request: Request, *, request_id: str, user: User | None, api_key: ApiKey | None, ) -> None: update_request_state(request, request_id=request_id) if user: update_request_state(request, user_id=user.id) if api_key: update_request_state(request, api_key_id=api_key.id) @dataclass class ApiRequestContext: """统一的API请求上下文,贯穿Pipeline与格式适配器。""" request: Request db: Session user: User | None api_key: ApiKey | None request_id: str start_time: float request_method: str request_path: str client_ip: str user_agent: str original_headers: dict[str, str] query_params: dict[str, str] request_content_type: str | None = None perf_metrics: dict[str, Any] | None = None raw_body: bytes | None = None json_body: dict[str, Any] | None = None balance_remaining: float | None = None prefetched_balance_remaining: float | None = None mode: str = "standard" # standard / proxy api_format_hint: str | None = None # URL 路径参数(如 Gemini API 的 /v1beta/models/{model}:generateContent) path_params: dict[str, Any] = field(default_factory=dict) # Management Token(用于管理 API 认证) management_token: ManagementToken | None = None # 供适配器扩展的状态存储 extra: dict[str, Any] = field(default_factory=dict) audit_metadata: dict[str, Any] = field(default_factory=dict) # 高频轮询端点日志抑制标志 quiet_logging: bool = False gateway_execution_path: str | None = None rate_limit_scope: str | None = None tx_committed_by_route: bool = False client_content_encoding: str | None = None client_accept_encoding: str | None = None def _get_perf_metrics(self) -> dict[str, Any] | None: if isinstance(self.perf_metrics, dict): return self.perf_metrics perf_metrics = getattr(self.request.state, "perf_metrics", None) if isinstance(perf_metrics, dict): self.perf_metrics = perf_metrics return perf_metrics return None def sync_runtime_state_from_request(self) -> None: request_state = getattr(self.request, "state", None) self.prefetched_balance_remaining = _snapshot_optional_float( getattr(request_state, "prefetched_balance_remaining", self.prefetched_balance_remaining) ) self.gateway_execution_path = _snapshot_optional_str( getattr(request_state, "gateway_execution_path", self.gateway_execution_path) ) self.rate_limit_scope = _snapshot_optional_str( getattr(request_state, "rate_limit_scope", self.rate_limit_scope) ) self.tx_committed_by_route = ( getattr(request_state, "tx_committed_by_route", self.tx_committed_by_route) is True ) async def ensure_raw_body_async(self) -> bytes: """按需读取原始请求体,避免所有请求都在 Pipeline 阶段预读。""" if self.raw_body is not None: return self.raw_body perf_metrics = self._get_perf_metrics() perf_sampled = isinstance(perf_metrics, dict) and bool(perf_metrics) body_start = PerfRecorder.start(force=perf_sampled) body_size = 0 try: self.raw_body = await asyncio.wait_for( self.request.body(), timeout=config.request_body_timeout ) body_size = len(self.raw_body or b"") except TimeoutError as exc: timeout_sec = int(config.request_body_timeout) logger.error("读取请求体超时({}s),可能客户端未发送完整请求体", timeout_sec) raise HTTPException( status_code=408, detail=f"Request timeout: body not received within {timeout_sec} seconds", ) from exc except ClientDisconnect: logger.warning( "[Context] 客户端在读取请求体期间断开连接: {} {}", self.request_method, self.request_path, ) raise HTTPException( status_code=499, detail="Client closed request", ) finally: body_duration = PerfRecorder.stop( body_start, "pipeline_body_read", labels={"mode": self.mode}, log_hint=f"size={body_size}", ) if isinstance(perf_metrics, dict): pipeline_metrics = perf_metrics.setdefault("pipeline", {}) pipeline_metrics["body_read_ms"] = int((body_duration or 0) * 1000) pipeline_metrics["body_bytes"] = int(body_size) return self.raw_body or b"" async def ensure_json_body_async(self) -> dict[str, Any]: """异步懒加载 JSON 请求体。""" await self.ensure_raw_body_async() return self.ensure_json_body() def ensure_json_body(self) -> dict[str, Any]: """确保请求体已解析为JSON并返回。""" if self.json_body is not None: return self.json_body if not self.raw_body: raise HTTPException(status_code=400, detail="请求体不能为空") perf_metrics = self._get_perf_metrics() perf_sampled = isinstance(perf_metrics, dict) and bool(perf_metrics) parse_start = PerfRecorder.start(force=perf_sampled) def _record_parse_duration(duration: float | None) -> None: if duration is None: return if not isinstance(perf_metrics, dict): return perf_metrics.setdefault("pipeline", {})["json_parse_ms"] = int(duration * 1000) body_to_parse = self.raw_body content_encoding = self.client_content_encoding or normalize_content_encoding( get_header_value(self.original_headers, "content-encoding") ) if is_gzip_content_encoding(content_encoding): try: body_to_parse = gzip.decompress(body_to_parse) except OSError as exc: parse_duration = PerfRecorder.stop( parse_start, "pipeline_json_parse", labels={"mode": self.mode}, ) _record_parse_duration(parse_duration) logger.warning("gzip 请求体解压失败: {}", exc) raise HTTPException(status_code=400, detail="gzip 请求体解压失败") from exc try: self.json_body = json.loads(body_to_parse.decode("utf-8")) parse_duration = PerfRecorder.stop( parse_start, "pipeline_json_parse", labels={"mode": self.mode}, ) _record_parse_duration(parse_duration) except json.JSONDecodeError as exc: parse_duration = PerfRecorder.stop( parse_start, "pipeline_json_parse", labels={"mode": self.mode}, ) _record_parse_duration(parse_duration) logger.warning(f"解析JSON失败: {exc}") raise HTTPException(status_code=400, detail="请求体必须是合法的JSON") from exc return self.json_body def add_audit_metadata(self, **values: Any) -> None: """向审计日志附加字段(会自动过滤 None)。""" for key, value in values.items(): if value is not None: self.audit_metadata[key] = value def extend_audit_metadata(self, data: dict[str, Any]) -> None: """批量附加审计字段。""" for key, value in data.items(): if value is not None: self.audit_metadata[key] = value @classmethod def build( cls, request: Request, db: Session, user: User | None, api_key: ApiKey | None, raw_body: bytes | None = None, mode: str = "standard", api_format_hint: str | None = None, path_params: dict[str, Any] | None = None, ) -> ApiRequestContext: """创建上下文实例并提前读取必要的元数据。""" snapshot = _ContextBuildSnapshot.from_request(request) _apply_context_state_markers( request, request_id=snapshot.request_id, user=user, api_key=api_key, ) context = cls( request=request, db=db, user=user, api_key=api_key, request_id=snapshot.request_id, start_time=snapshot.start_time, request_method=snapshot.request_method, request_path=snapshot.request_path, client_ip=snapshot.client_ip, user_agent=snapshot.user_agent, original_headers=snapshot.original_headers, request_content_type=snapshot.request_content_type, query_params=snapshot.query_params, perf_metrics=snapshot.perf_metrics, raw_body=raw_body, prefetched_balance_remaining=snapshot.prefetched_balance_remaining, mode=mode, api_format_hint=api_format_hint, path_params=dict(path_params or snapshot.path_params), gateway_execution_path=snapshot.gateway_execution_path, rate_limit_scope=snapshot.rate_limit_scope, tx_committed_by_route=snapshot.tx_committed_by_route, client_content_encoding=snapshot.client_content_encoding, client_accept_encoding=snapshot.client_accept_encoding, ) if snapshot.perf_metrics is not None: context.extra["perf"] = snapshot.perf_metrics return context