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

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