143 lines
4.2 KiB
Python
143 lines
4.2 KiB
Python
import asyncio
|
|
import logging
|
|
|
|
import pytest
|
|
|
|
from artemis.services import llm
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("replace_pause", [False, True])
|
|
async def test_cancelled_wait_clears_only_its_own_pause(tmp_path, monkeypatch, replace_pause):
|
|
pause = tmp_path / "pause"
|
|
pause.write_text("old request")
|
|
entered = asyncio.Event()
|
|
|
|
async def wait(_seconds):
|
|
entered.set()
|
|
await asyncio.Event().wait()
|
|
|
|
monkeypatch.setattr(llm.asyncio, "sleep", wait)
|
|
task = asyncio.create_task(llm._wait_for_resume(pause))
|
|
await entered.wait()
|
|
if replace_pause:
|
|
pause.write_text("a newer request owns this pause")
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
assert pause.exists() is replace_pause
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pause_deadline_clears_file(tmp_path, monkeypatch):
|
|
from unittest.mock import AsyncMock
|
|
|
|
pause = tmp_path / "pause"
|
|
pause.write_text("failed request")
|
|
monkeypatch.setattr(llm.settings, "LLM_PAUSE_TIMEOUT_SECONDS", 1)
|
|
monkeypatch.setattr(llm.asyncio, "sleep", AsyncMock())
|
|
assert await llm._wait_for_resume(pause) is False
|
|
assert not pause.exists()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_timeout_wrapper_cancels_llm_call_when_caller_is_cancelled():
|
|
started = asyncio.Event()
|
|
child_cancelled = asyncio.Event()
|
|
|
|
async def pending_llm_call():
|
|
started.set()
|
|
try:
|
|
await asyncio.Event().wait()
|
|
except asyncio.CancelledError:
|
|
child_cancelled.set()
|
|
raise
|
|
|
|
invocation = asyncio.create_task(
|
|
llm.invoke_llm_with_timeout_message(
|
|
pending_llm_call(),
|
|
timeout_seconds=60,
|
|
)
|
|
)
|
|
await started.wait()
|
|
invocation.cancel()
|
|
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await invocation
|
|
|
|
assert child_cancelled.is_set()
|
|
|
|
|
|
class _RecordingEngine:
|
|
current_step_id = "step-7"
|
|
current_session_id = "session-1"
|
|
|
|
def __init__(self):
|
|
self.trace_kwargs = None
|
|
self.published = None
|
|
|
|
def record_trace(self, **kwargs):
|
|
self.trace_kwargs = kwargs
|
|
return "trace-1"
|
|
|
|
def _publish(self, event_type, payload):
|
|
self.published = (event_type, payload)
|
|
|
|
|
|
def test_pause_error_is_persisted_and_published(monkeypatch, tmp_path):
|
|
engine = _RecordingEngine()
|
|
monkeypatch.setattr(llm, "_CURRENT_DATA_ENGINE", engine)
|
|
monkeypatch.setattr(llm.settings, "TRACES_PATH", str(tmp_path / "traces"))
|
|
|
|
pause_file = llm._handle_llm_pause_and_resume(RuntimeError("503 high demand"))
|
|
|
|
assert pause_file.read_text() == "LLM Error: 503 high demand"
|
|
assert engine.trace_kwargs == {
|
|
"type": "llm_call",
|
|
"name": "llm_pause",
|
|
"payload": {"error": "503 high demand", "pause": True},
|
|
"step_id": "step-7",
|
|
"status": "failed",
|
|
}
|
|
event_type, payload = engine.published
|
|
assert event_type == "task_paused"
|
|
assert payload["error"] == "503 high demand"
|
|
assert payload["step_id"] == "step-7"
|
|
assert isinstance(payload["timestamp"], float)
|
|
|
|
|
|
def test_terminal_failure_aggregates_observed_retries_and_total_wait(monkeypatch, tmp_path):
|
|
engine = _RecordingEngine()
|
|
monkeypatch.setattr(llm, "_CURRENT_DATA_ENGINE", engine)
|
|
monkeypatch.setattr(llm.settings, "TRACES_PATH", str(tmp_path / "traces"))
|
|
monkeypatch.setattr(llm.time, "time", lambda: 120.0)
|
|
retries = [
|
|
{
|
|
"error": "503 high demand",
|
|
"delay": 10.0,
|
|
"provider": "openai",
|
|
"source": "provider_sdk",
|
|
"request_id": "request-2",
|
|
"scheduled_at": 101.0,
|
|
"timestamp": 101.0,
|
|
}
|
|
]
|
|
request_token = llm._ACTIVE_LLM_REQUEST.set(
|
|
{
|
|
"request_id": "request-2",
|
|
"provider": "openai",
|
|
"started_at": 100.0,
|
|
"retries": retries,
|
|
}
|
|
)
|
|
try:
|
|
llm._handle_llm_pause_and_resume(RuntimeError("503 high demand"))
|
|
finally:
|
|
llm._ACTIVE_LLM_REQUEST.reset(request_token)
|
|
|
|
payload = engine.trace_kwargs["payload"]
|
|
assert payload["request_id"] == "request-2"
|
|
assert payload["provider"] == "openai"
|
|
assert payload["waited_seconds"] == 20.0
|
|
assert payload["retries"] == retries
|