Artemis-Eapil/tests/unit/test_responses_wire_lifecycle.py

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