Aether/_deprecated_py_src/api/handlers/base/stream_telemetry.py

"""
流式遥测记录器 - 从 ChatHandlerBase 提取的统计记录逻辑

职责：
1. 记录流式请求的成功/失败统计
2. 更新 Usage 状态
3. 更新候选记录状态
"""

import asyncio
import json
import time
from typing import Any

from sqlalchemy.orm import Session

from src.api.handlers.base.base_handler import MessageTelemetry
from src.api.handlers.base.stream_context import StreamContext
from src.api.handlers.base.utils import filter_proxy_response_headers
from src.config.settings import config
from src.core.logger import logger
from src.database import get_db
from src.models.database import ApiKey, User
from src.services.system.config import SystemConfigService
from src.services.usage.telemetry_writer import (
    DbTelemetryWriter,
    QueueTelemetryWriter,
    TelemetryWriter,
)


class StreamTelemetryRecorder:
    """
    流式遥测记录器

    负责在流式请求完成后记录统计信息。
    从 ChatHandlerBase 中提取的 _record_stream_stats 逻辑。
    """

    def __init__(
        self,
        request_id: str,
        user_id: str,
        api_key_id: str,
        client_ip: str,
        format_id: str,
    ):
        """
        初始化遥测记录器

        Args:
            request_id: 请求 ID
            user_id: 用户 ID
            api_key_id: API Key ID
            client_ip: 客户端 IP
            format_id: API 格式标识
        """
        self.request_id = request_id
        self.user_id = user_id
        self.api_key_id = api_key_id
        self.client_ip = client_ip
        self.format_id = format_id

    async def record_stream_stats(
        self,
        ctx: StreamContext,
        original_headers: dict[str, str],
        original_request_body: dict[str, Any],
        start_time: float,
    ) -> None:
        """
        记录流式统计信息

        Args:
            ctx: 流式上下文
            original_headers: 原始请求头
            original_request_body: 原始请求体
            start_time: 请求开始时间 (time.time())
        """
        bg_db = None

        try:
            # 在流结束后计算响应时间，与首字时间使用相同的时间基准
            # 注意：不要把统计延迟（stream_stats_delay）算进响应时间里
            response_time_ms = int((time.time() - start_time) * 1000)

            await asyncio.sleep(config.stream_stats_delay)  # 等待流完全关闭

            if not ctx.provider_name:
                await self._update_usage_status_on_error(
                    response_time_ms=response_time_ms,
                    error_message="Provider name not available",
                )
                return

            db_gen = get_db()
            bg_db = next(db_gen)

            try:
                writer = await self._get_telemetry_writer(bg_db, ctx, response_time_ms)
                if writer is None:
                    ctx.release_recorded_chunks()
                    return
                # 兜底估算：流未正常完成且 token 均为 0 时，从请求体粗略估算。
                # 覆盖成功但缺少 completion，以及已传出部分数据后被中断的场景。
                if ctx.should_estimate_incomplete_tokens():
                    # 用实际发给 Provider 的请求体估算 token（格式转换时与客户端请求体不同）
                    self._estimate_tokens_for_incomplete_stream(
                        ctx, ctx.provider_request_body or original_request_body
                    )

                should_log_body = SystemConfigService.should_log_body(bg_db)
                include_bodies = (
                    writer.include_bodies
                    if isinstance(writer, QueueTelemetryWriter)
                    else should_log_body
                )
                with ctx.managed_recorded_bodies(
                    response_time_ms, include_bodies=include_bodies
                ) as recorded_bodies:
                    try:
                        await self._dispatch_record(
                            bg_db,
                            writer,
                            ctx,
                            original_headers,
                            original_request_body,
                            recorded_bodies.response_body,
                            response_time_ms,
                            client_response_body=recorded_bodies.client_response_body,
                        )
                    except Exception as writer_error:
                        if not isinstance(writer, QueueTelemetryWriter):
                            raise
                        logger.warning(
                            f"[{self.request_id}] Queue writer failed, falling back to DB: {writer_error}"
                        )
                        db_writer = self._build_db_writer(bg_db)
                        if db_writer is None:
                            await self._update_usage_status_directly(
                                bg_db,
                                status=self._get_status_from_ctx(ctx),
                                response_time_ms=response_time_ms,
                                status_code=ctx.status_code,
                            )
                            return
                        if should_log_body:
                            recorded_bodies.ensure_populated(ctx, response_time_ms)
                        await self._dispatch_record(
                            bg_db,
                            db_writer,
                            ctx,
                            original_headers,
                            original_request_body,
                            recorded_bodies.response_body,
                            response_time_ms,
                            client_response_body=recorded_bodies.client_response_body,
                        )

                # 更新候选记录状态
                await self._update_candidate_status(bg_db, ctx, response_time_ms, start_time)

            finally:
                if bg_db:
                    bg_db.close()

        except Exception as e:
            logger.exception("记录流式统计信息时出错")
            await self._update_usage_status_on_error(
                response_time_ms=response_time_ms,
                error_message=f"记录统计信息失败: {str(e)[:200]}",
            )
        finally:
            # 遥测写入后主动释放大列表，避免长流式请求对象滞留在 worker 堆中。
            ctx.release_recorded_chunks()

    async def _record_success(
        self,
        writer: TelemetryWriter,
        ctx: StreamContext,
        original_headers: dict[str, str],
        original_request_body: dict[str, Any],
        response_body: dict[str, Any] | None,
        response_time_ms: int,
        client_response_body: dict[str, Any] | None = None,
    ) -> None:
        """记录成功的请求"""
        # 流式成功时，返回给客户端的是提供商响应头 + SSE 必需头
        client_response_headers = filter_proxy_response_headers(ctx.response_headers)
        client_response_headers.update(
            {
                "Cache-Control": "no-cache, no-transform",
                "X-Accel-Buffering": "no",
                "content-type": "text/event-stream",
            }
        )
        metadata: dict[str, Any] = {"stream": True, "content_length": ctx.data_count}
        if ctx.perf_metrics:
            metadata["perf"] = ctx.perf_metrics
        if ctx.proxy_info:
            metadata["proxy"] = ctx.proxy_info
        if ctx.pool_summary:
            metadata["pool_summary"] = ctx.pool_summary
        if ctx.candidate_keys:
            metadata["candidate_keys"] = ctx.candidate_keys
        if ctx.scheduling_audit:
            metadata["scheduling_audit"] = ctx.scheduling_audit

        await writer.record_success(
            provider=ctx.provider_name or "unknown",
            model=ctx.model,
            input_tokens=ctx.input_tokens,
            output_tokens=ctx.output_tokens,
            response_time_ms=response_time_ms,
            first_byte_time_ms=ctx.first_byte_time_ms,  # 传递首字时间
            status_code=ctx.status_code,
            request_headers=original_headers,
            request_body=original_request_body,
            response_headers=ctx.response_headers,
            client_response_headers=client_response_headers,
            response_body=response_body,
            client_response_body=client_response_body,
            provider_request_body=ctx.provider_request_body,
            cache_creation_tokens=ctx.cache_creation_tokens,
            cache_read_tokens=ctx.cached_tokens,
            cache_creation_tokens_5m=ctx.cache_creation_tokens_5m,
            cache_creation_tokens_1h=ctx.cache_creation_tokens_1h,
            is_stream=True,
            provider_request_headers=ctx.provider_request_headers,
            api_format=ctx.api_format,
            api_family=ctx.api_family,
            endpoint_kind=ctx.endpoint_kind,
            provider_id=ctx.provider_id,
            provider_endpoint_id=ctx.endpoint_id,
            provider_api_key_id=ctx.key_id,
            target_model=ctx.mapped_model,
            request_type="chat",
            metadata=metadata,
            endpoint_api_format=ctx.provider_api_format,
            has_format_conversion=ctx.has_format_conversion,
        )

        logger.debug(f"{self.format_id} 流式响应完成")
        logger.info(ctx.get_log_summary(self.request_id, response_time_ms))

    async def _record_failure(
        self,
        writer: TelemetryWriter,
        ctx: StreamContext,
        original_headers: dict[str, str],
        original_request_body: dict[str, Any],
        response_body: dict[str, Any] | None,
        response_time_ms: int,
        client_response_body: dict[str, Any] | None = None,
    ) -> None:
        """记录失败的请求"""
        # 失败时返回给客户端的是 JSON 错误响应，如果没有设置则使用默认值
        client_response_headers = ctx.client_response_headers or {
            "content-type": "application/json"
        }
        metadata: dict[str, Any] = {"stream": True, "content_length": ctx.data_count}
        if ctx.perf_metrics:
            metadata["perf"] = ctx.perf_metrics
        if ctx.proxy_info:
            metadata["proxy"] = ctx.proxy_info
        if ctx.pool_summary:
            metadata["pool_summary"] = ctx.pool_summary
        if ctx.candidate_keys:
            metadata["candidate_keys"] = ctx.candidate_keys
        if ctx.scheduling_audit:
            metadata["scheduling_audit"] = ctx.scheduling_audit

        await writer.record_failure(
            provider=ctx.provider_name or "unknown",
            model=ctx.model,
            response_time_ms=response_time_ms,
            status_code=ctx.status_code,
            error_message=ctx.error_message or f"HTTP {ctx.status_code}",
            request_headers=original_headers,
            request_body=original_request_body,
            is_stream=True,
            api_format=ctx.api_format,
            api_family=ctx.api_family,
            endpoint_kind=ctx.endpoint_kind,
            provider_request_headers=ctx.provider_request_headers,
            provider_request_body=ctx.provider_request_body,
            input_tokens=ctx.input_tokens,
            output_tokens=ctx.output_tokens,
            cache_creation_tokens=ctx.cache_creation_tokens,
            cache_read_tokens=ctx.cached_tokens,
            cache_creation_tokens_5m=ctx.cache_creation_tokens_5m,
            cache_creation_tokens_1h=ctx.cache_creation_tokens_1h,
            response_body=response_body,
            client_response_body=client_response_body,
            response_headers=ctx.response_headers,
            client_response_headers=client_response_headers,
            provider_id=ctx.provider_id,
            provider_endpoint_id=ctx.endpoint_id,
            provider_api_key_id=ctx.key_id,
            target_model=ctx.mapped_model,
            request_type="chat",
            metadata=metadata,
            endpoint_api_format=ctx.provider_api_format,
            has_format_conversion=ctx.has_format_conversion,
        )

        logger.debug(f"{self.format_id} 流式响应中断")
        log_summary = ctx.get_log_summary(self.request_id, response_time_ms)
        # 对于失败日志，添加缓存信息
        logger.info(f"{log_summary} cache:{ctx.cached_tokens}")

    async def _record_cancelled(
        self,
        writer: TelemetryWriter,
        ctx: StreamContext,
        original_headers: dict[str, str],
        original_request_body: dict[str, Any],
        response_body: dict[str, Any] | None,
        response_time_ms: int,
        client_response_body: dict[str, Any] | None = None,
    ) -> None:
        """记录客户端取消的请求"""
        client_response_headers = ctx.client_response_headers or {
            "content-type": "application/json"
        }
        metadata: dict[str, Any] = {"stream": True, "content_length": ctx.data_count}
        if ctx.perf_metrics:
            metadata["perf"] = ctx.perf_metrics
        if ctx.proxy_info:
            metadata["proxy"] = ctx.proxy_info
        if ctx.pool_summary:
            metadata["pool_summary"] = ctx.pool_summary
        if ctx.candidate_keys:
            metadata["candidate_keys"] = ctx.candidate_keys
        if ctx.scheduling_audit:
            metadata["scheduling_audit"] = ctx.scheduling_audit

        await writer.record_cancelled(
            provider=ctx.provider_name or "unknown",
            model=ctx.model,
            response_time_ms=response_time_ms,
            first_byte_time_ms=ctx.first_byte_time_ms,
            status_code=ctx.status_code,
            request_headers=original_headers,
            request_body=original_request_body,
            is_stream=True,
            api_format=ctx.api_format,
            api_family=ctx.api_family,
            endpoint_kind=ctx.endpoint_kind,
            provider_request_headers=ctx.provider_request_headers,
            provider_request_body=ctx.provider_request_body,
            input_tokens=ctx.input_tokens,
            output_tokens=ctx.output_tokens,
            cache_creation_tokens=ctx.cache_creation_tokens,
            cache_read_tokens=ctx.cached_tokens,
            cache_creation_tokens_5m=ctx.cache_creation_tokens_5m,
            cache_creation_tokens_1h=ctx.cache_creation_tokens_1h,
            response_body=response_body,
            client_response_body=client_response_body,
            response_headers=ctx.response_headers,
            client_response_headers=client_response_headers,
            provider_id=ctx.provider_id,
            provider_endpoint_id=ctx.endpoint_id,
            provider_api_key_id=ctx.key_id,
            target_model=ctx.mapped_model,
            request_type="chat",
            metadata=metadata,
            endpoint_api_format=ctx.provider_api_format,
            has_format_conversion=ctx.has_format_conversion,
        )

        logger.debug(f"{self.format_id} 流式响应被客户端取消")
        logger.info(ctx.get_log_summary(self.request_id, response_time_ms))

    async def _update_candidate_status(
        self,
        db: Session,
        ctx: StreamContext,
        response_time_ms: int,
        request_start_time: float,
    ) -> None:
        """更新候选记录状态"""
        if not ctx.attempt_id:
            return

        # Capture all needed ctx attrs before handing off to thread
        attempt_id = ctx.attempt_id
        is_success = ctx.is_success()
        is_disconnected = ctx.is_client_disconnected()
        status_code = ctx.status_code
        data_count = ctx.data_count
        rectified = ctx.rectified
        proxy_info = ctx.proxy_info
        first_byte_time_ms_val = ctx.first_byte_time_ms
        upstream_response = ctx.upstream_response
        error_message = ctx.error_message

        def _sync() -> None:
            from src.services.request.candidate import RequestCandidateService

            extra_data: dict[str, Any] = {
                "stream_completed": is_success,
                "data_count": data_count,
            }
            if rectified:
                extra_data["rectified"] = True
            if proxy_info:
                extra_data["proxy"] = proxy_info
            if first_byte_time_ms_val is not None:
                candidate_ttfb = RequestCandidateService.calculate_candidate_ttfb(
                    db=db,
                    candidate_id=attempt_id,
                    request_start_time=request_start_time,
                    global_first_byte_time_ms=first_byte_time_ms_val,
                )
                extra_data["first_byte_time_ms"] = candidate_ttfb

            if is_success:
                RequestCandidateService.mark_candidate_success(
                    db=db,
                    candidate_id=attempt_id,
                    status_code=status_code,
                    latency_ms=response_time_ms,
                    extra_data=extra_data,
                )
            elif is_disconnected:
                RequestCandidateService.mark_candidate_cancelled(
                    db=db,
                    candidate_id=attempt_id,
                    status_code=status_code,
                    latency_ms=response_time_ms,
                    extra_data=extra_data,
                )
            else:
                trace_error_message = upstream_response or error_message or f"HTTP {status_code}"
                RequestCandidateService.mark_candidate_failed(
                    db=db,
                    candidate_id=attempt_id,
                    error_type="stream_error",
                    error_message=trace_error_message,
                    status_code=status_code,
                    latency_ms=response_time_ms,
                    extra_data=extra_data,
                )

        await asyncio.to_thread(_sync)

    async def _update_usage_status_on_error(
        self,
        response_time_ms: int,
        error_message: str,
    ) -> None:
        """在记录失败时更新 Usage 状态"""
        try:
            db_gen = get_db()
            error_db = next(db_gen)
            try:
                await self._update_usage_status_directly(
                    error_db,
                    status="failed",
                    response_time_ms=response_time_ms,
                    status_code=500,
                    error_message=error_message,
                )
            finally:
                error_db.close()
        except Exception as inner_e:
            logger.error(f"[{self.request_id}] 更新 Usage 状态失败: {inner_e}")

    async def _update_usage_status_directly(
        self,
        db: Session,
        status: str,
        response_time_ms: int,
        status_code: int = 200,
        error_message: str | None = None,
    ) -> None:
        """直接更新 Usage 表的状态字段"""
        request_id = self.request_id

        def _sync() -> None:
            from src.models.database import Usage

            usage = db.query(Usage).filter(Usage.request_id == request_id).first()
            if usage:
                if getattr(usage, "billing_status", None) in {"settled", "void"}:
                    logger.debug("[{}] Usage 已终态，跳过快速状态更新: {}", self.request_id, status)
                    return
                setattr(usage, "status", status)
                setattr(usage, "status_code", status_code)
                setattr(usage, "response_time_ms", response_time_ms)
                if error_message:
                    setattr(usage, "error_message", error_message)
                db.commit()
                logger.debug(f"[{request_id}] Usage 状态已更新: {status}")

        try:
            await asyncio.to_thread(_sync)
        except Exception as e:
            logger.error(f"[{self.request_id}] 直接更新 Usage 状态失败: {e}")

    async def _get_telemetry_writer(
        self, bg_db: Session, ctx: StreamContext, response_time_ms: int
    ) -> TelemetryWriter | None:
        if config.usage_queue_enabled and self.user_id and self.api_key_id:
            # Queue payload detail follows system config request_record_level.
            log_level = SystemConfigService.get_request_record_level(bg_db).value
            sensitive_headers = SystemConfigService.get_sensitive_headers(bg_db) or []
            max_request_body_size = int(
                SystemConfigService.get_config(bg_db, "max_request_body_size", 5242880) or 0
            )
            max_response_body_size = int(
                SystemConfigService.get_config(bg_db, "max_response_body_size", 5242880) or 0
            )
            return QueueTelemetryWriter(
                request_id=self.request_id,
                user_id=self.user_id,
                api_key_id=self.api_key_id,
                log_level=log_level,
                sensitive_headers=sensitive_headers,
                max_request_body_size=max_request_body_size,
                max_response_body_size=max_response_body_size,
            )
        db_writer = self._build_db_writer(bg_db)
        if db_writer is None:
            await self._update_usage_status_directly(
                bg_db,
                status=self._get_status_from_ctx(ctx),
                response_time_ms=response_time_ms,
                status_code=ctx.status_code,
            )
            return None
        return db_writer

    async def _dispatch_record(
        self,
        db: Session,
        writer: TelemetryWriter,
        ctx: StreamContext,
        original_headers: dict[str, str],
        original_request_body: dict[str, Any],
        response_body: dict[str, Any] | None,
        response_time_ms: int,
        client_response_body: dict[str, Any] | None = None,
    ) -> None:
        """根据上下文状态分发到对应的记录方法"""
        if ctx.is_success():
            await self._record_success(
                writer,
                ctx,
                original_headers,
                original_request_body,
                response_body,
                response_time_ms,
                client_response_body=client_response_body,
            )
            # Queue writer 异步落库可能造成 UI 延迟，先直接更新 Usage 状态
            if isinstance(writer, QueueTelemetryWriter):
                await self._update_usage_status_directly(
                    db=db,
                    status=self._get_status_from_ctx(ctx),
                    response_time_ms=response_time_ms,
                    status_code=ctx.status_code,
                )
        elif ctx.is_client_disconnected():
            await self._record_cancelled(
                writer,
                ctx,
                original_headers,
                original_request_body,
                response_body,
                response_time_ms,
                client_response_body=client_response_body,
            )
            # Queue writer 异步落库可能造成 UI 延迟，先直接更新 Usage 状态
            if isinstance(writer, QueueTelemetryWriter):
                await self._update_usage_status_directly(
                    db=db,
                    status="cancelled",
                    response_time_ms=response_time_ms,
                    status_code=ctx.status_code,
                )
        else:
            await self._record_failure(
                writer,
                ctx,
                original_headers,
                original_request_body,
                response_body,
                response_time_ms,
                client_response_body=client_response_body,
            )
            # Queue writer 异步落库可能造成 UI 延迟，先直接更新 Usage 状态
            if isinstance(writer, QueueTelemetryWriter):
                await self._update_usage_status_directly(
                    db=db,
                    status=self._get_status_from_ctx(ctx),
                    response_time_ms=response_time_ms,
                    status_code=ctx.status_code,
                    error_message=ctx.error_message or f"HTTP {ctx.status_code}",
                )

    def _get_status_from_ctx(self, ctx: StreamContext) -> str:
        """根据上下文获取状态字符串"""
        if ctx.is_success():
            return "completed"
        if ctx.is_client_disconnected():
            return "cancelled"
        return "failed"

    @staticmethod
    def _estimate_tokens_for_incomplete_stream(
        ctx: StreamContext,
        request_body: dict[str, Any],
    ) -> None:
        """
        流未正常完成（无 response.completed）且 token 均为 0 时的兜底估算。

        从已收集的输出文本和请求体粗略估算 token 数，确保 usage 记录不为 0。
        估算采用 ~4 字符/token 的保守比例。
        """
        # 输出 tokens：从已收集的文本估算
        if ctx.collected_text_length > 0:
            ctx.output_tokens = max(1, ctx.collected_text_length // 4)

        # 输入 tokens：从请求体文本内容估算
        try:
            total_input_len = 0
            instructions = request_body.get("instructions")
            if isinstance(instructions, str):
                total_input_len += len(instructions)
            # OpenAI Responses API 使用 input 字段；Claude 使用 messages
            input_items = request_body.get("input") or request_body.get("messages") or []
            if isinstance(input_items, list):
                for item in input_items:
                    if isinstance(item, str):
                        total_input_len += len(item)
                    elif isinstance(item, dict):
                        content = item.get("content", "")
                        if isinstance(content, str):
                            total_input_len += len(content)
                        elif isinstance(content, list):
                            for block in content:
                                if isinstance(block, dict):
                                    text = block.get("text", "")
                                    if isinstance(text, str):
                                        total_input_len += len(text)
            if total_input_len > 0:
                ctx.input_tokens = max(1, total_input_len // 4)
            else:
                # fallback: 整个请求体 JSON 大小
                body_str = json.dumps(request_body, ensure_ascii=False)
                ctx.input_tokens = max(1, len(body_str) // 4)
        except Exception:
            pass

        if ctx.input_tokens > 0 or ctx.output_tokens > 0:
            logger.warning(
                f"[{ctx.request_id}] 流未正常完成 (has_completion=False, data_count={ctx.data_count}), "
                f"使用估算 tokens: in={ctx.input_tokens}, out={ctx.output_tokens}"
            )

    def _build_db_writer(self, bg_db: Session) -> DbTelemetryWriter | None:
        user = bg_db.query(User).filter(User.id == self.user_id).first()
        api_key_obj = bg_db.query(ApiKey).filter(ApiKey.id == self.api_key_id).first()

        if not user or not api_key_obj:
            logger.warning(
                f"[{self.request_id}] User or ApiKey not found, updating status directly"
            )
            return None

        bg_telemetry = MessageTelemetry(bg_db, user, api_key_obj, self.request_id, self.client_ip)
        return DbTelemetryWriter(bg_telemetry)