refactor: pass retry exhaustion callbacks explicitly

nanobot/providers/base.py

@@ -7,9 +7,8 @@ import json
 import os
 import re
 from abc import ABC, abstractmethod
-from collections.abc import Awaitable, Callable, Generator
-from contextlib import contextmanager, suppress
-from contextvars import ContextVar
+from collections.abc import Awaitable, Callable
+from contextlib import suppress
 from copy import deepcopy
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
@@ -27,24 +26,6 @@ MAX_STREAM_IDLE_TIMEOUT_S = 3600.0
 RETRY_AFTER_BUFFER = 1
 
 RetryEventCallback = Callable[[str], Awaitable[None]]
-_RETRY_EXHAUSTED_CALLBACK: ContextVar[RetryEventCallback | None] = ContextVar(
-    "nanobot_retry_exhausted_callback",
-    default=None,
-)
-
-
-@contextmanager
-def retry_exhaustion_callback(callback: RetryEventCallback) -> Generator[None, None, None]:
-    """Redirect terminal retry events within one async call context.
-
-    Provider wrappers use this internal scope to defer a candidate's terminal
-    notification without changing the public retry-method signatures.
-    """
-    token = _RETRY_EXHAUSTED_CALLBACK.set(callback)
-    try:
-        yield
-    finally:
-        _RETRY_EXHAUSTED_CALLBACK.reset(token)
 
 
 def resolve_stream_idle_timeout_s(
@@ -893,8 +874,9 @@ class LLMProvider(ABC):
         on_tool_call_delta: Callable[[dict[str, Any]], Awaitable[None]] | None = None,
         on_stream_recover: Callable[[], Awaitable[None]] | None = None,
         retry_mode: str = "standard",
-        on_retry_wait: Callable[[str], Awaitable[None]] | None = None,
+        on_retry_wait: RetryEventCallback | None = None,
         provider_context: ProviderCallContext | None = None,
+        on_retry_exhausted: RetryEventCallback | None = None,
     ) -> LLMResponse:
         """Call chat_stream() with retry on transient provider failures."""
         if max_tokens is self._SENTINEL or max_tokens is None:
@@ -931,16 +913,13 @@ class LLMProvider(ABC):
             kw["provider_context"] = provider_context
         if on_stream_recover and getattr(self, "supports_stream_recover_callback", False):
             kw["on_stream_recover"] = _recover_stream
-        on_retry_exhausted = _RETRY_EXHAUSTED_CALLBACK.get()
         return await self._run_with_retry(
             self._safe_chat_stream,
             kw,
             messages,
             retry_mode=retry_mode,
             on_retry_wait=on_retry_wait,
-            on_retry_exhausted=(
-                on_retry_exhausted if on_retry_exhausted is not None else on_retry_wait
-            ),
+            on_retry_exhausted=on_retry_exhausted or on_retry_wait,
             should_retry_guard=lambda: not has_streamed_content,
             on_stream_recover=_recover_stream if on_stream_recover else None,
         )
@@ -955,8 +934,9 @@ class LLMProvider(ABC):
         reasoning_effort: object = _SENTINEL,
         tool_choice: str | dict[str, Any] | None = None,
         retry_mode: str = "standard",
-        on_retry_wait: Callable[[str], Awaitable[None]] | None = None,
+        on_retry_wait: RetryEventCallback | None = None,
         provider_context: ProviderCallContext | None = None,
+        on_retry_exhausted: RetryEventCallback | None = None,
     ) -> LLMResponse:
         """Call chat() with retry on transient provider failures.
 
@@ -981,16 +961,13 @@ class LLMProvider(ABC):
         )
         if provider_context is not None:
             kw["provider_context"] = provider_context
-        on_retry_exhausted = _RETRY_EXHAUSTED_CALLBACK.get()
         return await self._run_with_retry(
             self._safe_chat,
             kw,
             messages,
             retry_mode=retry_mode,
             on_retry_wait=on_retry_wait,
-            on_retry_exhausted=(
-                on_retry_exhausted if on_retry_exhausted is not None else on_retry_wait
-            ),
+            on_retry_exhausted=on_retry_exhausted or on_retry_wait,
         )
 
     @classmethod
@@ -1095,8 +1072,8 @@ class LLMProvider(ABC):
         original_messages: list[dict[str, Any]],
         *,
         retry_mode: str,
-        on_retry_wait: Callable[[str], Awaitable[None]] | None,
-        on_retry_exhausted: Callable[[str], Awaitable[None]] | None,
+        on_retry_wait: RetryEventCallback | None,
+        on_retry_exhausted: RetryEventCallback | None,
         should_retry_guard: Callable[[], bool] | None = None,
         on_stream_recover: Callable[[], Awaitable[None]] | None = None,
     ) -> LLMResponse:

nanobot/providers/fallback_provider.py

@@ -17,7 +17,7 @@ from nanobot.providers.base import (
     LLMResponse,
     ProviderCallContext,
     ProviderConversationState,
-    retry_exhaustion_callback,
+    RetryEventCallback,
 )
 
 # Circuit breaker tuned to match OpenAICompatProvider's Responses API breaker.
@@ -206,8 +206,9 @@ class FallbackProvider(LLMProvider):
         reasoning_effort: object = _UNSET,
         tool_choice: str | dict[str, Any] | None = None,
         retry_mode: str = "standard",
-        on_retry_wait: Callable[[str], Awaitable[None]] | None = None,
+        on_retry_wait: RetryEventCallback | None = None,
         provider_context: ProviderCallContext | None = None,
+        on_retry_exhausted: RetryEventCallback | None = None,
     ) -> LLMResponse:
         """Exhaust each provider's retries before moving to the next fallback."""
         call_kwargs: dict[str, Any] = {
@@ -217,6 +218,7 @@ class FallbackProvider(LLMProvider):
             "tool_choice": tool_choice,
             "retry_mode": retry_mode,
             "on_retry_wait": on_retry_wait,
+            "on_retry_exhausted": on_retry_exhausted,
         }
         if max_tokens is not _UNSET:
             call_kwargs["max_tokens"] = max_tokens
@@ -235,6 +237,7 @@ class FallbackProvider(LLMProvider):
             lambda p, kw: p.chat_with_retry(**kw),
             call_kwargs,
             has_streamed=None,
+            on_retry_exhausted=on_retry_exhausted or on_retry_wait,
         )
 
     async def chat_with_context(
@@ -292,8 +295,9 @@ class FallbackProvider(LLMProvider):
         on_tool_call_delta: Callable[[dict[str, Any]], Awaitable[None]] | None = None,
         on_stream_recover: Callable[[], Awaitable[None]] | None = None,
         retry_mode: str = "standard",
-        on_retry_wait: Callable[[str], Awaitable[None]] | None = None,
+        on_retry_wait: RetryEventCallback | None = None,
         provider_context: ProviderCallContext | None = None,
+        on_retry_exhausted: RetryEventCallback | None = None,
     ) -> LLMResponse:
         """Exhaust streaming retries on one provider before failing over."""
         call_kwargs: dict[str, Any] = {
@@ -306,6 +310,7 @@ class FallbackProvider(LLMProvider):
             "on_tool_call_delta": on_tool_call_delta,
             "retry_mode": retry_mode,
             "on_retry_wait": on_retry_wait,
+            "on_retry_exhausted": on_retry_exhausted,
         }
         if max_tokens is not _UNSET:
             call_kwargs["max_tokens"] = max_tokens
@@ -350,6 +355,7 @@ class FallbackProvider(LLMProvider):
             has_streamed=has_streamed,
             on_stream_recover=_recover_stream if on_stream_recover is not None else None,
             persistent_retry_guard=lambda: not has_unrecovered_stream[0],
+            on_retry_exhausted=on_retry_exhausted or on_retry_wait,
         )
 
     async def _route_with_retry_fallback(
@@ -359,16 +365,17 @@ class FallbackProvider(LLMProvider):
         has_streamed: list[bool] | None,
         on_stream_recover: Callable[[], Awaitable[None]] | None = None,
         persistent_retry_guard: Callable[[], bool] | None = None,
+        on_retry_exhausted: RetryEventCallback | None = None,
     ) -> LLMResponse:
         """Apply finite retries per provider and persistence to the whole chain."""
-        on_retry_wait: Callable[[str], Awaitable[None]] | None = kwargs.get("on_retry_wait")
+        on_retry_wait: RetryEventCallback | None = kwargs.get("on_retry_wait")
         if kwargs.get("retry_mode", "standard") != "persistent":
             return await self._try_with_retry_fallback(
                 call,
                 kwargs,
                 has_streamed=has_streamed,
                 on_stream_recover=on_stream_recover,
-                on_retry_exhausted=on_retry_wait,
+                on_retry_exhausted=on_retry_exhausted,
             )
 
         async def _call_chain(**chain_kwargs: Any) -> LLMResponse:
@@ -387,7 +394,7 @@ class FallbackProvider(LLMProvider):
             kwargs["messages"],
             retry_mode="persistent",
             on_retry_wait=on_retry_wait,
-            on_retry_exhausted=on_retry_wait,
+            on_retry_exhausted=on_retry_exhausted,
             should_retry_guard=persistent_retry_guard,
             on_stream_recover=on_stream_recover,
         )
@@ -398,7 +405,7 @@ class FallbackProvider(LLMProvider):
         kwargs: dict[str, Any],
         has_streamed: list[bool] | None,
         on_stream_recover: Callable[[], Awaitable[None]] | None,
-        on_retry_exhausted: Callable[[str], Awaitable[None]] | None,
+        on_retry_exhausted: RetryEventCallback | None,
     ) -> LLMResponse:
         """Defer a provider's terminal retry event until the chain fails."""
         last_exhausted_message: str | None = None
@@ -413,8 +420,11 @@ class FallbackProvider(LLMProvider):
         ) -> LLMResponse:
             nonlocal last_exhausted_message
             last_exhausted_message = None
-            with retry_exhaustion_callback(_capture_exhaustion):
-                return await call(provider, call_kwargs)
+            candidate_kwargs = {
+                **call_kwargs,
+                "on_retry_exhausted": _capture_exhaustion,
+            }
+            return await call(provider, candidate_kwargs)
 
         response = await self._try_with_fallback(
             _call_with_deferred_exhaustion,

tests/agent/test_runner_fallback.py

@@ -858,10 +858,14 @@ class TestRetryBeforeFailover:
             _make_response("fallback unavailable", finish_reason="error", error_status_code=503),
         )
         retry_events: list[str] = []
+        terminal_events: list[str] = []
 
         async def _on_retry_event(message: str) -> None:
             retry_events.append(message)
 
+        async def _on_retry_exhausted(message: str) -> None:
+            terminal_events.append(message)
+
         fb = FallbackProvider(
             primary=primary,
             fallback_presets=[_fallback("fallback-a")],
@@ -872,12 +876,13 @@ class TestRetryBeforeFailover:
             result = await fb.chat_with_retry(
                 messages=[{"role": "user", "content": "hi"}],
                 on_retry_wait=_on_retry_event,
+                on_retry_exhausted=_on_retry_exhausted,
             )
 
-        terminal_events = [event for event in retry_events if "giving up" in event]
         assert result.finish_reason == "error"
         assert len(primary.chat_calls) == 4
         assert len(fallback.chat_calls) == 4
+        assert not any("giving up" in event for event in retry_events)
         assert terminal_events == ["Model request failed after 4 attempts, giving up."]
 
     @pytest.mark.asyncio
@@ -1027,41 +1032,6 @@ class TestRetryBeforeFailover:
             "identical errors."
         ]
 
-    @pytest.mark.asyncio
-    async def test_legacy_retry_override_does_not_receive_internal_callback(self) -> None:
-        class LegacyRetryProvider(_FakeProvider):
-            def __init__(self) -> None:
-                super().__init__("legacy", _make_response("legacy ok"))
-                self.retry_calls = 0
-
-            async def chat_with_retry(
-                self,
-                messages: list[dict[str, Any]],
-                tools: list[dict[str, Any]] | None = None,
-                model: str | None = None,
-                max_tokens: object = None,
-                temperature: object = None,
-                reasoning_effort: object = None,
-                tool_choice: str | dict[str, Any] | None = None,
-                retry_mode: str = "standard",
-                on_retry_wait: Any = None,
-                provider_context: ProviderCallContext | None = None,
-            ) -> LLMResponse:
-                self.retry_calls += 1
-                return await self.chat(messages=messages)
-
-        primary = LegacyRetryProvider()
-        fb = FallbackProvider(
-            primary=primary,
-            fallback_presets=[_fallback("fallback-a")],
-            provider_factory=MagicMock(),
-        )
-
-        result = await fb.chat_with_retry(messages=[{"role": "user", "content": "hi"}])
-
-        assert result.content == "legacy ok"
-        assert primary.retry_calls == 1
-
     @pytest.mark.asyncio
     async def test_fallback_retries_before_trying_next_model(self) -> None:
         primary = _FakeProvider(

tests/providers/test_provider_retry.py

@@ -132,6 +132,39 @@ async def test_chat_with_retry_emits_terminal_progress_when_standard_retries_exh
     assert progress[-1] == "Model request failed after 4 attempts, giving up."
 
 
+@pytest.mark.asyncio
+async def test_chat_with_retry_routes_terminal_progress_to_explicit_callback(monkeypatch) -> None:
+    provider = ScriptedProvider([
+        LLMResponse(content="429 rate limit a", finish_reason="error"),
+        LLMResponse(content="429 rate limit b", finish_reason="error"),
+        LLMResponse(content="429 rate limit c", finish_reason="error"),
+        LLMResponse(content="503 final server error", finish_reason="error"),
+    ])
+    retry_progress: list[str] = []
+    terminal_progress: list[str] = []
+
+    async def _fake_sleep(delay: int) -> None:
+        return None
+
+    async def _retry_progress(msg: str) -> None:
+        retry_progress.append(msg)
+
+    async def _terminal_progress(msg: str) -> None:
+        terminal_progress.append(msg)
+
+    monkeypatch.setattr("nanobot.providers.base.asyncio.sleep", _fake_sleep)
+
+    response = await provider.chat_with_retry(
+        messages=[{"role": "user", "content": "hello"}],
+        on_retry_wait=_retry_progress,
+        on_retry_exhausted=_terminal_progress,
+    )
+
+    assert response.content == "503 final server error"
+    assert not any("giving up" in message for message in retry_progress)
+    assert terminal_progress == ["Model request failed after 4 attempts, giving up."]
+
+
 @pytest.mark.asyncio
 async def test_chat_with_retry_preserves_cancelled_error() -> None:
     provider = ScriptedProvider([asyncio.CancelledError()])