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, ToolRegistry, tool
from plyngent.cli.retry import (
    DEFAULT_MAX_AUTO_RETRIES,
    DEFAULT_RETRY_DELAYS_SECONDS,
    _wait_for_retry,
    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,
    AssistantFunctionTool,
    AssistantFunctionToolCall,
    ChatCompletionChoice,
    ChatCompletionChunk,
    ChatCompletionResponse,
    ChatCompletionsParam,
    ChunkChoice,
    DeltaMessage,
    UserChatMessage,
)
from plyngent.memory import MemoryStore

if TYPE_CHECKING:
    from collections.abc import AsyncIterator

    from plyngent.agent import AgentEvent


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_second_ctrl_c_during_cancel_message_stays_in_repl(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """Regression: a second Ctrl+C during the cancel message must not exit.

    After the turn task is cancelled the asyncio SIGINT handler is gone, so the
    second Ctrl+C arrives as a KeyboardInterrupt while the "cancelled" message
    is printing. It must be swallowed and control must return to the REPL
    (``False``), not propagate out of ``run_turn_with_retries``.
    """
    import asyncio

    calls: list[int] = [0]

    def raiser(message: object = None, **kwargs: object) -> None:
        calls[0] += 1
        if calls[0] == 1:
            raise KeyboardInterrupt

    monkeypatch.setattr("plyngent.cli.retry.click.secho", raiser)

    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_turn_with_retries(agent, starter=lambda: agent.run("cancel-me"), delays=()))
    await asyncio.sleep(0.05)
    _ = turn.cancel()  # first Ctrl+C: cancels the in-flight turn
    ok = await turn  # must return False instead of raising KeyboardInterrupt
    assert ok is False
    assert agent.pending_retry_text == "cancel-me"
    await store.close()


async def test_second_ctrl_c_during_retry_cancel_message_stays_benign(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """Regression: a second Ctrl+C during the retry-cancel message is benign."""

    async def interrupted_sleep(_seconds: float) -> bool:
        return False  # first Ctrl+C interrupted the retry wait

    monkeypatch.setattr("plyngent.cli.retry.sleep_cancellable", interrupted_sleep)

    calls: list[int] = [0]

    def raiser(message: object = None, **kwargs: object) -> None:
        calls[0] += 1
        if calls[0] == 2:  # second Ctrl+C lands on the "auto-retry cancelled" line
            raise KeyboardInterrupt

    monkeypatch.setattr("plyngent.cli.retry.click.secho", raiser)

    assert await _wait_for_retry(1, 3, 0.01) is False


async def test_second_ctrl_c_during_retry_announce_stays_benign(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """A Ctrl+C during the "retrying" announcement does not abort the turn."""

    async def ok_sleep(_seconds: float) -> bool:
        return True

    monkeypatch.setattr("plyngent.cli.retry.sleep_cancellable", ok_sleep)

    calls: list[int] = [0]

    def raiser(message: object = None, **kwargs: object) -> None:
        calls[0] += 1
        if calls[0] == 2:  # Ctrl+C lands on the "retrying" line
            raise KeyboardInterrupt

    monkeypatch.setattr("plyngent.cli.retry.click.secho", raiser)

    assert await _wait_for_retry(1, 3, 0.01) is True


async def test_turn_boundaries_reset_auto_continue() -> None:
    from plyngent.cli.limits import (
        auto_continue_enabled,
        prompt_continue_limit,
        reset_auto_continue,
    )
    from plyngent.prompting import temporary_backend
    from tests.test_prompting import ScriptedBackend

    reset_auto_continue()
    try:
        with temporary_backend(ScriptedBackend(["yyy"])):
            assert prompt_continue_limit("limit") is True
        assert auto_continue_enabled() is True

        store = await MemoryStore.open(DatabaseConfig())
        try:
            session = await store.create_session(name="t")
            client = FlakyClient(fail_times=0)
            agent = ChatAgent(client, model="m", memory=store, session_id=session.sid)
            seen: list[bool] = []

            async def starter() -> AsyncIterator[AgentEvent]:
                seen.append(auto_continue_enabled())
                async for event in agent.run("hi"):
                    yield event

            ok = await run_turn_with_retries(agent, starter=starter, delays=())
            assert ok is True
            assert seen == [False]  # cleared before the turn starts
            assert auto_continue_enabled() is False  # cleared after the turn ends
        finally:
            await store.close()
    finally:
        reset_auto_continue()


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()


class StepClient:
    """Scripted chat completions: each step is a response or a raised exception."""

    steps: list[ChatCompletionResponse | Exception]
    calls: int

    def __init__(self, steps: list[ChatCompletionResponse | Exception]) -> None:
        self.steps = list(steps)
        self.calls = 0

    @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
        step = self.steps.pop(0) if self.steps else RuntimeError("no more scripted steps")
        if isinstance(step, Exception):
            raise step
        if stream:
            return _as_stream(step)
        return step


def _tool_call(name: str) -> AssistantChatMessage:
    return AssistantChatMessage(
        content="",
        tool_calls=[
            AssistantFunctionToolCall(
                id="1",
                function=AssistantFunctionTool(name=name, arguments="{}"),
            )
        ],
    )


def _capture_waits(monkeypatch: pytest.MonkeyPatch) -> list[float]:
    """Record retry waits instead of sleeping; returns the collected delays."""
    waits: list[float] = []

    async def fake_sleep(seconds: float) -> bool:
        waits.append(seconds)
        return True

    monkeypatch.setattr("plyngent.cli.retry.sleep_cancellable", fake_sleep)
    return waits


async def test_retry_budget_resets_after_a_completed_round(
    monkeypatch: pytest.MonkeyPatch,
    capsys: pytest.CaptureFixture[str],
) -> None:
    """A round the model answered restarts the retry budget and its delays."""
    waits = _capture_waits(monkeypatch)

    @tool(register=False)
    def ping() -> str:
        return "pong"

    client = StepClient(
        [
            RuntimeError("connection lost"),  # attempt 1: no round completed
            _response(_tool_call("ping")),  # attempt 2, round 1: answered
            RuntimeError("connection lost"),  # attempt 2, round 2: dropped again
            _response(AssistantChatMessage(content="done")),  # attempt 3: completes
        ]
    )
    agent = ChatAgent(client, model="m", tools=ToolRegistry([ping]), stream=False)

    ok = await run_turn_with_retries(
        agent,
        starter=lambda: agent.run("hold-me"),
        delays=(5.0, 10.0, 15.0),
    )
    out = capsys.readouterr().out
    assert ok is True
    assert waits == [5.0, 5.0]  # the delay schedule restarted from the first step
    assert out.count("auto-retry 1/3") == 2
    assert "auto-retry 2/3" not in out


async def test_retry_budget_keeps_counting_without_progress(
    monkeypatch: pytest.MonkeyPatch,
    capsys: pytest.CaptureFixture[str],
) -> None:
    """Consecutive failures with no completed round share one budget."""
    waits = _capture_waits(monkeypatch)
    client = StepClient(
        [
            RuntimeError("connection lost"),
            RuntimeError("connection lost"),
            _response(AssistantChatMessage(content="done")),
        ]
    )
    agent = ChatAgent(client, model="m", stream=False)

    ok = await run_turn_with_retries(
        agent,
        starter=lambda: agent.run("hold-me"),
        delays=(5.0, 10.0, 15.0),
    )
    out = capsys.readouterr().out
    assert ok is True
    assert waits == [5.0, 10.0]
    assert "auto-retry 1/3" in out
    assert "auto-retry 2/3" in out