fix: bound model stream aggregation and close failed responses

tests/unit/test_chunk_accumulator.py

@@ -0,0 +1,83 @@
+import itertools
+import json
+
+from langchain_core.messages import AIMessageChunk
+import langchain_core.messages.ai as ai
+import pytest
+
+from artemis.llm.chunk_accumulator import ChunkAccumulator
+from artemis.llm.response_guard import ResponseIntegrityError, ensure_response_complete
+
+
+def test_batching_preserves_interleaved_tools_unicode_metadata_and_usage():
+    arguments = [json.dumps({"target": "中文🙂" * 80}), json.dumps({"target": 3})]
+    chunks = [AIMessageChunk(content="Live text", id="resp_one")]
+    for index in range(2):
+        chunks.append(
+            AIMessageChunk(
+                content=[],
+                tool_call_chunks=[
+                    {"name": "click", "args": "", "id": f"call_{index}", "index": index}
+                ],
+            )
+        )
+    for pair in itertools.zip_longest(*arguments):
+        for index, value in enumerate(pair):
+            if value is not None:
+                chunks.append(
+                    AIMessageChunk(content=[], tool_call_chunks=[{"args": value, "index": index}])
+                )
+    chunks.append(
+        AIMessageChunk(
+            content=[],
+            response_metadata={"id": "resp_one", "status": "completed"},
+            usage_metadata={"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
+            chunk_position="last",
+        )
+    )
+    expected = chunks[0]
+    accumulator = ChunkAccumulator()
+    max_pending = 0
+    for index, chunk in enumerate(chunks):
+        if index:
+            expected = expected + chunk
+        accumulator.append(chunk)
+        if accumulator.ready:
+            accumulator.flush()
+        max_pending = max(max_pending, len(accumulator.pending))
+    result = accumulator.flush()
+    assert result.model_dump() == expected.model_dump()
+    assert [t["args"] for t in result.tool_calls] == [json.loads(a) for a in arguments]
+    assert result.usage_metadata["total_tokens"] == 15
+    assert max_pending < 256
+    assert not accumulator.pending
+
+
+def test_partial_json_is_not_reparsed_for_every_single_character(monkeypatch):
+    chunks = [
+        AIMessageChunk(
+            content=[], tool_call_chunks=[{"name": "click", "args": '{"target":x', "index": 0}]
+        )
+    ]
+    chunks.extend(
+        AIMessageChunk(content=[], tool_call_chunks=[{"args": "x", "index": 0}]) for _ in range(800)
+    )
+    original = ai.parse_partial_json
+    parsed_lengths = []
+
+    def counted(arguments, *args, **kwargs):
+        parsed_lengths.append(len(arguments))
+        return original(arguments, *args, **kwargs)
+
+    monkeypatch.setattr(ai, "parse_partial_json", counted)
+    accumulator = ChunkAccumulator()
+    for chunk in chunks:
+        accumulator.append(chunk)
+        if accumulator.ready:
+            accumulator.flush()
+    result = accumulator.flush()
+    assert result.tool_call_chunks[0]["args"] == '{"target":x' + "x" * 800
+    assert len(parsed_lengths) < 20
+    result.response_metadata["status"] = "completed"
+    with pytest.raises(ResponseIntegrityError):
+        ensure_response_complete(result)