fix: improve cross-device evidence, request diagnostics and prompt efficiency

tests/unit/test_request_diagnostics.py

@@ -0,0 +1,340 @@
+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_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_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"