nisse/nisse-core/tests/test_llm.py

import json
from pathlib import Path
from typing import Any

import httpx
import pytest

from nisse_core.envfile import load_env_file
from nisse_core.log.canonical import (
    AssistantText,
    StoredEvent,
    ToolCall,
    UserMessage,
)
from nisse_core.models.llm import (
    DEFAULT_MAX_OUTPUT_TOKENS,
    OUTAGE_BACKOFF_S,
    RATE_LIMIT_RETRY_S,
    AnthropicPassthroughClient,
    BudgetExhausted,
    LiteLLMClient,
    LLMError,
    ModelRefused,
    RoleMap,
    TruncatedResponse,
    UpstreamUnavailable,
)

EVENTS = [StoredEvent(id=1, ts="t", event=UserMessage(text="hi", source="m:1"))]


def make_client(handler, sleep=lambda _s: None) -> LiteLLMClient:
    return LiteLLMClient(
        "http://litellm.test/v1",
        "sk-test",
        transport=httpx.MockTransport(handler),
        sleep=sleep,  # never wait in unit tests
    )


def completion(message: dict) -> httpx.Response:
    return httpx.Response(
        200, json={"model": "gpt-test-1", "choices": [{"message": message}]}
    )


def test_request_shape_and_auth() -> None:
    seen: dict = {}

    def handler(request: httpx.Request) -> httpx.Response:
        seen["url"] = str(request.url)
        seen["auth"] = request.headers.get("authorization")
        seen["body"] = json.loads(request.content)
        return completion({"content": "hello"})

    make_client(handler).complete(EVENTS, tools=None, model="local/qwen", system="s")
    assert seen["url"] == "http://litellm.test/v1/chat/completions"
    assert seen["auth"] == "Bearer sk-test"
    assert seen["body"]["model"] == "local/qwen"
    assert seen["body"]["messages"][0] == {"role": "system", "content": "s"}
    assert seen["body"]["messages"][1] == {"role": "user", "content": "hi"}


def test_assistant_text_parsed_with_producing_model() -> None:
    client = make_client(lambda request: completion({"content": "hello"}))
    [event] = client.complete(EVENTS, tools=None, model="local/qwen")
    assert event == AssistantText(text="hello", model="gpt-test-1")


def test_tool_calls_get_locally_minted_canonical_ids() -> None:
    message = {
        "content": None,
        "tool_calls": [
            {
                "id": "call_provider_abc",
                "type": "function",
                "function": {
                    "name": "run_tests",
                    "arguments": '{"path": "tests/"}',
                },
            }
        ],
    }
    client = make_client(lambda request: completion(message))
    [event] = client.complete(EVENTS, tools=None, model="local/qwen")
    assert isinstance(event, ToolCall)
    assert event.tool == "run_tests"
    assert event.arguments == {"path": "tests/"}
    assert event.call_id != "call_provider_abc"  # provider ID never a join key
    assert event.call_id.startswith("call-")


def test_truncated_tool_arguments_raise_llm_error() -> None:
    message = {
        "content": None,
        "tool_calls": [
            {
                "id": "c",
                "type": "function",
                "function": {"name": "submit_plan", "arguments": '{"content": "'},
            }
        ],
    }
    client = make_client(lambda request: completion(message))
    with pytest.raises(LLMError, match="non-JSON arguments"):
        client.complete(EVENTS, tools=None, model="local/qwen")


def test_non_object_tool_arguments_raise_llm_error() -> None:
    message = {
        "content": None,
        "tool_calls": [
            {
                "id": "c",
                "type": "function",
                "function": {"name": "submit_plan", "arguments": '"just a string"'},
            }
        ],
    }
    client = make_client(lambda request: completion(message))
    with pytest.raises(LLMError, match="expected object"):
        client.complete(EVENTS, tools=None, model="local/qwen")


def test_http_error_status_raises_llm_error() -> None:
    client = make_client(lambda request: httpx.Response(500, text="boom"))
    with pytest.raises(LLMError, match="500"):
        client.complete(EVENTS, tools=None, model="local/qwen")


def test_rate_limit_is_retried_once_then_surfaced() -> None:
    # A 429 is transient. Losing a long delegation to one costs far more
    # than waiting out the window, but retrying forever would hide a real
    # ceiling, so it is retried exactly once.
    attempts: list[int] = []
    waited: list[float] = []

    def handler(request: httpx.Request) -> httpx.Response:
        attempts.append(1)
        return httpx.Response(429, text="rate limited")

    client = make_client(handler, sleep=waited.append)
    with pytest.raises(LLMError, match="429"):
        client.complete(EVENTS, tools=None, model="local/qwen")
    assert len(attempts) == 2
    assert waited == [RATE_LIMIT_RETRY_S]
    # Reported even though the call failed, so a caller can leave the wait
    # out of how long the model took.
    assert client.last_rate_limit_wait_s == RATE_LIMIT_RETRY_S


def test_rate_limit_that_clears_succeeds() -> None:
    responses = [httpx.Response(429, text="slow down"), completion({"content": "hi"})]
    client = make_client(lambda request: responses.pop(0), sleep=lambda _s: None)
    [event] = client.complete(EVENTS, tools=None, model="local/qwen")
    assert isinstance(event, AssistantText)
    assert client.last_rate_limit_wait_s == RATE_LIMIT_RETRY_S


def test_the_rate_limit_wait_is_per_call() -> None:
    responses = [
        httpx.Response(429, text="slow down"),
        completion({"content": "one"}),
        completion({"content": "two"}),
    ]
    client = make_client(lambda request: responses.pop(0), sleep=lambda _s: None)
    client.complete(EVENTS, tools=None, model="local/qwen")
    client.complete(EVENTS, tools=None, model="local/qwen")
    assert client.last_rate_limit_wait_s == 0.0


def test_connection_failure_raises_llm_error() -> None:
    def handler(request: httpx.Request) -> httpx.Response:
        raise httpx.ConnectError("refused")

    with pytest.raises(LLMError, match="request failed"):
        make_client(handler).complete(EVENTS, tools=None, model="local/qwen")


def test_malformed_response_raises_llm_error() -> None:
    client = make_client(lambda request: httpx.Response(200, json={"oops": True}))
    with pytest.raises(LLMError, match="malformed"):
        client.complete(EVENTS, tools=None, model="local/qwen")


ROLES = RoleMap(
    driver="local/qwen",
    helper="anthropic/claude-sonnet-5",
    expert="anthropic/claude-opus-5",
)


def test_role_map_resolves_and_pins() -> None:
    roles = ROLES
    assert roles.resolve("driver") == "local/qwen"
    assert roles.pinned() == {
        "driver": "local/qwen",
        "helper": "anthropic/claude-sonnet-5",
        "expert": "anthropic/claude-opus-5",
        # Unset, the classifier falls back to the expert's model (L9).
        "classifier": "anthropic/claude-opus-5",
    }


def test_a_pinned_map_from_before_the_driver_reads_its_own_roles() -> None:
    # Pinned before M2.8: the old loop's roles are there, and are ignored.
    old = {"planner": "p", "executor": "e", "driver": "d", "helper": "h"}
    roles = RoleMap.from_pinned(old, ROLES)
    assert (roles.driver, roles.helper, roles.expert) == ("d", "h", ROLES.expert)


def test_unknown_role_raises() -> None:
    with pytest.raises(LLMError, match="unknown role"):
        ROLES.resolve("boss")


def test_load_env_file(tmp_path: Path) -> None:
    envfile = tmp_path / ".env"
    envfile.write_text(
        "# comment\n\nNISSE_LLM_TOKEN=sk-abc\nNISSE_LLM_BASE_URL = http://x \n"
    )
    env: dict[str, str] = {"NISSE_LLM_TOKEN": "sk-already-set"}
    load_env_file(envfile, env)
    assert env["NISSE_LLM_TOKEN"] == "sk-already-set"  # real env wins
    assert env["NISSE_LLM_BASE_URL"] == "http://x"


def test_load_env_file_missing_is_noop(tmp_path: Path) -> None:
    env: dict[str, str] = {}
    load_env_file(tmp_path / "nope.env", env)
    assert env == {}


# --- truncation (silent until measured against Opus) ---------------------


def test_openai_route_raises_on_truncation() -> None:
    # finish_reason "length" means the answer was cut off; parsing it as
    # complete yields a tool call with required arguments simply absent.
    truncated = httpx.Response(
        200,
        json={
            "model": "gpt-test-1",
            "choices": [
                {
                    "finish_reason": "length",
                    "message": {
                        "content": None,
                        "tool_calls": [
                            {
                                "id": "c",
                                "type": "function",
                                "function": {
                                    "name": "submit_plan",
                                    "arguments": '{"type": "arch"}',
                                },
                            }
                        ],
                    },
                }
            ],
        },
    )
    client = make_client(lambda request: truncated)
    with pytest.raises(TruncatedResponse, match="output limit"):
        client.complete(EVENTS, tools=None, model="gpt-6-astra")


def test_anthropic_route_raises_on_truncation() -> None:
    # Observed for real: Opus spent 4,064 output tokens against a 4,096
    # ceiling, and a longer plan came back as submit_plan with no content.
    def handler(request: httpx.Request) -> httpx.Response:
        return httpx.Response(
            200,
            json={
                "model": "claude-opus-5",
                "stop_reason": "max_tokens",
                "content": [
                    {
                        "type": "tool_use",
                        "id": "t",
                        "name": "submit_plan",
                        "input": {"type": "arch"},
                    }
                ],
            },
        )

    client = AnthropicPassthroughClient(
        "https://litellm.test",
        "tok",
        transport=httpx.MockTransport(handler),
        sleep=lambda _s: None,
    )
    with pytest.raises(TruncatedResponse, match="output limit"):
        client.complete(EVENTS, tools=None, model="claude-opus-5")


# --- the provider's safety filter -------------------------------------------


def test_openai_route_raises_when_the_filter_declines() -> None:
    # As LiteLLM reports Opus's refusal (seen live, #164 follow-up): a 200
    # with no text, which used to read as an empty answer.
    declined = httpx.Response(
        200,
        json={
            "model": "claude-opus-5-5",
            "choices": [
                {
                    "finish_reason": "content_filter",
                    "message": {"role": "assistant", "content": None},
                }
            ],
            "usage": {"prompt_tokens": 612, "completion_tokens": 0},
        },
    )
    client = make_client(lambda request: declined)
    with pytest.raises(ModelRefused, match="safety filter"):
        client.complete(EVENTS, tools=None, model="claude-opus-5-5")


def test_anthropic_route_raises_when_the_filter_declines() -> None:
    def handler(request: httpx.Request) -> httpx.Response:
        return httpx.Response(
            200,
            json={"model": "claude-opus-5-5", "stop_reason": "refusal", "content": []},
        )

    client = AnthropicPassthroughClient(
        "https://litellm.test",
        "tok",
        transport=httpx.MockTransport(handler),
        sleep=lambda _s: None,
    )
    with pytest.raises(ModelRefused, match="declined to answer"):
        client.complete(EVENTS, tools=None, model="claude-opus-5-5")


def test_a_refusal_is_an_llm_error_so_callers_already_handle_it() -> None:
    assert issubclass(ModelRefused, LLMError)


def test_a_complete_answer_is_not_mistaken_for_truncation() -> None:
    def handler(request: httpx.Request) -> httpx.Response:
        return httpx.Response(
            200,
            json={
                "model": "claude-opus-5",
                "stop_reason": "tool_use",
                "content": [
                    {
                        "type": "tool_use",
                        "id": "t",
                        "name": "submit_plan",
                        "input": {"type": "arch", "content": "a real plan"},
                    }
                ],
            },
        )

    client = AnthropicPassthroughClient(
        "https://litellm.test", "tok", transport=httpx.MockTransport(handler)
    )
    [call] = client.complete(EVENTS, tools=None, model="claude-opus-5")
    assert call.arguments["content"] == "a real plan"


def test_output_ceiling_is_generous_by_default() -> None:
    # Output is billed by what is produced, so a low ceiling buys nothing and
    # costs silent truncation. The configured value is tested in test_config.
    assert DEFAULT_MAX_OUTPUT_TOKENS >= 8192
    assert LiteLLMClient("http://x", "t")._max_tokens == DEFAULT_MAX_OUTPUT_TOKENS


def test_truncation_is_an_llm_error_so_callers_already_handle_it() -> None:
    # The planner retries, the executor and delegation gate. None of them
    # need to learn a new exception type.
    assert issubclass(TruncatedResponse, LLMError)


def test_a_5xx_is_upstream_trouble() -> None:
    client = make_client(lambda request: httpx.Response(503, text="bad gateway"))
    with pytest.raises(UpstreamUnavailable, match="503"):
        client.complete(EVENTS, tools=None, model="local-fast")


def test_a_connection_failure_is_upstream_trouble() -> None:
    def handler(request: httpx.Request) -> httpx.Response:
        raise httpx.ConnectError("refused")

    with pytest.raises(UpstreamUnavailable, match="request failed"):
        make_client(handler).complete(EVENTS, tools=None, model="local-fast")


def test_a_persistent_rate_limit_is_upstream_trouble() -> None:
    # It cleared for neither attempt, so the caller should keep its work
    # rather than discard it.
    client = make_client(lambda request: httpx.Response(429, text="slow down"))
    with pytest.raises(UpstreamUnavailable):
        client.complete(EVENTS, tools=None, model="local-fast")


def test_a_bad_request_is_not_upstream_trouble() -> None:
    client = make_client(lambda request: httpx.Response(400, text="bad schema"))
    with pytest.raises(LLMError) as excinfo:
        client.complete(EVENTS, tools=None, model="local-fast")
    assert not isinstance(excinfo.value, UpstreamUnavailable)


def test_budget_refusal_is_a_kind_of_upstream_trouble() -> None:
    # So one catch covers the family, but it can still be told apart where
    # the difference matters — it will not clear within the hour.
    assert issubclass(BudgetExhausted, UpstreamUnavailable)


# --- tool_choice (L9: constrained decoding) --------------------------------

TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "finish",
            "description": "done",
            "parameters": {"type": "object", "properties": {}},
        },
    }
]


def capture_bodies() -> tuple[list[dict], Any]:
    bodies: list[dict] = []

    def handler(request: httpx.Request) -> httpx.Response:
        bodies.append(json.loads(request.content))
        return completion({"role": "assistant", "content": "ok"})

    return bodies, handler


def test_required_is_sent_only_when_asked_and_tools_are_offered() -> None:
    bodies, handler = capture_bodies()
    client = make_client(handler)
    client.complete(EVENTS, tools=TOOLS, model="m", require_tool_call=True)
    client.complete(EVENTS, tools=TOOLS, model="m")
    client.complete(EVENTS, tools=None, model="m", require_tool_call=True)
    assert bodies[0]["tool_choice"] == "required"
    assert "tool_choice" not in bodies[1]  # the server's default, auto
    assert "tool_choice" not in bodies[2]  # nothing to choose from


def test_the_passthrough_asks_in_anthropics_terms() -> None:
    bodies: list[dict] = []

    def handler(request: httpx.Request) -> httpx.Response:
        bodies.append(json.loads(request.content))
        return httpx.Response(
            200,
            json={
                "model": "claude-haiku-4-5",
                "stop_reason": "tool_use",
                "content": [
                    {"type": "tool_use", "id": "t", "name": "finish", "input": {}}
                ],
            },
        )

    client = AnthropicPassthroughClient(
        "https://litellm.test", "tok", transport=httpx.MockTransport(handler)
    )
    client.complete(
        EVENTS, tools=TOOLS, model="claude-haiku-4-5", require_tool_call=True
    )
    assert bodies[0]["tool_choice"] == {"type": "any"}


def test_the_answering_model_is_recorded() -> None:
    # A proxy fallback answers under another model's name; this is how a
    # caller can tell.
    client = make_client(
        lambda request: completion({"role": "assistant", "content": "hi"})
    )
    client.complete(EVENTS, tools=None, model="local-fast")
    assert client.last_model == "gpt-test-1"


def test_usage_includes_what_the_server_reused_from_its_cache() -> None:
    client = make_client(
        lambda request: httpx.Response(
            200,
            json={
                "model": "qwen",
                "choices": [{"message": {"role": "assistant", "content": "hi"}}],
                "usage": {
                    "prompt_tokens": 900,
                    "completion_tokens": 40,
                    "prompt_tokens_details": {"cached_tokens": 800},
                },
            },
        )
    )
    client.complete(EVENTS, tools=None, model="qwen")
    assert client.last_usage == {"input": 900, "output": 40, "cached": 800}

    passthrough = AnthropicPassthroughClient(
        "https://litellm.test",
        "tok",
        transport=httpx.MockTransport(
            lambda request: httpx.Response(
                200,
                json={
                    "model": "claude-opus-5-5",
                    "stop_reason": "end_turn",
                    "content": [{"type": "text", "text": "hi"}],
                    "usage": {
                        "input_tokens": 50,
                        "output_tokens": 7,
                        "cache_read_input_tokens": 30,
                    },
                },
            )
        ),
        sleep=lambda _s: None,
    )
    passthrough.complete(EVENTS, tools=None, model="claude-opus-5-5")
    assert passthrough.last_usage == {"input": 50, "output": 7, "cached": 30}


# --- outages delay, never fail (L5) --------------------------------------------


def test_a_gateway_error_is_waited_out_with_backoff() -> None:
    responses = [
        httpx.Response(502, text="bad gateway"),
        httpx.Response(503, text="unavailable"),
        completion({"content": "back"}),
    ]
    waited: list[float] = []
    client = make_client(lambda request: responses.pop(0), sleep=waited.append)
    [event] = client.complete(EVENTS, tools=None, model="local/qwen")
    assert isinstance(event, AssistantText) and event.text == "back"
    assert waited == list(OUTAGE_BACKOFF_S[:2])
    # Waiting isn't the model's time.
    assert client.last_rate_limit_wait_s == sum(OUTAGE_BACKOFF_S[:2])


def test_an_outage_that_outlasts_the_backoff_is_surfaced() -> None:
    attempts: list[int] = []
    waited: list[float] = []

    def handler(request: httpx.Request) -> httpx.Response:
        attempts.append(1)
        return httpx.Response(502, text="bad gateway")

    client = make_client(handler, sleep=waited.append)
    with pytest.raises(UpstreamUnavailable, match="502"):
        client.complete(EVENTS, tools=None, model="local/qwen")
    assert len(attempts) == len(OUTAGE_BACKOFF_S) + 1
    assert waited == list(OUTAGE_BACKOFF_S)


def test_a_dropped_connection_is_retried() -> None:
    calls: list[int] = []

    def handler(request: httpx.Request) -> httpx.Response:
        calls.append(1)
        if len(calls) == 1:
            raise httpx.ConnectError("refused")
        return completion({"content": "ok"})

    [event] = make_client(handler).complete(EVENTS, tools=None, model="local/qwen")
    assert isinstance(event, AssistantText) and len(calls) == 2


def test_a_500_or_a_timeout_is_not_retried() -> None:
    # A 500 may be the request's fault; a timeout already had its full
    # request_timeout_s.
    for failure in (
        lambda request: httpx.Response(500, text="boom"),
        lambda request: (_ for _ in ()).throw(httpx.ReadTimeout("slow")),
    ):
        calls: list[int] = []

        def handler(
            request: httpx.Request, failure=failure, calls=calls
        ) -> httpx.Response:
            calls.append(1)
            return failure(request)

        with pytest.raises(LLMError):
            make_client(handler).complete(EVENTS, tools=None, model="local/qwen")
        assert len(calls) == 1


def test_a_budget_refusal_is_not_waited_on() -> None:
    calls: list[int] = []

    def handler(request: httpx.Request) -> httpx.Response:
        calls.append(1)
        return httpx.Response(503, text="Budget has been exceeded! Current cost: 5.1")

    with pytest.raises(BudgetExhausted):
        make_client(handler).complete(EVENTS, tools=None, model="local/qwen")
    assert len(calls) == 1


# --- streamed replies on the chat route (#174) -----------------------------------


def sse(*chunks: dict, done: bool = True) -> httpx.Response:
    lines = "".join(f"data: {json.dumps(c)}\n\n" for c in chunks)
    return httpx.Response(
        200,
        headers={"content-type": "text/event-stream"},
        content=lines + ("data: [DONE]\n\n" if done else ""),
    )


def piece(delta: dict | None = None, finish: str | None = None, **extra) -> dict:
    choice: dict = {"index": 0, "delta": delta or {}}
    if finish:
        choice["finish_reason"] = finish
    return {"model": "qwen-real", "choices": [choice], **extra}


class Failing(httpx.SyncByteStream):
    """A reply that sends its first piece, then fails."""

    def __init__(self, first: bytes, failure: Exception) -> None:
        self._first, self._failure = first, failure

    def __iter__(self):
        yield self._first
        raise self._failure


def test_the_chat_route_asks_for_a_stream_with_usage() -> None:
    seen: dict = {}

    def handler(request: httpx.Request) -> httpx.Response:
        seen.update(json.loads(request.content))
        return sse(piece({"content": "hi"}, "stop"))

    make_client(handler).complete(EVENTS, tools=None, model="local/qwen")
    assert seen["stream"] is True
    assert seen["stream_options"] == {"include_usage": True}


def test_a_streamed_reply_is_put_back_together() -> None:
    usage = {"prompt_tokens": 7, "completion_tokens": 3, "cost": 0.25}
    usage["prompt_tokens_details"] = {"cached_tokens": 4}
    client = make_client(
        lambda request: sse(
            piece({"role": "assistant", "reasoning_content": "think"}),
            piece({"reasoning_content": "ing"}),
            piece({"content": "hel"}),
            piece({"content": "lo"}),
            piece(finish="stop"),
            {"model": "qwen-real", "choices": [], "usage": usage},
        )
    )
    [event] = client.complete(EVENTS, tools=None, model="local/qwen")
    assert event == AssistantText(text="hello", model="qwen-real")
    assert client.last_usage == {"input": 7, "output": 3, "cached": 4}
    assert client.last_cost == 0.25
    assert client.last_model == "qwen-real"
    # The call record keeps the reply as one response, reasoning included.
    message = client.last_exchange["response"]["choices"][0]["message"]
    assert message["reasoning_content"] == "thinking"
    assert message["content"] == "hello"


def test_streamed_tool_calls_are_joined_from_their_pieces() -> None:
    def call(index: int, **function) -> dict:
        return {"tool_calls": [{"index": index, "function": function}]}

    client = make_client(
        lambda request: sse(
            piece(call(0, name="write_file", arguments='{"path":')),
            piece(call(0, arguments=' "a.txt"}')),
            piece(call(1, name="read_file", arguments="{")),
            piece(call(1, arguments='"path": "b"}')),
            piece(finish="tool_calls"),
        )
    )
    first, second = client.complete(EVENTS, tools=None, model="local/qwen")
    assert isinstance(first, ToolCall) and isinstance(second, ToolCall)
    assert (first.tool, first.arguments) == ("write_file", {"path": "a.txt"})
    assert (second.tool, second.arguments) == ("read_file", {"path": "b"})


def test_a_streamed_reply_cut_off_at_the_limit_is_truncated() -> None:
    client = make_client(lambda request: sse(piece({"content": "half"}, "length")))
    with pytest.raises(TruncatedResponse):
        client.complete(EVENTS, tools=None, model="local/qwen")


def test_a_stream_that_stops_before_it_finishes_is_upstream_trouble() -> None:
    client = make_client(lambda request: sse(piece({"content": "half"}), done=False))
    with pytest.raises(UpstreamUnavailable, match="stopped before it finished"):
        client.complete(EVENTS, tools=None, model="local/qwen")


def test_an_error_mid_stream_is_upstream_trouble() -> None:
    broke = make_client(lambda request: sse({"error": {"message": "backend died"}}))
    with pytest.raises(UpstreamUnavailable, match="backend died"):
        broke.complete(EVENTS, tools=None, model="local/qwen")
    spent = make_client(
        lambda request: sse({"error": {"message": "Budget has been exceeded"}})
    )
    with pytest.raises(BudgetExhausted):
        spent.complete(EVENTS, tools=None, model="local/qwen")


def test_a_model_gone_quiet_mid_reply_times_out_once() -> None:
    calls: list[int] = []

    def handler(request: httpx.Request) -> httpx.Response:
        calls.append(1)
        first = f"data: {json.dumps(piece({'content': 'hel'}))}\n\n".encode()
        return httpx.Response(
            200,
            headers={"content-type": "text/event-stream"},
            stream=Failing(first, httpx.ReadTimeout("quiet")),
        )

    with pytest.raises(UpstreamUnavailable, match="said nothing for 600 s"):
        make_client(handler).complete(EVENTS, tools=None, model="local/qwen")
    assert len(calls) == 1  # a hung model is not asked again


def test_a_connection_lost_mid_reply_is_retried() -> None:
    calls: list[int] = []

    def handler(request: httpx.Request) -> httpx.Response:
        calls.append(1)
        if len(calls) == 1:
            first = f"data: {json.dumps(piece({'content': 'hel'}))}\n\n".encode()
            return httpx.Response(
                200,
                headers={"content-type": "text/event-stream"},
                stream=Failing(first, httpx.RemoteProtocolError("peer closed")),
            )
        return sse(piece({"content": "hello"}, "stop"))

    [event] = make_client(handler).complete(EVENTS, tools=None, model="local/qwen")
    assert event == AssistantText(text="hello", model="qwen-real") and len(calls) == 2


def test_the_passthrough_still_asks_for_the_whole_reply() -> None:
    seen: dict = {}

    def handler(request: httpx.Request) -> httpx.Response:
        seen.update(json.loads(request.content))
        return httpx.Response(
            200,
            json={"model": "m", "content": [{"type": "text", "text": "hi"}]},
        )

    AnthropicPassthroughClient(
        "https://litellm.test", "tok", transport=httpx.MockTransport(handler)
    ).complete(EVENTS, tools=None, model="claude")
    assert "stream" not in seen


class _Server:
    """A real HTTP server that streams a reply in pieces, ``gap`` seconds
    apart: the timeout is httpx's own, so only a real socket shows what it
    measures."""

    def __init__(self, pieces: int, gap: float) -> None:
        import threading
        from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

        class Handler(BaseHTTPRequestHandler):
            def do_POST(self) -> None:
                self.rfile.read(int(self.headers["content-length"]))
                self.send_response(200)
                self.send_header("content-type", "text/event-stream")
                self.end_headers()
                import time

                try:
                    for n in range(pieces):
                        time.sleep(gap)
                        chunk = piece({"content": str(n)})
                        self.wfile.write(f"data: {json.dumps(chunk)}\n\n".encode())
                        self.wfile.flush()
                    self.wfile.write(b"data: [DONE]\n\n")
                except (BrokenPipeError, ConnectionResetError):
                    pass  # the client gave up, as it should

            def log_message(self, *args: Any) -> None:
                pass

        self.server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
        self.url = f"http://127.0.0.1:{self.server.server_address[1]}"
        threading.Thread(target=self.server.serve_forever, daemon=True).start()

    def close(self) -> None:
        self.server.shutdown()
        self.server.server_close()


def test_a_slow_reply_that_keeps_coming_outlasts_the_timeout() -> None:
    # Six pieces 0.1 s apart take longer than the 0.35 s timeout, but no gap
    # does: a model still writing is never cut off (#174).
    server = _Server(pieces=6, gap=0.1)
    try:
        client = LiteLLMClient(server.url, "tok", timeout_s=0.35)
        [event] = client.complete(EVENTS, tools=None, model="local/qwen")
    finally:
        server.close()
    assert isinstance(event, AssistantText) and event.text == "012345"


def test_a_reply_that_goes_quiet_times_out() -> None:
    server = _Server(pieces=2, gap=1.0)
    try:
        client = LiteLLMClient(server.url, "tok", timeout_s=0.2)
        with pytest.raises(UpstreamUnavailable, match="said nothing"):
            client.complete(EVENTS, tools=None, model="local/qwen")
    finally:
        server.close()