Files
Artemis-Eapil/tests/unit/test_request_diagnostics.py
T

564 lines
22 KiB
Python

import asyncio
from copy import deepcopy
import json
import hashlib
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
from artemis.runtime.cancellation import CancellationReason
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.parametrize(
("headers", "expected"),
[
({"X-Oneapi-Request-Id": "gateway-request"}, "gateway-request"),
(
{"X-Oneapi-Request-Id": "gateway-request", "X-Request-Id": "provider-request"},
"gateway-request",
),
({"X-Request-Id": "provider-request"}, "provider-request"),
({"Request-Id": "legacy-request"}, "legacy-request"),
(
{"X-Oneapi-Request-Id": "invalid header value", "X-Request-Id": "standard-request"},
"standard-request",
),
],
)
def test_gateway_request_id_preserves_the_standard_header_without_assuming_its_source(
headers, expected
):
attempt = RequestAttempt(
Engine(), request_id="local-request", attempt=1, timeout=300, max_retries=1, metadata={}
)
attempt.headers(httpx.Response(200, headers=headers))
assert attempt.data["server_request_id"] == expected
assert attempt.data["response_x_request_id"] == headers.get("X-Request-Id")
def test_sdk_error_cannot_overwrite_the_observed_gateway_request_id():
class RequestError(Exception):
request_id = "sdk-request"
attempt = RequestAttempt(
Engine(), request_id="local-request", attempt=1, timeout=300, max_retries=1, metadata={}
)
attempt.headers(httpx.Response(400, headers={"X-Oneapi-Request-Id": "gateway-request"}))
attempt.finish("failed", error=RequestError("failed"))
assert attempt.data["server_request_id"] == "gateway-request"
assert attempt.data["error_request_id"] == "sdk-request"
@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 request.headers["x-client-attempt-id"] == str(attempt.trace_id)
assert attempt.data["client_attempt_id_matches"] is True
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"
request.headers["X-Client-Attempt-Id"] = "caller-attempt"
await observe_request_headers(request)
assert request.headers["x-client-request-id"] == "caller-id"
assert attempt.data["client_request_id_matches"] is False
assert request.headers["X-Client-Attempt-Id"] == "caller-attempt"
assert attempt.data["client_attempt_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
wire_ids = []
class StalledBody(httpx.AsyncByteStream):
async def __aiter__(self):
await asyncio.sleep(10)
yield b""
async def transport(request):
nonlocal calls
calls += 1
wire_ids.append(
(request.headers["X-Client-Request-Id"], request.headers["X-Client-Attempt-Id"])
)
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={"request": [observe_request_headers], "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["previous_attempt_id"] == failed["attempt_id"]
assert wire_ids == [
(failed["request_id"], failed["attempt_id"]),
(final["request_id"], final["attempt_id"]),
]
retry = next(r["payload"] for r in engine.records if r["name"] == "llm_retry")
assert retry["request_id"] == failed["request_id"]
assert retry["attempt_id"] == failed["attempt_id"]
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}')
for tool, name, arguments, call_id, index, count in zip(
attempt.data["tools"],
["click", "wait"],
['{"secret":"private"}', '{"seconds":1}'],
["c1", "c2"],
[0, 1],
[2, 1],
strict=True,
):
assert tool == {
"name": name,
"call_id": call_id,
"chunk_index": index,
"argument_chars": len(arguments),
"argument_bytes": len(arguments.encode()),
"argument_sha256": hashlib.sha256(arguments.encode()).hexdigest(),
"argument_deltas": count,
"argument_prefixes": [],
"argument_encoding_valid": True,
}
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 len(failed["tools"]) == 1
tool = failed["tools"][0]
assert tool["name"] == "click" and tool["argument_chars"] == tool["argument_bytes"] == 5
assert tool["argument_sha256"] == hashlib.sha256(b'{"x":').hexdigest()
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"
@pytest.mark.asyncio
async def test_diagnostic_ids_survive_streaming_and_structured_completion_without_header_bodies():
engine = Engine()
attempt = RequestAttempt(
engine, request_id="local", attempt=1, timeout=300, max_retries=1, metadata={}
)
token = ACTIVE_ATTEMPT.set(attempt)
traceparent = "00-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa-bbbbbbbbbbbbbbbb-01"
try:
request = httpx.Request(
"POST", "https://example.test/responses", headers={"x-client-request-id": "caller-id"}
)
await observe_request_headers(request)
await observe_response_headers(
httpx.Response(
200,
headers={
"x-request-id": "server-id",
"x-trace-id": "gateway-id",
"traceparent": traceparent,
"set-cookie": "private cookie",
"authorization": "private credential",
},
)
)
attempt.chunk(AIMessageChunk(content="hello", response_metadata={"id": "resp-first"}))
attempt.finish(
"completed",
response={
"raw": AIMessageChunk(
content="hello", response_metadata={"id": "resp-final", "status": "completed"}
),
"parsed": {"private": "business details"},
},
)
finally:
ACTIVE_ATTEMPT.reset(token)
assert request.headers["x-client-request-id"] == "caller-id"
assert attempt.data["client_request_id"] == "caller-id"
assert attempt.data["server_request_id"] == "server-id"
assert attempt.data["server_trace_id"] == "gateway-id"
assert attempt.data["traceparent"] == traceparent
assert attempt.data["response_id"] == "resp-final"
assert attempt.data["response_status"] == "completed"
assert "private" not in json.dumps(engine.records, default=str)
def test_invalid_or_overlong_correlation_headers_are_omitted():
attempt = RequestAttempt(
Engine(), request_id="local", attempt=1, timeout=300, max_retries=1, metadata={}
)
attempt.headers(
httpx.Response(
200,
headers={
"x-request-id": "x" * 257,
"x-trace-id": "not an id",
"traceparent": "not a trace",
},
)
)
assert attempt.data["server_request_id"] is None
assert attempt.data["server_trace_id"] is None
assert "traceparent" not in attempt.data
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"
@pytest.mark.parametrize(
"kind, key", [("llm_diagnostic", "requests"), ("final_review_diagnostic", "final_reviews")]
)
def test_snapshots_survive_out_of_order_writes_and_export(tmp_path, kind, key):
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=kind,
name="llm_request",
timestamp=10 + sequence,
status="success" if sequence == 3 else "running",
payload={
"sequence": sequence,
"phase": phase,
"status": "completed" if sequence == 3 else "running",
"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[key]) == 1
assert result[key][0]["phase"] == "completed"
assert result[key][0]["active"] is False
assert result[key][0]["timestamps_utc"]["started_at"].endswith("+00:00")
assert result["total"] == (1 if kind == "llm_diagnostic" else 0)
def test_review_diagnostics_show_evidence_stage_without_a_live_model_request(tmp_path, monkeypatch):
from artemis.data_engine.storage import StorageManager
from apps.admin_console.database.connection import db_session
from apps.admin_console.services import request_diagnostics as service
path = tmp_path / "data.db"
StorageManager(db_path=path, base_trace_dir=tmp_path)
monkeypatch.setattr(service.time, "time", lambda: 130)
with db_session(path) as conn:
conn.execute("INSERT INTO sessions(session_id,status) VALUES ('task', 'running')")
value = {
"review_id": "final-review",
"status": "running",
"stage": "reading_evidence",
"started_at": 100,
"deadline_at": 200,
"total_seconds": 100,
"updated_at": 110,
}
conn.execute(
"INSERT INTO traces(trace_id,session_id,type,timestamp,payload) "
"VALUES ('review', 'task', 'final_review_diagnostic', 100, ?)",
(json.dumps(value),),
)
conn.commit()
result = service.request_diagnostics("task", db_path=path, limit=100)
assert result["requests"] == []
review = result["final_reviews"][0]
assert review["active"] is True
assert review["stage"] == "reading_evidence"
assert review["elapsed_seconds"] == 30
assert review["remaining_seconds"] == 70
assert review["updated_at"] == 110 # preserve the actual last persisted event time
assert review["timestamps_utc"]["deadline_at"].endswith("+00:00")
assert (
service.request_diagnostics("task", db_path=path, offset=1)["final_reviews"][0]["active"]
is False
)
@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(CancellationReason.GUIDANCE_REPLACED)
with pytest.raises(asyncio.CancelledError):
await task
assert engine.records[-1]["payload"]["phase"] == "cancelled"
assert "模型请求" in engine.records[-1]["payload"]["error"]
assert engine.records[-1]["payload"]["cancel_reason"] == "guidance_replaced"
assert engine.records[-1]["payload"]["cancel_phase"] == "retry_wait"
assert "timeout_scope" not in engine.records[-1]["payload"]
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"