diff --git a/nanobot/providers/base.py b/nanobot/providers/base.py index 20364b687..f603119b9 100644 --- a/nanobot/providers/base.py +++ b/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: diff --git a/nanobot/providers/fallback_provider.py b/nanobot/providers/fallback_provider.py index 513cf0d01..d59217843 100644 --- a/nanobot/providers/fallback_provider.py +++ b/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, diff --git a/tests/agent/test_runner_fallback.py b/tests/agent/test_runner_fallback.py index 30f8eadc0..667349a41 100644 --- a/tests/agent/test_runner_fallback.py +++ b/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( diff --git a/tests/providers/test_provider_retry.py b/tests/providers/test_provider_retry.py index dfba6a1f5..e3e32c28f 100644 --- a/tests/providers/test_provider_retry.py +++ b/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()])