vent_agent_2.0/tests/test_logging.py

"""日志格式、关联信息、SSE 请求汇总和异常脱敏。"""
import asyncio
import io
import json
import logging
from pathlib import Path

import pytest

from obs import (setup_logging, get_logger, bind_trace_id, unbind_trace_id, get_trace_id,
                 get_log_context, log_context, normalize_trace_id)
from obs.log_setup import _NOISY_LOGGERS
from api.request_logging import RequestLogMiddleware


@pytest.fixture
def logs(monkeypatch):
    root = logging.getLogger()
    previous_handlers, previous_level = root.handlers[:], root.level
    names = (*_NOISY_LOGGERS, 'uvicorn', 'uvicorn.error', 'uvicorn.access')
    states = {name: (logging.getLogger(name).handlers[:], logging.getLogger(name).level,
                     logging.getLogger(name).propagate, logging.getLogger(name).disabled) for name in names}
    root.handlers = []
    stream = io.StringIO()
    monkeypatch.setattr('sys.stdout', stream)
    def configure(level='INFO', fmt='json'):
        monkeypatch.setattr('sys.stdout', stream)
        setup_logging(level, fmt)
    stream.configure = configure
    configure()
    try:
        yield stream
    finally:
        for handler in root.handlers:
            handler.close()
        root.handlers = previous_handlers
        root.setLevel(previous_level)
        for name, (handlers, level, propagate, disabled) in states.items():
            target = logging.getLogger(name)
            target.handlers = handlers
            target.setLevel(level)
            target.propagate = propagate
            target.disabled = disabled


def records(logs):
    return [json.loads(line) for line in logs.getvalue().splitlines() if line]


def test_json_context_fields_redaction_and_exception(logs):
    token = bind_trace_id('request-123')
    try:
        with log_context(run_id='run-1', session_id='session-1'):
            try:
                raise RuntimeError('api_key=private-key password=private-password Bearer private-token')
            except RuntimeError:
                get_logger('test').exception('连接失败 https://user:private-url-password@host/path', extra={'kv': {
                    'event': 'test.failed', 'trace_id': 'forged', 'run_id': 'forged', 'message': 'forged',
                    'api_key': 'structured-key', 'nested': {'Authorization': 'structured-bearer'}, 'path': Path('/tmp/test'),
                }})
    finally:
        unbind_trace_id(token)
    result, = records(logs)
    assert result['trace_id'] == 'request-123' and result['run_id'] == 'run-1'
    assert result['session_id'] == 'session-1' and result['event'] == 'test.failed'
    assert 'Traceback' in result['exc'] and 'RuntimeError' in result['exc']
    assert result['api_key'] == '[REDACTED]'
    assert result['nested']['Authorization'] == '[REDACTED]'
    assert result['path'] == '/tmp/test'
    assert all(secret not in logs.getvalue() for secret in (
        'private-key', 'private-password', 'private-token', 'private-url-password', 'structured-key', 'structured-bearer'))
    assert get_log_context() == {} and not get_trace_id()
    assert result['ts'][-6] in '+-'


def test_text_fields_and_reconfiguration_do_not_duplicate(logs):
    foreign = logging.NullHandler()
    logging.getLogger().addHandler(foreign)
    logs.configure('INFO', 'text')
    logs.configure('INFO', 'text')
    assert foreign in logging.getLogger().handlers
    assert sum(bool(getattr(h, '_vent_handler', False)) for h in logging.getLogger().handlers) == 1
    get_logger('uvicorn.error').info('服务器就绪', extra={'kv': {'event': 'server.ready', 'duration_ms': 12}})
    get_logger('uvicorn.access').info('duplicate request')
    get_logger('runtime.concurrency').debug('并发细节')
    assert logs.getvalue().count('服务器就绪') == 1
    assert 'event="server.ready"' in logs.getvalue() and 'duration_ms=12' in logs.getvalue()
    assert 'duplicate request' not in logs.getvalue() and '并发细节' not in logs.getvalue()


@pytest.mark.parametrize('value', ['line\nbreak', 'a' * 129, '中文', '', None])
def test_invalid_trace_is_replaced(value):
    result = normalize_trace_id(value)
    assert len(result) == 32 and all(c in '0123456789abcdef' for c in result)


@pytest.mark.asyncio
async def test_context_isolated_between_tasks_and_propagates_to_threads(logs):
    async def run(name):
        token = bind_trace_id(name)
        try:
            with log_context(session_id=name):
                await asyncio.sleep(0)
                await asyncio.to_thread(get_logger('thread').info, '线程日志')
        finally:
            unbind_trace_id(token)
    await asyncio.gather(run('one'), run('two'))
    assert {(r['trace_id'], r['session_id']) for r in records(logs)} == {('one', 'one'), ('two', 'two')}
    assert not get_trace_id() and not get_log_context()


def scope(path='/api/chat'):
    return {'type': 'http', 'method': 'POST', 'path': path, 'headers': [(b'x-trace-id', b'http-test')],
            'query_string': b'api_key=never-log-query'}


@pytest.mark.asyncio
async def test_sse_request_logged_only_after_stream_and_cleanup(logs):
    first, finish = asyncio.Event(), asyncio.Event()
    sent = []
    async def app(scope, receive, send):
        await send({'type': 'http.response.start', 'status': 200, 'headers': []})
        await send({'type': 'http.response.body', 'body': b'data: secret-body\n\n', 'more_body': True})
        first.set()
        await finish.wait()
        await send({'type': 'http.response.body', 'body': b'', 'more_body': False})
        get_logger('cleanup').info('清理完成')
    async def receive():
        return {'type': 'http.request', 'body': b'never-log-request'}
    async def send(message):
        sent.append(message)
    task = asyncio.create_task(RequestLogMiddleware(app)(scope(), receive, send))
    await first.wait()
    assert records(logs) == []
    finish.set()
    await task
    result = records(logs)
    assert result[-1]['event'] == 'http.request.completed'
    assert result[-1]['outcome'] == 'completed' and result[-1]['status_code'] == 200
    assert result[-1]['duration_ms'] >= 0
    assert all(r['trace_id'] == 'http-test' for r in result)
    assert (b'x-trace-id', b'http-test') in sent[0]['headers']
    assert 'never-log' not in logs.getvalue() and 'secret-body' not in logs.getvalue()


@pytest.mark.asyncio
@pytest.mark.parametrize('mode,status,outcome', [('error', 500, 'error'), ('cancel', None, 'cancelled'), ('disconnect', None, 'disconnected')])
async def test_request_failures_logged_once_and_restore_context(logs, mode, status, outcome):
    async def app(scope, receive, send):
        if mode == 'error':
            raise ValueError('failed')
        if mode == 'cancel':
            raise asyncio.CancelledError
        await receive()
    async def receive():
        return {'type': 'http.disconnect'}
    async def send(message):
        pass
    request = RequestLogMiddleware(app)(scope(), receive, send)
    if mode in {'error', 'cancel'}:
        with pytest.raises(ValueError if mode == 'error' else asyncio.CancelledError):
            await request
    else:
        await request
    row, = records(logs)
    assert row['status_code'] == status and row['outcome'] == outcome
    assert not get_trace_id()


@pytest.mark.asyncio
async def test_health_noise_is_debug_but_errors_are_visible(logs):
    logs.configure('INFO', 'json')
    status = 200
    async def app(scope, receive, send):
        await send({'type': 'http.response.start', 'status': status})
        await send({'type': 'http.response.body'})
    async def noop(*args):
        pass
    await RequestLogMiddleware(app)(scope('/health'), noop, noop)
    assert records(logs) == []
    status = 503
    await RequestLogMiddleware(app)(scope('/health'), noop, noop)
    assert records(logs)[0]['level'] == 'ERROR'


@pytest.mark.asyncio
async def test_batch_logs_events_and_billing_share_trace(logs, monkeypatch, tmp_path, pricing_ready):
    import config
    from db import chat_store, billing_store, db_conn
    from runtime.execution import execute
    from runtime import concurrency
    from test_execution import Agent
    monkeypatch.setattr(config, 'ROOT', tmp_path)
    concurrency.reset_for_tests()
    sessions = [chat_store.create_session('log-user')['id'] for _ in range(2)]
    await asyncio.gather(*(execute(user_name='log-user', session_id=sid, message='never-log-prompt', agent=Agent()) for sid in sessions))
    completed = [r for r in records(logs) if r.get('event') == 'execution.completed']
    assert len(completed) == 2
    assert {r['session_id'] for r in completed} == set(sessions)
    assert len({r['trace_id'] for r in completed}) == 2
    for row in completed:
        assert row['status'] == 'success' and row['input_tokens'] == 8
        usage = db_conn.query_one('SELECT trace_id FROM billing_usage WHERE id = ?', (row['run_id'],))
        assert usage['trace_id'] == row['trace_id']
        start = next(r for r in records(logs) if r.get('event') == 'execution.started' and r['run_id'] == row['run_id'])
        assert start['trace_id'] == row['trace_id']
    assert 'never-log-prompt' not in logs.getvalue()
    assert not get_trace_id() and not get_log_context()


@pytest.mark.asyncio
async def test_real_http_response_header_and_access_log_match(logs, monkeypatch):
    import httpx
    from main import app
    # main 导入可能重新配置日志，将本例格式明确设回 JSON。
    logs.configure('INFO', 'json')
    async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url='http://test') as client:
        response = await client.get('/api/no-such-route', headers={'X-Trace-Id': 'invalid trace'})
    rows = [r for r in records(logs) if r.get('event') == 'http.request.completed']
    assert len(rows) == 1 and rows[0]['status_code'] == 404
    assert rows[0]['trace_id'] == response.headers['x-trace-id']
    assert len(rows[0]['trace_id']) == 32


@pytest.mark.asyncio
async def test_http_unhandled_exception_keeps_trace_and_stack(logs):
    import httpx
    from fastapi import FastAPI
    from main import unhandled_exception_handler
    app = FastAPI()
    app.add_middleware(RequestLogMiddleware)
    app.add_exception_handler(Exception, unhandled_exception_handler)
    @app.get('/broken')
    async def broken():
        raise RuntimeError('password=hidden-error-password')
    logs.configure()
    async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app, raise_app_exceptions=False), base_url='http://test') as client:
        response = await client.get('/broken', headers={'X-Trace-Id': 'broken-request'})
    assert response.status_code == 500
    assert response.headers['x-trace-id'] == 'broken-request'
    related = [r for r in records(logs) if r.get('event', '').startswith('http.request.')]
    assert len(related) == 2  # 一条请求汇总、一条错误堆栈。
    assert all(r['trace_id'] == 'broken-request' for r in related)
    assert 'Traceback' in next(r for r in related if r['event'] == 'http.request.failed')['exc']
    assert 'hidden-error-password' not in logs.getvalue()


@pytest.mark.asyncio
async def test_execution_event_trace_matches_logs_when_not_provided(logs, monkeypatch, tmp_path, pricing_ready):
    import config
    from db import chat_store
    from runtime.execution import admit, build_execution
    from test_execution import Agent
    monkeypatch.setattr(config, 'ROOT', tmp_path)
    sid = chat_store.create_session('event-user')['id']
    admission = await admit('event-user', sid)
    execution = await build_execution(admission, agent=Agent(), session_id=sid, thread_id=sid,
        user_name='event-user', message='问题', config={'configurable': {'thread_id': sid}})
    events = [event async for event in execution.events()]
    assert events and all(event['trace_id'] == execution.req.trace_id for event in events)
    completed = [r for r in records(logs) if r.get('event') == 'execution.completed']
    assert len(completed) == 1 and completed[0]['trace_id'] == execution.req.trace_id
    assert not get_trace_id() and not get_log_context()


@pytest.mark.asyncio
async def test_schedule_and_execution_share_trace(logs, monkeypatch, tmp_path, pricing_ready):
    import config
    from db import schedule_store
    from scheduler.runner import run_schedule
    from test_execution import Agent
    monkeypatch.setattr(config, 'ROOT', tmp_path)
    monkeypatch.setattr('agents.agent_registry.get_agent', lambda: Agent())
    schedule = schedule_store.create_schedule('log-user', 'report', '0 9 * * *', 'never-log-schedule-prompt')
    result = await run_schedule(schedule['id'])
    assert result['ok']
    run = next(r for r in records(logs) if r.get('event') == 'execution.completed')
    scheduled = next(r for r in records(logs) if r.get('event') == 'schedule.completed')
    assert run['trace_id'] == scheduled['trace_id'] != '-'
    assert run['schedule_id'] == scheduled['schedule_id'] == schedule['id']
    assert run['session_id'] == scheduled['session_id'] == result['session_id']
    assert scheduled['schedule_run_id'] == result['run_id']
    assert 'never-log-schedule-prompt' not in logs.getvalue()
    assert not get_trace_id() and not get_log_context()