Artemis-Eapil/tests/unit/test_llm_stream_guard.py

"""Runaway click arguments must never reach the executor or endless retries."""

from copy import deepcopy
import json
from types import SimpleNamespace
from unittest.mock import AsyncMock

from langchain_core.messages import AIMessageChunk
import pytest

from artemis.llm.reliability import FailureCategory, LLMExhaustedError, retry_policy_for
from artemis.llm.router import ModelEndpoint
from artemis.llm.stream_guard import (
    SMALL_ACTION_ARGUMENT_LIMIT,
    ToolArgumentGuard,
    ToolArgumentOverflow,
)
from artemis.services import llm


def delta(arguments, *, name=None, index=0):
    return AIMessageChunk(
        content=[],
        tool_call_chunks=[{"name": name, "args": arguments, "index": index}],
    )


@pytest.fixture(autouse=True)
def isolate(monkeypatch, tmp_path):
    llm._ENDPOINT_BREAKER._states.clear()
    llm._NON_STREAMING_ENDPOINTS.clear()
    monkeypatch.setattr(llm, "PAUSE_FILE", tmp_path / "pause")
    monkeypatch.setattr(
        llm, "retry_policy_for", lambda _: SimpleNamespace(max_attempts=2, delay_for=lambda _: 0)
    )
    monkeypatch.setattr(llm, "_wait_for_resume", AsyncMock(return_value=False))
    monkeypatch.setattr(llm, "_CURRENT_DATA_ENGINE", None)


@pytest.mark.parametrize("name", ["save_note", "input_text", "report_task_status", "plugin_tool"])
def test_free_text_tools_and_regular_text_are_not_limited(name):
    guard = ToolArgumentGuard()
    text = "报告内容" * SMALL_ACTION_ARGUMENT_LIMIT
    guard.observe(AIMessageChunk(content=text))
    guard.observe(delta(json.dumps({"text": text}, ensure_ascii=False), name=name))


def test_split_and_interleaved_arguments_are_counted_per_call_without_storing_body():
    guard = ToolArgumentGuard()
    for index in (0, 1):
        guard.observe(delta('{"target_description":"', name="click", index=index))
    for index in (0, 1):
        guard.observe(delta("private" * 1200, index=index))
    with pytest.raises(ToolArgumentOverflow) as caught:
        guard.observe(delta("private" * 1200, index=1))
    assert caught.value.details["tool"] == "click"
    assert caught.value.details["argument_chars"] > SMALL_ACTION_ARGUMENT_LIMIT
    assert "private" not in str(caught.value)
    assert "private" not in repr(guard.calls)


def test_name_arriving_after_arguments_still_triggers_guard():
    guard = ToolArgumentGuard()
    guard.observe(delta("x" * (SMALL_ACTION_ARGUMENT_LIMIT + 1)))
    with pytest.raises(ToolArgumentOverflow):
        guard.observe(delta("", name="click"))


def test_invalid_output_retry_policy_is_bounded_and_honors_zero_retries(monkeypatch):
    monkeypatch.setattr(llm.settings, "OPENAI_MAX_RETRIES", 5)
    assert retry_policy_for(FailureCategory.INVALID_OUTPUT).max_attempts == 2
    monkeypatch.setattr(llm.settings, "OPENAI_MAX_RETRIES", 0)
    assert retry_policy_for(FailureCategory.INVALID_OUTPUT).max_attempts == 1


@pytest.mark.asyncio
@pytest.mark.parametrize("recover", [False, True])
@pytest.mark.parametrize("first_chunk_overflows", [False, True])
async def test_gateway_closes_bad_stream_discards_partial_and_bounds_retry(
    monkeypatch, recover, first_chunk_overflows
):
    records, resets, closed = [], [], []

    class Engine:
        current_session_id = "test-session"
        current_step_id = None

        def record_trace(self, **record):
            records.append(deepcopy(record))

        def _publish(self, *_):
            pass

    class Model:
        calls = 0

        async def astream(self, *_args, **_kwargs):
            self.calls += 1
            try:
                if recover and self.calls == 2:
                    yield delta('{"target":3}', name="click")
                    return
                if not first_chunk_overflows:
                    yield delta('{"target_description":"', name="click")
                yield delta(
                    "private" * (SMALL_ACTION_ARGUMENT_LIMIT // 7 + 1),
                    name="click" if first_chunk_overflows else None,
                )
                pytest.fail("Runaway stream was consumed after the safety bound")
            finally:
                closed.append(self.calls)

    monkeypatch.setattr(llm, "_CURRENT_DATA_ENGINE", Engine())
    monkeypatch.setattr(llm, "_record_llm_event", lambda *a, **_: resets.append(a))
    model = Model()
    wrapper = llm.RobustChatModelWrapper(model, endpoint=ModelEndpoint(max_retries=5))
    if recover:
        response = await wrapper.complete([])
        assert response.tool_calls == [
            {"name": "click", "args": {"target": 3}, "id": None, "type": "tool_call"}
        ]
        llm._wait_for_resume.assert_not_awaited()
    else:
        with pytest.raises(LLMExhaustedError) as caught:
            await wrapper.complete([])
        assert caught.value.failure.category is FailureCategory.INVALID_OUTPUT
        llm._wait_for_resume.assert_awaited_once()
    assert model.calls == 2
    assert closed == [1, 2]
    failure = next(
        r["payload"]
        for r in records
        if r["type"] == "llm_diagnostic" and r["payload"]["phase"] == "failed"
    )
    assert failure["stream_guard"]["tool"] == "click"
    assert failure["stream_guard"]["argument_chars"] > SMALL_ACTION_ARGUMENT_LIMIT
    assert failure["failure_phase"] == "generating_tool_arguments"
    assert failure["category"] == "invalid_output"
    assert failure["usage"] is None
    assert "private" not in json.dumps(records + resets, default=str)
    assert any(name == "llm_stream_reset" and p["action"] == "discard" for name, p in resets)