fix: bound model stream aggregation and close failed responses

tests/unit/test_responses_diagnostics.py

@@ -0,0 +1,263 @@
+import json
+import gzip
+from types import SimpleNamespace
+
+import httpx2
+import pytest
+
+from artemis.llm.responses_diagnostics import MAX_FRAME_BYTES, ResponsesObservation
+from artemis.llm.diagnostics import ACTIVE_ATTEMPT, observe_response_headers
+from artemis.llm.tool_constraints import MAX_REQUEST_BYTES, request_constraints, tool_constraints
+
+
+def frame(kind, sequence, **fields):
+    return (
+        "data: "
+        + json.dumps({"type": kind, "sequence_number": sequence, **fields}, ensure_ascii=False)
+        + "\r\n\r\n"
+    ).encode()
+
+
+@pytest.mark.parametrize("size", [1, 37, 100000])
+def test_raw_frames_split_unicode_and_done_snapshots_are_not_double_counted(size):
+    arguments = '{"text":"私人🙂"}'
+    data = b"".join(
+        [
+            frame("response.created", 0, response={"id": "resp_one"}),
+            frame(
+                "response.output_item.added",
+                1,
+                output_index=2,
+                item={
+                    "type": "function_call",
+                    "id": "fc_one",
+                    "call_id": "call_one",
+                    "name": "click",
+                    "arguments": "",
+                },
+            ),
+            frame(
+                "response.function_call_arguments.delta",
+                2,
+                output_index=2,
+                item_id="fc_one",
+                delta=arguments,
+            ),
+            frame(
+                "response.function_call_arguments.done",
+                3,
+                output_index=2,
+                item_id="fc_one",
+                arguments=arguments,
+            ),
+            frame(
+                "response.output_item.done",
+                4,
+                output_index=2,
+                item={"type": "function_call", "id": "fc_one", "arguments": arguments},
+            ),
+            frame(
+                "response.completed",
+                5,
+                response={
+                    "id": "resp_one",
+                    "tools": [{"type": "function", "name": "click", "strict": True}],
+                },
+            ),
+        ]
+    )
+    observer = ResponsesObservation()
+    for offset in range(0, len(data), size):
+        observer.feed(data[offset : offset + size])
+    observer.end(eof=True)
+    stats = observer.snapshot()
+    assert stats["body_bytes"] == len(data)
+    assert stats["parse_errors"] == 0
+    assert stats["incomplete_sse_frame"] is False
+    (tool,) = stats["tools"]
+    assert tool["argument_chars"] == len(arguments)
+    assert (
+        tool["delta_events"]
+        == tool["arguments_done_events"]
+        == tool["output_item_done_events"]
+        == 1
+    )
+    assert tool["call_id"] == "call_one"
+    (response,) = stats["responses"]
+    assert response["sequence_first"] == 0 and response["sequence_last"] == 5
+    assert response["sequence_forward_gaps"] == response["sequence_non_increasing"] == 0
+    assert response["returned_tool_constraints"]["tools"][0]["strict"] == "true"
+    assert "私人" not in json.dumps(stats, ensure_ascii=False)
+    assert not observer.line and not observer.lines
+
+
+def test_sequence_statistics_are_per_response_and_describe_observed_jumps():
+    observer = ResponsesObservation()
+    observer.event({"type": "response.in_progress", "sequence_number": 1})
+    observer.event({"type": "response.created", "sequence_number": 0, "response": {"id": "resp_a"}})
+    for seq in [2, 2, 1, 3, None]:
+        observer.event({"type": "response.in_progress", "sequence_number": seq})
+    observer.event({"type": "response.created", "sequence_number": 0, "response": {"id": "resp_b"}})
+    a, b = observer.snapshot()["responses"]
+    assert a["sequence_forward_gaps"] == a["sequence_forward_missing_total"] == 2
+    assert a["sequence_non_increasing"] == 2
+    assert a["sequence_equal_previous"] == 1
+    assert a["sequence_missing_events"] == 1
+    assert b["sequence_non_increasing"] == b["sequence_forward_gaps"] == 0
+    assert observer.data["unknown_response_events"] == 1
+
+
+def test_oversized_frame_recovery_and_partial_eof_keep_memory_bounded():
+    observer = ResponsesObservation()
+    observer.feed(b"data: " + b"x" * MAX_FRAME_BYTES)
+    observer.feed(b"\r\n\r\n")
+    observer.feed(b"data: invalid\n\n")
+    observer.feed(frame("response.created", 0, response={"id": "resp_valid"}))
+    observer.feed(b'data: {"private":')
+    observer.end(eof=True)
+    stats = observer.snapshot()
+    assert stats["oversized_frames"] == 1
+    assert stats["parse_errors"] == 1
+    assert len(stats["responses"]) == 1
+    assert stats["incomplete_sse_frame"] is True
+    assert "private" not in repr(vars(observer))
+    assert not observer.line and not observer.lines
+    observer.end()  # HTTP close can follow EOF; retain the original coverage gap.
+    assert observer.data["incomplete_sse_frame"] is True
+
+
+def test_raw_tool_and_response_statistics_have_fixed_caps():
+    observer = ResponsesObservation()
+    observer.event({"type": "response.created", "response": {"id": "resp_one"}})
+    for index in range(100):
+        observer.event(
+            {
+                "type": "response.function_call_arguments.delta",
+                "item_id": f"fc_{index}",
+                "output_index": index,
+                "delta": " ",
+            }
+        )
+    for index in range(20):
+        observer.event({"type": "response.created", "response": {"id": f"resp_{index}"}})
+    assert len(observer.tools) == len(observer.fingerprints) == len(observer.aliases) == 64
+    assert len(observer.responses) == 16
+    assert observer.data["statistics_truncated"] is True
+
+
+@pytest.mark.parametrize("size", [1, 37, MAX_FRAME_BYTES * 2])
+def test_oversized_crlf_frame_recovers_across_arbitrary_read_boundaries(size):
+    observer = ResponsesObservation()
+    data = b"data: " + b"x" * MAX_FRAME_BYTES + b"\r\n\r\n"
+    data += frame("response.created", 0, response={"id": "resp_after_large"})
+    for offset in range(0, len(data), size):
+        observer.feed(data[offset : offset + size])
+    assert observer.data["oversized_frames"] == 1
+    assert observer.data["parse_errors"] == 0
+    assert list(observer.responses) == ["resp_after_large"]
+    assert not observer.line and not observer.lines and not observer.dropping
+
+
+def test_alias_statistics_are_bounded_even_when_item_keeps_changing_index():
+    observer = ResponsesObservation()
+    observer.event({"response": {"id": "resp_one"}})
+    for index in range(1000):
+        observer.event(
+            {
+                "type": "response.function_call_arguments.delta",
+                "item_id": "fc_same",
+                "output_index": index,
+                "delta": " ",
+            }
+        )
+    assert len(observer.tools) == 1
+    assert len(observer.aliases) == 128
+    assert observer.data["statistics_truncated"] is True
+
+
+def test_malformed_event_type_cannot_break_passive_observation():
+    observer = ResponsesObservation()
+    observer.feed(frame([], 0, response={"id": "resp_one"}))
+    assert observer.snapshot()["responses"][0]["events"] == 1
+
+
+@pytest.mark.asyncio
+async def test_compressed_body_is_explicitly_skipped_and_sdk_decoding_is_unchanged():
+    body = frame("response.created", 0, response={"id": "resp_one"})
+    compressed = gzip.compress(body)
+
+    class Stream(httpx2.AsyncByteStream):
+        async def __aiter__(self):
+            yield compressed
+
+    attempt = SimpleNamespace(
+        wire_streams=[],
+        data={},
+        headers=lambda _: None,
+        stream_work=lambda *_: None,
+        publish_if_due=lambda: None,
+    )
+    response = httpx2.Response(
+        200,
+        headers={"content-type": "text/event-stream", "content-encoding": "gzip"},
+        stream=Stream(),
+    )
+    token = ACTIVE_ATTEMPT.set(attempt)
+    try:
+        await observe_response_headers(response)
+        assert b"".join([chunk async for chunk in response.aiter_bytes()]) == body
+        stats = attempt.wire_streams[0].observation.snapshot()
+        assert stats["sse_observation"] == "skipped_content_encoding"
+        assert stats["body_bytes"] == len(compressed)
+        assert not stats["responses"] and not stats["parse_errors"]
+    finally:
+        await response.aclose()
+        ACTIVE_ATTEMPT.reset(token)
+
+
+@pytest.mark.parametrize("strict,expected", [(None, "omitted"), (False, "false"), (True, "true")])
+def test_actual_strict_state_and_optional_schema_are_observed_without_mutation(strict, expected):
+    tool = {
+        "type": "function",
+        "name": "click",
+        "parameters": {
+            "type": "object",
+            "properties": {"target": {"anyOf": [{"type": "string"}, {"type": "null"}]}},
+            "required": [],
+        },
+    }
+    if strict is not None:
+        tool["strict"] = strict
+    body = json.dumps(
+        {"tools": [tool], "input": "private prompt", "max_output_tokens": 2048}
+    ).encode()
+    result = request_constraints(body)
+    assert result["max_output_tokens"] == 2048
+    (observed,) = result["tools"]
+    assert observed["strict"] == expected
+    assert observed["schema"]["optional_properties"] == observed["schema"]["open_objects"] == 1
+    assert "strict" not in tool if strict is None else tool["strict"] is strict
+    assert "private" not in json.dumps(result)
+    assert json.loads(body)["tools"][0] == tool
+
+
+def test_constraints_bound_large_requests_namespaces_and_deep_schemas():
+    assert request_constraints(b"x" * (MAX_REQUEST_BYTES + 1)) == {"observation": "body_too_large"}
+    assert request_constraints(b"no JSON") == {"observation": "invalid_json"}
+    tools = [
+        {
+            "type": "namespace",
+            "name": "functions",
+            "tools": [{"type": "function", "name": "click"}],
+        },
+        {"type": "function", "name": "drag", "strict": False},
+    ]
+    result = tool_constraints(tools)
+    assert {t["name"]: t["strict"] for t in result["tools"]} == {
+        "click": "omitted",
+        "drag": "false",
+    }
+    schema = {"type": "object", "properties": {str(i): {"type": "string"} for i in range(2000)}}
+    result = tool_constraints([{"type": "function", "name": "wide", "parameters": schema}])
+    assert result["tools"][0]["schema"]["truncated"] is True
+    assert result["tools"][0]["schema"]["nodes"] <= 1024