plyngent/tests/test_cli/test_retry.py

from __future__ import annotations

from typing import TYPE_CHECKING, Literal, overload

import pytest

from plyngent.agent import ChatAgent
from plyngent.cli.retry import (
    DEFAULT_MAX_AUTO_RETRIES,
    DEFAULT_RETRY_DELAYS_SECONDS,
    default_retry_delays,
    retry_pending_with_retries,
    run_turn_with_retries,
    sleep_cancellable,
)
from plyngent.config.models import DatabaseConfig
from plyngent.lmproto.openai_compatible.model import (
    AssistantChatMessage,
    ChatCompletionChoice,
    ChatCompletionChunk,
    ChatCompletionResponse,
    ChatCompletionsParam,
    ChunkChoice,
    DeltaMessage,
    UserChatMessage,
)
from plyngent.memory import MemoryStore

if TYPE_CHECKING:
    from collections.abc import AsyncIterator


def test_default_retry_delays_schedule() -> None:
    assert DEFAULT_MAX_AUTO_RETRIES == 10
    assert default_retry_delays(0) == ()
    assert default_retry_delays(4) == (5.0, 10.0, 15.0, 20.0)
    delays = default_retry_delays(10)
    assert delays[:4] == (5.0, 10.0, 15.0, 20.0)
    assert delays[4:] == (30.0, 40.0, 50.0, 60.0, 70.0, 80.0)
    assert delays == DEFAULT_RETRY_DELAYS_SECONDS
    assert len(DEFAULT_RETRY_DELAYS_SECONDS) == 10


def _response(message: AssistantChatMessage) -> ChatCompletionResponse:
    return ChatCompletionResponse(
        id="1",
        object="chat.completion",
        created=0,
        model="test",
        choices=[
            ChatCompletionChoice(
                index=0,
                message=message,
                logprobs={},
                finish_reason="stop",
            )
        ],
        system_fingerprint="",
        usage={},
    )


async def _as_stream(response: ChatCompletionResponse) -> AsyncIterator[ChatCompletionChunk]:
    message = response.choices[0].message
    content = message.content if isinstance(message.content, str) else ""
    yield ChatCompletionChunk(
        id=response.id,
        object="chat.completion.chunk",
        created=response.created,
        model=response.model,
        choices=[
            ChunkChoice(
                index=0,
                delta=DeltaMessage(content=content),
                finish_reason="stop",
            )
        ],
    )


class FlakyClient:
    calls: int
    fail_times: int

    def __init__(self, fail_times: int) -> None:
        self.calls = 0
        self.fail_times = fail_times

    @overload
    async def chat_completions(
        self, param: ChatCompletionsParam, *, stream: Literal[False] = False
    ) -> ChatCompletionResponse: ...

    @overload
    async def chat_completions(
        self, param: ChatCompletionsParam, *, stream: Literal[True]
    ) -> AsyncIterator[ChatCompletionChunk]: ...

    async def chat_completions(
        self, param: ChatCompletionsParam, *, stream: bool = False
    ) -> ChatCompletionResponse | AsyncIterator[ChatCompletionChunk]:
        del param
        self.calls += 1
        if self.calls <= self.fail_times:
            msg = f"fail-{self.calls}"
            raise RuntimeError(msg)
        response = _response(AssistantChatMessage(content="ok"))
        if stream:
            return _as_stream(response)
        return response


async def test_sleep_cancellable_completes() -> None:
    assert await sleep_cancellable(0.01) is True


async def test_auto_retry_eventually_succeeds(monkeypatch: pytest.MonkeyPatch) -> None:
    async def no_wait(_seconds: float) -> bool:
        return True

    monkeypatch.setattr("plyngent.cli.retry.sleep_cancellable", no_wait)
    store = await MemoryStore.open(DatabaseConfig())
    session = await store.create_session(name="t")
    client = FlakyClient(fail_times=2)
    agent = ChatAgent(client, model="m", memory=store, session_id=session.sid)

    ok = await run_turn_with_retries(
        agent,
        starter=lambda: agent.run("hello"),
        delays=(0.01, 0.01, 0.01),
    )
    assert ok is True
    assert client.calls == 3
    loaded = await store.list_messages(session.sid)
    assert len(loaded) == 2
    assert isinstance(loaded[0], UserChatMessage)
    assert loaded[0].content == "hello"
    await store.close()


async def test_run_cancellable_cancels_task() -> None:
    import asyncio

    from plyngent.cli.retry import run_cancellable

    started = asyncio.Event()

    async def hang() -> None:
        started.set()
        await asyncio.sleep(60)

    task = asyncio.create_task(run_cancellable(hang()))
    _ = await started.wait()
    _ = task.cancel()
    with pytest.raises(asyncio.CancelledError):
        await task


async def test_cancel_turn_rolls_back_and_sets_pending() -> None:
    import asyncio

    from plyngent.cli.display import render_events
    from plyngent.cli.retry import run_cancellable

    store = await MemoryStore.open(DatabaseConfig())
    session = await store.create_session(name="t")

    class HangClient:
        @overload
        async def chat_completions(
            self, param: ChatCompletionsParam, *, stream: Literal[False] = False
        ) -> ChatCompletionResponse: ...

        @overload
        async def chat_completions(
            self, param: ChatCompletionsParam, *, stream: Literal[True]
        ) -> AsyncIterator[ChatCompletionChunk]: ...

        async def chat_completions(
            self, param: ChatCompletionsParam, *, stream: bool = False
        ) -> ChatCompletionResponse | AsyncIterator[ChatCompletionChunk]:
            del param, stream
            await asyncio.sleep(60)
            return _response(AssistantChatMessage(content="never"))

    agent = ChatAgent(HangClient(), model="m", memory=store, session_id=session.sid)
    turn = asyncio.create_task(run_cancellable(render_events(agent.run("cancel-me"))))
    await asyncio.sleep(0.05)
    _ = turn.cancel()
    with pytest.raises(asyncio.CancelledError):
        await turn

    assert agent.pending_retry_text == "cancel-me"
    assert len(agent.messages) == 1
    assert isinstance(agent.messages[0], UserChatMessage)
    loaded = await store.list_messages(session.sid)
    assert len(loaded) == 1
    assert isinstance(loaded[0], UserChatMessage)
    await store.close()


async def test_manual_retry_after_exhausted(monkeypatch: pytest.MonkeyPatch) -> None:
    async def no_wait(_seconds: float) -> bool:
        return True

    monkeypatch.setattr("plyngent.cli.retry.sleep_cancellable", no_wait)
    store = await MemoryStore.open(DatabaseConfig())
    session = await store.create_session(name="t")
    client = FlakyClient(fail_times=5)
    agent = ChatAgent(client, model="m", memory=store, session_id=session.sid)

    ok = await run_turn_with_retries(
        agent,
        starter=lambda: agent.run("hold-me"),
        delays=(0.01,),  # one auto-retry only → 2 attempts total, both fail
    )
    assert ok is False
    assert agent.pending_retry_text == "hold-me"
    assert len(await store.list_messages(session.sid)) == 1

    # Now succeed on manual /retry path (no second user message)
    client.fail_times = client.calls  # next call succeeds
    ok2 = await retry_pending_with_retries(agent)
    assert ok2 is True
    assert agent.pending_retry_text is None
    loaded = await store.list_messages(session.sid)
    assert len(loaded) == 2
    assert isinstance(loaded[0], UserChatMessage)
    assert loaded[0].content == "hold-me"
    assert sum(1 for m in loaded if isinstance(m, UserChatMessage)) == 1
    await store.close()