Artemis-Eapil/tests/unit/test_request_diagnostics.py

import asyncio
from copy import deepcopy
import json
from types import SimpleNamespace
from uuid import uuid4

import httpx
from langchain_core.messages import AIMessageChunk, HumanMessage
import pytest

from artemis.llm.diagnostics import (
    ACTIVE_ATTEMPT,
    RequestAttempt,
    error_text,
    input_images,
    observe_request_headers,
    observe_response_headers,
)
from artemis.llm.router import ModelEndpoint
from artemis.services import llm


class Engine:
    current_session_id = "diagnostic-test"
    current_step_id = None

    def __init__(self):
        self.records = []

    def record_trace(self, **record):
        self.records.append(deepcopy(record))
        return record.get("trace_id", uuid4())

    def _publish(self, *_):
        pass


@pytest.mark.asyncio
async def test_request_hook_correlates_api_calls_without_logging_header_values():
    engine = Engine()
    attempt = RequestAttempt(
        engine, request_id="local-request", attempt=1, timeout=300, max_retries=1, metadata={}
    )
    token = ACTIVE_ATTEMPT.set(attempt)
    try:
        request = httpx.Request(
            "POST",
            "https://example.test/v1/responses",
            headers={
                "Thread-Id": "private-thread",
                "Session-Id": "private-session",
                "Authorization": "Bearer private-key",
            },
        )
        await observe_request_headers(request)
        assert request.headers["x-client-request-id"] == "local-request"
        assert attempt.data["client_request_id_matches"] is True
        assert len(attempt.data["routing_header_fingerprints"]) == 2
        assert "private" not in json.dumps(engine.records, default=str)
        request.headers["X-Client-Request-Id"] = "caller-id"
        await observe_request_headers(request)
        assert request.headers["x-client-request-id"] == "caller-id"
        assert attempt.data["client_request_id_matches"] is False
    finally:
        ACTIVE_ATTEMPT.reset(token)


@pytest.mark.asyncio
async def test_headers_without_body_timeout_then_retry_keeps_evidence(monkeypatch):
    engine = Engine()
    monkeypatch.setattr(llm, "_CURRENT_DATA_ENGINE", engine)
    monkeypatch.setattr(
        llm, "retry_policy_for", lambda _: SimpleNamespace(delay_for=lambda _: 0, max_attempts=2)
    )
    llm._ENDPOINT_BREAKER._states.clear()
    calls = 0

    class StalledBody(httpx.AsyncByteStream):
        async def __aiter__(self):
            await asyncio.sleep(10)
            yield b""

    async def transport(request):
        nonlocal calls
        calls += 1
        headers = {"x-request-id": f"server-{calls}", "set-cookie": "secret-cookie"}
        return httpx.Response(
            200, headers=headers, stream=StalledBody() if calls == 1 else httpx.ByteStream(b"ok")
        )

    async with httpx.AsyncClient(
        transport=httpx.MockTransport(transport),
        event_hooks={"response": [observe_response_headers]},
    ) as client:

        class Model:
            async def astream(self, *_args, **_kwargs):
                async with client.stream(
                    "POST",
                    "https://example.test/responses",
                    headers={"Authorization": "Bearer secret-key"},
                ) as response:
                    await response.aread()
                yield AIMessageChunk(
                    content="ready",
                    usage_metadata={"input_tokens": 10, "output_tokens": 2, "total_tokens": 12},
                )

        wrapper = llm.RobustChatModelWrapper(
            Model(), endpoint=ModelEndpoint(timeout_seconds=0.05, max_retries=1)
        )
        result = await wrapper.complete([])
    assert result.content == "ready"
    snapshots = [r["payload"] for r in engine.records if r["type"] == "llm_diagnostic"]
    failed = next(p for p in snapshots if p["phase"] == "failed")
    final = snapshots[-1]
    assert failed["http_status"] == 200
    assert failed["server_request_id"] == "server-1"
    assert failed["failure_phase"] == "waiting_output"
    assert failed["error_type"] == "TimeoutError" and "0.05 秒" in failed["error"]
    assert failed["usage"] is None and "first_output_at" not in failed
    assert final["request_id"] == failed["request_id"]
    assert final["attempt_id"] != failed["attempt_id"] and final["attempt"] == 2
    assert final["phase"] == "completed" and final["usage"]["total_tokens"] == 12
    assert (
        final["started_at"] <= final["headers_at"] <= final["first_output_at"] <= final["ended_at"]
    )
    assert "secret" not in json.dumps(snapshots)
    assert ACTIVE_ATTEMPT.get() is None
    assert any(r["name"] == "llm_retry" and r["payload"]["error"] for r in engine.records)


def test_tool_deltas_are_counted_once_without_retaining_arguments():
    engine = Engine()
    attempt = RequestAttempt(
        engine, request_id="tools", attempt=1, timeout=300, max_retries=1, metadata={}
    )
    deltas = [
        AIMessageChunk(
            content=[], tool_call_chunks=[{"name": "click", "args": "", "id": "c1", "index": 0}]
        ),
        AIMessageChunk(
            content=[{"type": "function_call", "arguments": '{"secret":'}],
            tool_call_chunks=[{"args": '{"secret":', "index": 0}],
        ),
        AIMessageChunk(content=[], tool_call_chunks=[{"args": '"private"}', "index": 0}]),
        AIMessageChunk(
            content=[],
            tool_call_chunks=[{"name": "wait", "args": '{"seconds":1}', "id": "c2", "index": 1}],
        ),
    ]
    combined = deltas[0]
    for i, delta in enumerate(deltas):
        attempt.chunk(delta)
        if i:
            combined += delta
    attempt.chunk(AIMessageChunk(content=[], response_metadata={"status": "completed"}))
    attempt.finish("completed", response=combined)
    assert attempt.data["output_chunks"] == 4
    assert attempt.data["tool_argument_chars"] == len('{"secret":"private"}{"seconds":1}')
    assert attempt.data["tools"] == [
        {"name": "click", "argument_chars": len('{"secret":"private"}')},
        {"name": "wait", "argument_chars": len('{"seconds":1}')},
    ]
    assert attempt.data["response_status"] == "completed"
    assert "private" not in json.dumps(engine.records, default=str)
    assert "secret" not in json.dumps(engine.records, default=str)


@pytest.mark.asyncio
async def test_timeout_during_tool_generation_keeps_last_output_and_tool_name(monkeypatch):
    engine = Engine()
    monkeypatch.setattr(llm, "_CURRENT_DATA_ENGINE", engine)
    llm._ENDPOINT_BREAKER._states.clear()
    calls = 0

    class Model:
        async def astream(self, *_args, **_kwargs):
            nonlocal calls
            calls += 1
            yield AIMessageChunk(
                content=[],
                tool_call_chunks=[{"name": "click", "args": '{"x":', "id": "c1", "index": 0}],
            )
            if calls == 1:
                await asyncio.sleep(10)
            yield AIMessageChunk(content=[], tool_call_chunks=[{"args": "1}", "index": 0}])

    monkeypatch.setattr(
        llm, "retry_policy_for", lambda _: SimpleNamespace(delay_for=lambda _: 0, max_attempts=2)
    )
    wrapper = llm.RobustChatModelWrapper(
        Model(), endpoint=ModelEndpoint(timeout_seconds=0.03, max_retries=1, max_tokens=12000)
    )
    await wrapper.complete([])
    failed = next(
        r["payload"]
        for r in engine.records
        if r["type"] == "llm_diagnostic" and r["payload"]["phase"] == "failed"
    )
    assert failed["failure_phase"] == "generating_tool_arguments"
    assert failed["last_output_at"] >= failed["first_output_at"]
    assert failed["tools"] == [{"name": "click", "argument_chars": 5}]
    assert failed["usage"] is None
    assert failed["max_output_tokens"] == 12000


@pytest.mark.asyncio
async def test_concurrent_requests_do_not_mix_server_ids():
    engine = Engine()

    async def run(name):
        attempt = RequestAttempt(
            engine, request_id=name, attempt=1, timeout=5, max_retries=0, metadata={"model": name}
        )
        token = ACTIVE_ATTEMPT.set(attempt)
        try:
            await asyncio.sleep(0)
            await observe_response_headers(httpx.Response(200, headers={"x-request-id": name}))
        finally:
            ACTIVE_ATTEMPT.reset(token)
        return attempt.data

    first, second = await asyncio.gather(run("one"), run("two"))
    assert first["server_request_id"] == "one"
    assert second["server_request_id"] == "two"


def test_error_redaction_and_image_headers():
    import base64
    import io
    from PIL import Image

    out = io.BytesIO()
    Image.new("RGB", (32, 48)).save(out, "PNG")
    message = HumanMessage(
        content=[
            {
                "type": "image_url",
                "image_url": {
                    "url": "data:image/png;base64," + base64.b64encode(out.getvalue()).decode()
                },
            }
        ]
    )
    assert input_images(([message],), {}) == [{"width": 32, "height": 48}]
    text = error_text(
        RuntimeError("Authorization: Bearer sk-my-secret https://host/path?token=private")
    )
    assert "sk-my-secret" not in text and "private" not in text
    assert error_text(RuntimeError()) == "RuntimeError"


def test_snapshots_survive_out_of_order_writes_and_export(tmp_path):
    from artemis.data_engine.models import TraceRecord
    from artemis.data_engine.storage import StorageManager
    from apps.admin_console.database.connection import db_session
    from apps.admin_console.services.request_diagnostics import request_diagnostics

    path = tmp_path / "data.db"
    storage = StorageManager(db_path=path, base_trace_dir=tmp_path)
    session_id, trace_id = uuid4(), uuid4()
    with db_session(path) as conn:
        conn.execute(
            "INSERT INTO sessions(session_id,status) VALUES (?, 'completed')", (str(session_id),)
        )
        conn.commit()
    for sequence, phase in [(3, "completed"), (2, "receiving_output"), (1, "waiting_output")]:
        storage.create_trace(
            TraceRecord(
                trace_id=trace_id,
                session_id=session_id,
                type="llm_diagnostic",
                name="llm_request",
                timestamp=10 + sequence,
                status="success" if sequence == 3 else "running",
                payload={
                    "sequence": sequence,
                    "phase": phase,
                    "request_id": "call-one",
                    "started_at": 10,
                    "ended_at": 13 if sequence == 3 else None,
                },
            )
        )
    result = request_diagnostics(str(session_id), db_path=path)
    assert len(result["requests"]) == 1
    assert result["requests"][0]["phase"] == "completed"
    assert result["requests"][0]["active"] is False
    assert result["requests"][0]["timestamps_utc"]["started_at"].endswith("+00:00")


@pytest.mark.asyncio
async def test_cancellation_closes_attempt_and_does_not_retry(monkeypatch):
    engine = Engine()
    monkeypatch.setattr(llm, "_CURRENT_DATA_ENGINE", engine)
    llm._ENDPOINT_BREAKER._states.clear()

    async def cancelled():
        raise asyncio.CancelledError()

    with pytest.raises(asyncio.CancelledError):
        await llm._run_with_recovery(cancelled, max_retries=5)
    assert engine.records[-1]["payload"]["phase"] == "cancelled"
    assert ACTIVE_ATTEMPT.get() is None


@pytest.mark.asyncio
async def test_cancelling_retry_wait_does_not_leave_a_live_request(monkeypatch):
    waiting = asyncio.Event()

    class RetryEngine(Engine):
        def record_trace(self, **record):
            result = super().record_trace(**record)
            if record["payload"].get("phase") == "retry_wait":
                waiting.set()
            return result

    engine = RetryEngine()
    monkeypatch.setattr(llm, "_CURRENT_DATA_ENGINE", engine)
    monkeypatch.setattr(
        llm, "retry_policy_for", lambda _: SimpleNamespace(delay_for=lambda _: 60, max_attempts=2)
    )

    async def timeout():
        raise TimeoutError()

    task = asyncio.create_task(llm._run_with_recovery(timeout, max_retries=1))
    await asyncio.wait_for(waiting.wait(), 2)
    task.cancel()
    with pytest.raises(asyncio.CancelledError):
        await task
    assert engine.records[-1]["payload"]["phase"] == "cancelled"
    assert "模型请求" in engine.records[-1]["payload"]["error"]


def test_diagnostic_api_limits_preview_but_downloads_complete_history(tmp_path, monkeypatch):
    from fastapi import FastAPI
    from fastapi.testclient import TestClient
    from apps.admin_console.database import connection
    from apps.admin_console.routers.sessions import router
    from artemis.data_engine.storage import StorageManager

    path = tmp_path / "data.db"
    StorageManager(db_path=path, base_trace_dir=tmp_path)
    monkeypatch.setattr(connection, "DB_PATH", path)
    with connection.db_session() as conn:
        conn.execute("INSERT INTO sessions(session_id,status) VALUES ('task', 'running')")
        for i in range(105):
            value = {
                "request_id": "one-request",
                "attempt_id": str(i),
                "started_at": i,
                "phase": "completed" if i == 104 else "retry_wait",
            }
            conn.execute(
                "INSERT INTO traces(trace_id,session_id,type,timestamp,payload) VALUES (?, 'task', 'llm_diagnostic', ?, ?)",
                (str(i), i, json.dumps(value)),
            )
        conn.commit()
    app = FastAPI()
    app.include_router(router)
    with TestClient(app) as client:
        preview = client.get("/api/sessions/task/request-diagnostics")
        assert preview.status_code == 200
        assert len(preview.json()["requests"]) == 100
        assert not any(r["active"] for r in preview.json()["requests"])
        download = client.get("/api/sessions/task/request-diagnostics?download=true")
        assert download.headers["content-disposition"].startswith("attachment;")
        assert len(download.json()["requests"]) == 105
        assert download.json()["session_id"] == "task"