cli/retry: reset the auto-retry budget after a recovered round

src/plyngent/cli/retry.py

@@ -8,6 +8,7 @@ from typing import TYPE_CHECKING
 
 import click
 
+from plyngent.agent import AssistantMessageEvent
 from plyngent.cli.display import render_events
 from plyngent.cli.interrupt import allow_task_cancel, set_sigint_reinstall
 from plyngent.cli.limits import reset_auto_continue_turn
@@ -186,6 +187,10 @@ async def run_turn_with_retries(
 
     Ctrl+C cancels the in-flight task; user message stays in DB for ``/retry``.
 
+    The retry budget counts *consecutive* failures: an attempt that completed a
+    model round proves the connection answered, so the next failure starts a
+    fresh budget (and the short first delay) again.
+
     The ``yyy`` limit auto-continue is turn-scoped: cleared before and after the
     turn, so a retry of the same turn keeps it but the next user turn prompts.
     """
@@ -196,7 +201,7 @@ async def run_turn_with_retries(
         reset_auto_continue_turn()
 
 
-async def _run_turn_with_retries(
+async def _run_turn_with_retries(  # noqa: C901 — attempt / cancel / error state machine
     agent: ChatAgent,
     *,
     starter: Callable[[], AsyncIterator[AgentEvent]],
@@ -204,10 +209,21 @@ async def _run_turn_with_retries(
 ) -> bool:
     max_retries = len(delays)
     attempt = 0
+    recovered = False
     current: Callable[[], AsyncIterator[AgentEvent]] = starter
+
+    async def note_completed_round(events: AsyncIterator[AgentEvent]) -> AsyncIterator[AgentEvent]:
+        """Yield *events*, recording that a full model round came back."""
+        nonlocal recovered
+        async for event in events:
+            if isinstance(event, AssistantMessageEvent):
+                recovered = True
+            yield event
+
     while True:
+        recovered = False
         try:
-            await run_cancellable(render_events(current()))
+            await run_cancellable(render_events(note_completed_round(current())))
         except asyncio.CancelledError:
             # Do not auto-retry cancelled turns — user intent was stop, not retry.
             _echo_cancel_lines("", "cancelled; user message kept — use /retry to try again", "")
@@ -219,6 +235,10 @@ async def _run_turn_with_retries(
             click.secho(f"error: {exc}", fg="red")
             if agent.pending_retry_text is not None:
                 current = agent.retry
+            if recovered:
+                # The connection answered a round before failing again, so this
+                # is a fresh outage: restart the budget and the delay schedule.
+                attempt = 0
             if attempt >= max_retries:
                 if agent.pending_retry_text is not None:
                     click.secho(

tests/test_cli/test_retry.py

@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Literal, overload
 
 import pytest
 
-from plyngent.agent import ChatAgent
+from plyngent.agent import ChatAgent, ToolRegistry, tool
 from plyngent.cli.retry import (
     DEFAULT_MAX_AUTO_RETRIES,
     DEFAULT_RETRY_DELAYS_SECONDS,
@@ -17,6 +17,8 @@ from plyngent.cli.retry import (
 from plyngent.config.models import DatabaseConfig
 from plyngent.lmproto.openai_compatible.model import (
     AssistantChatMessage,
+    AssistantFunctionTool,
+    AssistantFunctionToolCall,
     ChatCompletionChoice,
     ChatCompletionChunk,
     ChatCompletionResponse,
@@ -365,3 +367,120 @@ async def test_manual_retry_after_exhausted(monkeypatch: pytest.MonkeyPatch) ->
     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