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"