vent_agent_2.0/api/request_logging.py

"""ASGI 请求日志：覆盖完整 SSE 生命周期，不读取请求正文或查询参数。"""
import asyncio
import logging
import time

from obs import bind_trace_id, unbind_trace_id, normalize_trace_id, get_logger

logger = get_logger('api.access')


class RequestLogMiddleware:
    def __init__(self, app, *, trace_enabled=True):
        self.app = app
        self.trace_enabled = trace_enabled

    async def __call__(self, scope, receive, send):
        if scope['type'] != 'http':
            return await self.app(scope, receive, send)
        started = time.perf_counter()
        headers = dict(scope.get('headers', []))
        trace_id = normalize_trace_id(headers.get(b'x-trace-id', b'').decode('latin1'))
        token = bind_trace_id(trace_id) if self.trace_enabled else None
        if self.trace_enabled:
            scope.setdefault('state', {})['trace_id'] = trace_id
        status, completed, disconnected, failure = None, False, False, None

        async def logged_receive():
            nonlocal disconnected
            message = await receive()
            if message['type'] == 'http.disconnect':
                disconnected = True
            return message

        async def logged_send(message):
            nonlocal status, completed
            if message['type'] == 'http.response.start':
                status = message['status']
                if self.trace_enabled:
                    message = {**message, 'headers': [(k, v) for k, v in message.get('headers', [])
                                                       if k.lower() != b'x-trace-id']
                               + [(b'x-trace-id', trace_id.encode('ascii'))]}
            await send(message)
            if message['type'] == 'http.response.body' and not message.get('more_body', False):
                completed = True

        try:
            await self.app(scope, logged_receive, logged_send)
        except BaseException as exc:
            failure = exc
            raise
        finally:
            try:
                outcome = ('cancelled' if isinstance(failure, asyncio.CancelledError)
                           else 'error' if failure else 'completed' if completed
                           else 'disconnected' if disconnected else 'incomplete')
                code = status if status is not None else (500 if failure and outcome == 'error' else None)
                level = logging.ERROR if outcome == 'error' or (code and code >= 500) else (
                    logging.WARNING if outcome != 'completed' or (code and code >= 400) else logging.INFO)
                path = scope.get('path', '')
                if level == logging.INFO and (path == '/health' or path.startswith('/static/') or path == '/favicon.ico'):
                    level = logging.DEBUG
                route = scope.get('route')
                logger.log(level, '请求结束', extra={'kv': {
                    'event': 'http.request.completed', 'method': scope.get('method'),
                    'path': getattr(route, 'path', path), 'status_code': code, 'outcome': outcome,
                    'duration_ms': round((time.perf_counter() - started) * 1000, 1),
                    **({'error_type': type(failure).__name__} if failure else {}),
                }})
            finally:
                if token is not None:
                    unbind_trace_id(token)