fix(llm): isolate invalid output failures and trace argument streams

tests/unit/test_argument_fingerprint.py

@@ -0,0 +1,59 @@
+import hashlib
+import json
+
+import pytest
+from langchain_core.messages import AIMessageChunk
+
+from artemis.llm.argument_fingerprint import ArgumentFingerprint, MAX_PREFIXES, PREFIX_BYTES
+from artemis.llm.diagnostics import RequestAttempt
+
+
+@pytest.mark.parametrize("fragment_size", [1, 37, 4096, 100000])
+def test_fingerprints_match_at_byte_offsets_despite_unicode_and_different_fragmentation(
+    fragment_size,
+):
+    arguments = '{"target":"' + "私人内容🙂" * 8000 + '"}'
+    expected = arguments.encode("utf-8")
+    fingerprint = ArgumentFingerprint()
+    for offset in range(0, len(arguments), fragment_size):
+        fingerprint.observe(arguments[offset : offset + fragment_size])
+    result = fingerprint.snapshot()
+    assert result["argument_chars"] == len(arguments)
+    assert result["argument_bytes"] == len(expected)
+    assert result["argument_sha256"] == hashlib.sha256(expected).hexdigest()
+    assert result["argument_encoding_valid"] is True
+    assert result["argument_prefixes"] == [
+        {"bytes": size, "sha256": hashlib.sha256(expected[:size]).hexdigest()}
+        for size in range(PREFIX_BYTES, PREFIX_BYTES * MAX_PREFIXES + 1, PREFIX_BYTES)
+    ]
+    assert "私人" not in json.dumps(result, ensure_ascii=False)
+
+
+def test_invalid_utf8_is_marked_without_failing_or_forging_a_digest():
+    fingerprint = ArgumentFingerprint()
+    for value in ["safe", "\ud800", "later"]:
+        fingerprint.observe(value)
+    assert fingerprint.snapshot() == {
+        "argument_chars": 10,
+        "argument_bytes": None,
+        "argument_sha256": None,
+        "argument_deltas": 3,
+        "argument_prefixes": [],
+        "argument_encoding_valid": False,
+    }
+
+
+def test_tool_tracking_is_bounded_but_total_count_keeps_all_received_arguments():
+    attempt = RequestAttempt(
+        None, request_id="bounded", attempt=1, timeout=300, max_retries=0, metadata={}
+    )
+    for index in range(65):
+        attempt.chunk(
+            AIMessageChunk(
+                content=[],
+                tool_call_chunks=[{"name": "click", "args": "{}", "index": index}],
+            )
+        )
+    assert len(attempt.data["tools"]) == len(attempt.tool_fingerprints) == 64
+    assert attempt.data["tool_stats_truncated"] is True
+    assert attempt.data["tool_argument_chars"] == 130

tests/unit/test_llm_response_integrity.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import json
+import hashlib
 from types import SimpleNamespace
 from unittest.mock import AsyncMock
 
@@ -15,6 +16,7 @@ import pytest_asyncio
 from artemis.llm.reliability import FailureCategory, LLMExhaustedError
 from artemis.llm.response_guard import ResponseIntegrityError, ensure_response_complete
 from artemis.llm.router import ModelEndpoint, ModelFactory
+from artemis.llm.diagnostics import RequestAttempt
 from artemis.services import llm as llm_service
 from artemis.services.llm import RobustChatModelWrapper, acomplete
 
@@ -195,6 +197,33 @@ async def test_completed_stream_returns_original_valid_arguments(transport_case)
     assert result.tool_calls[0]["args"] == {"passed": True}
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("transport_case", ["fragmented"], indirect=True)
+async def test_real_sdk_fingerprint_counts_received_arguments_once(transport_case):
+    case = transport_case
+    arguments = json.dumps({"passed": True, "evidence": "中文🙂" * 1000}, ensure_ascii=False)
+    case.outputs.append(("completed", arguments))
+    observation = RequestAttempt(
+        None, request_id="sdk-test", attempt=1, timeout=300, max_retries=0, metadata={}
+    )
+    merged = None
+    async for chunk in case.model.astream([]):
+        observation.chunk(chunk)
+        merged = chunk if merged is None else merged + chunk
+    observation.finish("completed", response=merged)
+    assert len(observation.data["tools"]) == 1
+    tool = observation.data["tools"][0]
+    assert tool["call_id"] == "call_test"
+    assert tool["argument_chars"] == len(arguments)
+    assert tool["argument_bytes"] == len(arguments.encode())
+    assert tool["argument_sha256"] == hashlib.sha256(arguments.encode()).hexdigest()
+    assert tool["argument_prefixes"][0] == {
+        "bytes": 4096,
+        "sha256": hashlib.sha256(arguments.encode()[:4096]).hexdigest(),
+    }
+    assert tool["argument_deltas"] == 1
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("transport_case", ["fragmented"], indirect=True)
 @pytest.mark.parametrize("status", ["completed", None])

tests/unit/test_llm_stream_guard.py

@@ -124,7 +124,9 @@ async def test_gateway_closes_bad_stream_discards_partial_and_bounds_retry(
         with pytest.raises(LLMExhaustedError) as caught:
             await wrapper.complete([])
         assert caught.value.failure.category is FailureCategory.INVALID_OUTPUT
-        llm._wait_for_resume.assert_awaited_once()
+        llm._wait_for_resume.assert_not_awaited()
+        assert not llm.PAUSE_FILE.exists()
+        assert any(name == "llm_gave_up" and p["handover"] == "task_failed" for name, p in resets)
     assert model.calls == 2
     assert closed == [1, 2]
     failure = next(
@@ -139,3 +141,47 @@ async def test_gateway_closes_bad_stream_discards_partial_and_bounds_retry(
     assert failure["usage"] is None
     assert "private" not in json.dumps(records + resets, default=str)
     assert any(name == "llm_stream_reset" and p["action"] == "discard" for name, p in resets)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("fallback_kind", ["same", "failing", "working"])
+@pytest.mark.parametrize("already_paused", [False, True])
+async def test_invalid_output_fallback_does_not_pause_or_clear_another_request(
+    fallback_kind, already_paused
+):
+    calls = []
+    if already_paused:
+        llm.PAUSE_FILE.write_text("another request or user pause", encoding="utf-8")
+
+    async def main():
+        calls.append("main")
+        raise ToolArgumentOverflow("click", SMALL_ACTION_ARGUMENT_LIMIT + 1)
+
+    async def fallback():
+        calls.append("fallback")
+        if fallback_kind == "working":
+            return "valid response"
+        raise ToolArgumentOverflow("click", SMALL_ACTION_ARGUMENT_LIMIT + 2)
+
+    async def run():
+        return await llm.with_fallback(
+            lambda: llm._run_with_recovery(main, endpoint_key="main", max_retries=5),
+            lambda: llm._run_with_recovery(
+                fallback,
+                endpoint_key="main" if fallback_kind == "same" else "fallback",
+                max_retries=5,
+            ),
+        )
+
+    if fallback_kind == "working":
+        assert await run() == "valid response"
+    else:
+        with pytest.raises(LLMExhaustedError) as caught:
+            await run()
+        assert caught.value.failure.category is FailureCategory.INVALID_OUTPUT
+    assert calls.count("main") == 2
+    assert calls.count("fallback") == {"same": 0, "failing": 2, "working": 1}[fallback_kind]
+    llm._wait_for_resume.assert_not_awaited()
+    assert llm.PAUSE_FILE.exists() is already_paused
+    if already_paused:
+        assert llm.PAUSE_FILE.read_text(encoding="utf-8") == "another request or user pause"

tests/unit/test_request_diagnostics.py

@@ -1,6 +1,7 @@
 import asyncio
 from copy import deepcopy
 import json
+import hashlib
 from types import SimpleNamespace
 from uuid import uuid4
 
@@ -36,6 +37,46 @@ class Engine:
         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()
@@ -55,13 +96,18 @@ async def test_request_hook_correlates_api_calls_without_logging_header_values()
         )
         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)
 
@@ -75,6 +121,7 @@ async def test_headers_without_body_timeout_then_retry_keeps_evidence(monkeypatc
     )
     llm._ENDPOINT_BREAKER._states.clear()
     calls = 0
+    wire_ids = []
 
     class StalledBody(httpx.AsyncByteStream):
         async def __aiter__(self):
@@ -84,6 +131,9 @@ async def test_headers_without_body_timeout_then_retry_keeps_evidence(monkeypatc
     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")
@@ -91,7 +141,7 @@ async def test_headers_without_body_timeout_then_retry_keeps_evidence(monkeypatc
 
     async with httpx.AsyncClient(
         transport=httpx.MockTransport(transport),
-        event_hooks={"response": [observe_response_headers]},
+        event_hooks={"request": [observe_request_headers], "response": [observe_response_headers]},
     ) as client:
 
         class Model:
@@ -123,6 +173,10 @@ async def test_headers_without_body_timeout_then_retry_keeps_evidence(monkeypatc
     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"]
@@ -163,10 +217,26 @@ def test_tool_deltas_are_counted_once_without_retaining_arguments():
     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}')},
-    ]
+    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)
@@ -205,7 +275,10 @@ async def test_timeout_during_tool_generation_keeps_last_output_and_tool_name(mo
     )
     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 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