"""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_not_awaited() assert not llm.PAUSE_FILE.exists() assert any(name == "llm_gave_up" and p["handover"] == "task_failed" for name, p in resets) 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) @pytest.mark.asyncio @pytest.mark.parametrize("fallback_kind", ["same", "failing", "working"]) @pytest.mark.parametrize("already_paused", [False, True]) async def test_invalid_output_fallback_does_not_pause_or_clear_another_request( fallback_kind, already_paused ): calls = [] if already_paused: llm.PAUSE_FILE.write_text("another request or user pause", encoding="utf-8") async def main(): calls.append("main") raise ToolArgumentOverflow("click", SMALL_ACTION_ARGUMENT_LIMIT + 1) async def fallback(): calls.append("fallback") if fallback_kind == "working": return "valid response" raise ToolArgumentOverflow("click", SMALL_ACTION_ARGUMENT_LIMIT + 2) async def run(): return await llm.with_fallback( lambda: llm._run_with_recovery(main, endpoint_key="main", max_retries=5), lambda: llm._run_with_recovery( fallback, endpoint_key="main" if fallback_kind == "same" else "fallback", max_retries=5, ), ) if fallback_kind == "working": assert await run() == "valid response" else: with pytest.raises(LLMExhaustedError) as caught: await run() assert caught.value.failure.category is FailureCategory.INVALID_OUTPUT assert calls.count("main") == 2 assert calls.count("fallback") == {"same": 0, "failing": 2, "working": 1}[fallback_kind] llm._wait_for_resume.assert_not_awaited() assert llm.PAUSE_FILE.exists() is already_paused if already_paused: assert llm.PAUSE_FILE.read_text(encoding="utf-8") == "another request or user pause"