fix: bound model stream aggregation and close failed responses

tests/unit/test_responses_wire_lifecycle.py

@@ -0,0 +1,323 @@
+import asyncio
+from copy import deepcopy
+import json
+
+import httpx2 as httpx
+from langchain_core.messages import AIMessageChunk
+from langchain_core.callbacks import AsyncCallbackHandler
+from pydantic import SecretStr
+import pytest
+import pytest_asyncio
+
+from artemis.llm.diagnostics import ACTIVE_ATTEMPT
+from artemis.llm.reliability import LLMExhaustedError
+from artemis.llm.router import ModelEndpoint, ModelFactory
+from artemis.mcp.action_specs import tool_declaration
+from artemis.runtime.cancellation import CancellationReason
+from artemis.services import llm
+
+
+def event(kind, sequence, **fields):
+    return (
+        "data: " + json.dumps({"type": kind, "sequence_number": sequence, **fields}) + "\n\n"
+    ).encode()
+
+
+def initial_events():
+    return event(
+        "response.created",
+        0,
+        response={
+            "id": "resp_test",
+            "object": "response",
+            "created_at": 1,
+            "status": "in_progress",
+            "model": "test-model",
+            "output": [],
+        },
+    ) + event(
+        "response.output_item.added",
+        1,
+        output_index=0,
+        item={
+            "type": "function_call",
+            "id": "fc_test",
+            "call_id": "call_test",
+            "name": "action_sequence",
+            "arguments": "",
+            "status": "in_progress",
+        },
+    )
+
+
+class Stream(httpx.AsyncByteStream):
+    def __init__(self):
+        self.fragments = [
+            event(
+                "response.function_call_arguments.delta",
+                2,
+                output_index=0,
+                item_id="fc_test",
+                delta='{"steps":[',
+            )
+        ]
+        self.stalled = asyncio.Event()
+        self.close_calls = 0
+
+    async def __aiter__(self):
+        yield initial_events()
+        for fragment in self.fragments:
+            yield fragment
+        self.stalled.set()
+        await asyncio.Event().wait()
+
+    async def aclose(self):
+        self.close_calls += 1
+
+
+class Engine:
+    current_session_id = "wire-lifecycle"
+    current_step_id = None
+
+    def __init__(self):
+        self.records = []
+        self.retrying = asyncio.Event()
+
+    def record_trace(self, **record):
+        self.records.append(deepcopy(record))
+        if record.get("payload", {}).get("phase") == "retry_wait":
+            self.retrying.set()
+
+    def _publish(self, *_):
+        pass
+
+    @property
+    def data(self):
+        return next(r["payload"] for r in reversed(self.records) if r["type"] == "llm_diagnostic")
+
+
+@pytest_asyncio.fixture
+async def case(monkeypatch, tmp_path):
+    import openai
+    from langchain_openai import ChatOpenAI
+
+    stream, engine, clients, requests = Stream(), Engine(), [], []
+
+    def handle(request):
+        requests.append(json.loads(request.content))
+        return httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=stream)
+
+    def client(**kwargs):
+        value = httpx.AsyncClient(transport=httpx.MockTransport(handle), **kwargs)
+        clients.append(value)
+        return value
+
+    monkeypatch.setattr(openai, "DefaultAsyncHttpxClient", client)
+    monkeypatch.setattr(llm.settings, "OPENAI_API_KEY", SecretStr("test-secret"))
+    monkeypatch.setattr(llm.settings, "OPENAI_BASE_URL", "https://responses.example/v1")
+    monkeypatch.delenv("ARTEMIS_FAKE_LLM", raising=False)
+    monkeypatch.setattr(llm, "_get_current_data_engine", lambda: engine)
+    monkeypatch.setattr(llm, "PAUSE_FILE", tmp_path / "pause")
+    llm._ENDPOINT_BREAKER._states.clear()
+    llm._NON_STREAMING_ENDPOINTS.clear()
+    fallback = llm._FALLBACK_AVAILABLE.set(True)
+
+    def wrapper(timeout=30, max_retries=0):
+        endpoint = ModelEndpoint(
+            model_name="test-model", timeout_seconds=timeout, max_retries=max_retries
+        )
+        model = ModelFactory.create_model(endpoint)
+        assert isinstance(model, ChatOpenAI)
+        return llm.RobustChatModelWrapper(model, endpoint=endpoint).bind_tools(
+            [tool_declaration(n) for n in ("click", "action_sequence", "drag")]
+        )
+
+    yield stream, engine, requests, wrapper
+    llm._FALLBACK_AVAILABLE.reset(fallback)
+    for value in clients:
+        await value.aclose()
+    llm._ENDPOINT_BREAKER._states.clear()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("reason", ["outer_timeout", "request_timeout", "user_stop"])
+async def test_unfinished_arguments_below_length_limit_close_on_all_deadlines(case, reason):
+    stream, engine, requests, wrapper = case
+    if reason == "outer_timeout":
+        running = asyncio.create_task(
+            llm.invoke_llm_with_timeout_message(
+                wrapper().complete([]), timeout_seconds=0.1, hard_timeout=0.1
+            )
+        )
+    elif reason == "request_timeout":
+        running = asyncio.create_task(wrapper(0.1).complete([]))
+    else:
+        running = asyncio.create_task(wrapper().complete([]))
+        await asyncio.wait_for(stream.stalled.wait(), 2)
+        running.cancel(CancellationReason.USER_STOP)
+    with pytest.raises((TimeoutError, LLMExhaustedError, asyncio.CancelledError)):
+        await asyncio.wait_for(running, 3)
+    data = engine.data
+    assert data["tool_argument_chars"] == len('{"steps":[') < 16384
+    assert (
+        data["failure_phase"] == "generating_tool_arguments"
+        if reason == "request_timeout"
+        else data["cancel_phase"] == "generating_tool_arguments"
+    )
+    assert (
+        data["timeout_scope"] == ("outer" if reason == "outer_timeout" else "request")
+        if reason != "user_stop"
+        else data["cancel_reason"] == "user_stop"
+    )
+    assert data["cancel_observed_at"] <= data["ended_at"]
+    (wire,) = data["wire_responses"]
+    assert wire["closed_at"] >= wire["close_started_at"]
+    assert wire["body_eof_observed"] is False
+    assert wire["tools"][0].get("arguments_done_events", 0) == 0
+    assert stream.close_calls == len(requests) == 1
+    assert ACTIVE_ATTEMPT.get() is None
+    assert "test-secret" not in repr(data)
+    assert '{"steps":[' not in repr(data)
+    assert not llm.PAUSE_FILE.exists()
+
+
+@pytest.mark.asyncio
+async def test_overflow_closes_transport_even_while_consumer_is_between_sdk_yields(case):
+    stream, engine, requests, wrapper = case
+    stream.fragments = [
+        event(
+            "response.function_call_arguments.delta",
+            2,
+            output_index=0,
+            item_id="fc_test",
+            delta=" " * 17000,
+        )
+    ]
+    with pytest.raises(LLMExhaustedError):
+        await wrapper().complete([])
+    assert engine.data["stream_guard"]["tool"] == "action_sequence"
+    assert engine.data["wire_responses"][0]["closed_at"]
+    assert stream.close_calls == len(requests) == 1
+    assert not stream.stalled.is_set()
+
+
+@pytest.mark.asyncio
+async def test_overflow_closes_transport_before_waiting_for_retry(case):
+    stream, engine, requests, wrapper = case
+    stream.fragments = [
+        event(
+            "response.function_call_arguments.delta",
+            2,
+            output_index=0,
+            item_id="fc_test",
+            delta=" " * 17000,
+        )
+    ]
+    running = asyncio.create_task(wrapper(max_retries=1).complete([]))
+    await asyncio.wait_for(engine.retrying.wait(), 2)
+    assert stream.close_calls == len(requests) == 1
+    assert engine.data["wire_responses"][0]["closed_at"]
+    running.cancel(CancellationReason.USER_STOP)
+    with pytest.raises(asyncio.CancelledError):
+        await running
+    assert engine.data["cancel_phase"] == "retry_wait"
+
+
+@pytest.mark.asyncio
+async def test_actual_native_schemas_do_not_silently_force_upstream_strict(case):
+    stream, engine, requests, wrapper = case
+    running = asyncio.create_task(wrapper().complete([]))
+    await asyncio.wait_for(stream.stalled.wait(), 2)
+    running.cancel(CancellationReason.USER_STOP)
+    with pytest.raises(asyncio.CancelledError):
+        await running
+    constraints = engine.data["wire_request_constraints"]
+    assert constraints["max_output_tokens"] is None
+    assert {t["name"]: t["strict"] for t in constraints["tools"]} == {
+        "click": "omitted",
+        "action_sequence": "omitted",
+        "drag": "omitted",
+    }
+    assert any(t["schema"]["optional_properties"] for t in constraints["tools"])
+    assert all("strict" not in tool for tool in requests[0]["tools"])
+
+
+@pytest.mark.asyncio
+async def test_live_deltas_are_emitted_before_a_batch_is_flushed(monkeypatch):
+    emitted, resume = asyncio.Event(), asyncio.Event()
+    events = []
+
+    class Model:
+        async def astream(self, *_args, **_kwargs):
+            yield AIMessageChunk(content="first live text")
+            await resume.wait()
+
+    wrapper = llm.RobustChatModelWrapper(Model(), endpoint=ModelEndpoint(model_name="test"))
+    monkeypatch.setattr(wrapper, "_emit_stream_delta", lambda *_: emitted.set())
+    monkeypatch.setattr(llm, "_record_llm_event", lambda *args: events.append(args))
+    running = asyncio.create_task(wrapper._complete_attempt(True, []))
+    await asyncio.wait_for(emitted.wait(), 2)
+    assert not running.done()
+    running.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await running
+    assert any(name == "llm_stream_reset" for name, *_ in events)
+
+
+@pytest.mark.asyncio
+async def test_callback_backpressure_can_leave_wire_ahead_of_sdk_without_prefetch(case):
+    stream, engine, requests, wrapper = case
+    blocked = asyncio.Event()
+    stream.fragments = [
+        b"".join(
+            event(
+                "response.function_call_arguments.delta",
+                n + 2,
+                output_index=0,
+                item_id="fc_test",
+                delta=" ",
+            )
+            for n in range(200)
+        )
+    ]
+
+    class Callback(AsyncCallbackHandler):
+        async def on_llm_new_token(self, token, *, chunk=None, **kwargs):
+            if chunk is not None and any(t.get("args") for t in chunk.message.tool_call_chunks):
+                attempt = ACTIVE_ATTEMPT.get()
+                attempt.publish()
+                blocked.set()
+                await asyncio.Event().wait()
+
+    running = asyncio.create_task(wrapper().complete([], config={"callbacks": [Callback()]}))
+    await asyncio.wait_for(blocked.wait(), 2)
+    wire = engine.data["wire_responses"][0]
+    assert wire["tools"][0]["argument_chars"] == 200
+    assert engine.data["tool_argument_chars"] == 0
+    assert not stream.stalled.is_set()  # no second transport read or background reader
+    running.cancel(CancellationReason.USER_STOP)
+    with pytest.raises(asyncio.CancelledError):
+        await asyncio.wait_for(running, 2)
+    assert stream.close_calls == len(requests) == 1
+    assert engine.data["wire_responses"][0]["closed_at"]
+
+
+@pytest.mark.asyncio
+async def test_unresponsive_close_keeps_user_stop_and_records_cleanup_failure(case, monkeypatch):
+    stream, engine, requests, wrapper = case
+
+    async def stuck_close():
+        stream.close_calls += 1
+        await asyncio.Event().wait()
+
+    monkeypatch.setattr(stream, "aclose", stuck_close)
+    running = asyncio.create_task(wrapper().complete([]))
+    await asyncio.wait_for(stream.stalled.wait(), 2)
+    running.cancel(CancellationReason.USER_STOP)
+    with pytest.raises(asyncio.CancelledError):
+        await asyncio.wait_for(running, 3)
+    assert engine.data["cancel_reason"] == "user_stop"
+    wire = engine.data["wire_responses"][0]
+    assert wire["close_error_type"] == "TimeoutError"
+    assert "closed_at" not in wire
+    assert stream.close_calls == len(requests) == 1