diff --git a/services/ai/providers/anthropic.py b/services/ai/providers/anthropic.py index cf66a949..ec6d2a46 100644 --- a/services/ai/providers/anthropic.py +++ b/services/ai/providers/anthropic.py @@ -7,7 +7,14 @@ from collections.abc import AsyncIterator, Iterable from typing import Any, ClassVar, cast -from anthropic import APIStatusError, AsyncAnthropic, AsyncStream, MessageStreamEvent +import httpx +from anthropic import ( + APIConnectionError, + APIStatusError, + AsyncAnthropic, + AsyncStream, + MessageStreamEvent, +) from anthropic.types import MessageParam, ToolParam from . import LLMProvider, TokenUsage @@ -32,6 +39,13 @@ def _anthropic_context_overflow(e: BaseException) -> bool: return "prompt is too long" in str(e).lower() +def _anthropic_is_retryable(e: BaseException) -> bool: + """Transport-level failure (connect/read/timeout, dropped stream) that a + caller may safely retry once. Status errors are excluded: the SDK already + retries those at request-creation time.""" + return isinstance(e, (httpx.TransportError, APIConnectionError)) + + logger = logging.getLogger(__name__) @@ -189,6 +203,7 @@ async def stream_response( status_code=_anthropic_status_code(e), cause=e, is_context_overflow=_anthropic_context_overflow(e), + is_retryable=_anthropic_is_retryable(e), ) from e async def generate_response( @@ -238,6 +253,7 @@ async def generate_response( status_code=_anthropic_status_code(e), cause=e, is_context_overflow=_anthropic_context_overflow(e), + is_retryable=_anthropic_is_retryable(e), ) from e async def health_check( diff --git a/services/ai/providers/azure_foundry.py b/services/ai/providers/azure_foundry.py index 45403a01..9cb2bb30 100644 --- a/services/ai/providers/azure_foundry.py +++ b/services/ai/providers/azure_foundry.py @@ -117,6 +117,7 @@ async def stream_response( status_code=e.status_code, cause=e, is_context_overflow=e.is_context_overflow, + is_retryable=e.is_retryable, ) from e async def generate_response( @@ -140,6 +141,7 @@ async def generate_response( status_code=e.status_code, cause=e, is_context_overflow=e.is_context_overflow, + is_retryable=e.is_retryable, ) from e async def health_check( diff --git a/services/ai/providers/bedrock.py b/services/ai/providers/bedrock.py index 0486471f..33656f74 100644 --- a/services/ai/providers/bedrock.py +++ b/services/ai/providers/bedrock.py @@ -9,6 +9,8 @@ from typing import Any, ClassVar, cast import boto3 +import httpx +from anthropic import APIConnectionError as AnthropicConnectionError from anthropic import APIStatusError, AnthropicBedrock from anthropic.types import ( MessageParam, @@ -36,7 +38,13 @@ ) from anthropic.types.message_stream_event import MessageStreamEvent from anthropic.types.raw_message_delta_event import Delta -from botocore.exceptions import ClientError +from botocore.exceptions import ( + ClientError, + ConnectionError as BotoConnectionError, + HTTPClientError, + SSLError as BotoSSLError, +) + from . import LLMProvider, TokenUsage from .anthropic_message_adapter import ( @@ -48,6 +56,23 @@ logger = logging.getLogger(__name__) +def _bedrock_is_retryable(e: BaseException) -> bool: + """Transport-level failure (dropped/errored HTTP connection, mid-stream + disconnect) that a caller may safely retry once. ``ClientError`` (API + status errors) is excluded: boto3 already retries those at request-creation + time.""" + return isinstance( + e, + ( + httpx.TransportError, + AnthropicConnectionError, + HTTPClientError, + BotoConnectionError, + BotoSSLError, + ), + ) + + def sanitize_document_name(name: str) -> str: """ Sanitize document name for AWS Bedrock requirements. @@ -113,7 +138,9 @@ def __init__(self, model_id: str, region_name: str | None = None): else: self.client = boto3.client("bedrock-runtime", region_name=region_name) - def _to_provider_error(self, e: Exception, model: str | None = None) -> ProviderError: + def _to_provider_error( + self, e: Exception, model: str | None = None + ) -> ProviderError: is_context_overflow = False if isinstance(e, ClientError): error = e.response.get("Error", {}) @@ -128,9 +155,7 @@ def _to_provider_error(self, e: Exception, model: str | None = None) -> Provider else: body = str(e) status = e.status_code if isinstance(e, APIStatusError) else None - is_context_overflow = ( - isinstance(e, APIStatusError) and e.status_code == 413 - ) + is_context_overflow = isinstance(e, APIStatusError) and e.status_code == 413 return ProviderError( body, provider_type=self.provider_type, @@ -138,6 +163,7 @@ def _to_provider_error(self, e: Exception, model: str | None = None) -> Provider status_code=status, cause=e, is_context_overflow=is_context_overflow, + is_retryable=_bedrock_is_retryable(e), ) def _determine_model_family(self, model_id: str) -> str: diff --git a/services/ai/providers/gemini.py b/services/ai/providers/gemini.py index dcf1158d..b6293836 100644 --- a/services/ai/providers/gemini.py +++ b/services/ai/providers/gemini.py @@ -18,6 +18,7 @@ from collections.abc import AsyncIterator from typing import Any, ClassVar +import httpx from google import genai from google.genai import types from google.genai.errors import APIError @@ -47,6 +48,14 @@ def _gemini_status_code(e: BaseException) -> int | None: return e.code if isinstance(e, APIError) else None +def _gemini_is_retryable(e: BaseException) -> bool: + """Transport-level failure (connect/read/timeout, dropped stream) that a + caller may safely retry once. The genai client is httpx-based; status + errors surface as ``APIError`` and are already retried by the client at + request-creation time.""" + return isinstance(e, httpx.TransportError) + + logger = logging.getLogger(__name__) THOUGHT_SIGNATURE_KEY = "_gemini_thought_signature" @@ -460,6 +469,7 @@ async def stream_response( model=model or self.model_name, status_code=_gemini_status_code(e), cause=e, + is_retryable=_gemini_is_retryable(e), ) from e async def generate_response( @@ -504,6 +514,7 @@ async def generate_response( model=model or self.model_name, status_code=_gemini_status_code(e), cause=e, + is_retryable=_gemini_is_retryable(e), ) from e async def health_check( diff --git a/services/ai/providers/openai.py b/services/ai/providers/openai.py index c9965d08..3e6f8027 100644 --- a/services/ai/providers/openai.py +++ b/services/ai/providers/openai.py @@ -9,7 +9,8 @@ from collections.abc import AsyncIterator from typing import Any, ClassVar -from openai import APIStatusError, AsyncOpenAI +import httpx +from openai import APIConnectionError, APIStatusError, AsyncOpenAI from anthropic.types import ( Message, MessageDeltaUsage, @@ -55,6 +56,13 @@ def _openai_context_overflow(e: BaseException) -> bool: return _openai_error_code(e) == "context_length_exceeded" +def _openai_is_retryable(e: BaseException) -> bool: + """Transport-level failure (connect/read/timeout, dropped stream) that a + caller may safely retry once. Status errors are excluded: the SDK already + retries those at request-creation time.""" + return isinstance(e, (httpx.TransportError, APIConnectionError)) + + logger = logging.getLogger(__name__) MIN_REASONING_OUTPUT_TOKENS = 1024 @@ -295,7 +303,9 @@ async def stream_response( ) elif event_type == "error": - raise RuntimeError(getattr(event, "message", "Unknown stream error")) + raise RuntimeError( + getattr(event, "message", "Unknown stream error") + ) yield RawMessageStopEvent(type="message_stop") @@ -310,6 +320,7 @@ async def stream_response( status_code=_openai_status_code(e), cause=e, is_context_overflow=_openai_context_overflow(e), + is_retryable=_openai_is_retryable(e), ) from e def _convert_messages(self, messages: list[dict[str, Any]]) -> list[dict[str, Any]]: @@ -481,6 +492,7 @@ async def generate_response( status_code=_openai_status_code(e), cause=e, is_context_overflow=_openai_context_overflow(e), + is_retryable=_openai_is_retryable(e), ) from e async def health_check( diff --git a/services/ai/providers/openai_compatible.py b/services/ai/providers/openai_compatible.py index eba8bc59..de40c05a 100644 --- a/services/ai/providers/openai_compatible.py +++ b/services/ai/providers/openai_compatible.py @@ -12,7 +12,8 @@ from collections.abc import AsyncIterator from typing import Any, ClassVar, cast -from openai import APIStatusError, AsyncOpenAI +import httpx +from openai import APIConnectionError, APIStatusError, AsyncOpenAI from openai.types.chat import ( ChatCompletionAssistantMessageParam, ChatCompletionChunk, @@ -75,6 +76,13 @@ def _openai_compat_context_overflow(e: BaseException) -> bool: return _openai_compat_error_code(e) == "context_length_exceeded" +def _openai_compat_is_retryable(e: BaseException) -> bool: + """Transport-level failure (connect/read/timeout, dropped stream) that a + caller may safely retry once. Status errors are excluded: the SDK already + retries those at request-creation time.""" + return isinstance(e, (httpx.TransportError, APIConnectionError)) + + logger = logging.getLogger(__name__) # Some OpenAI-compatible providers expose non-standard assistant-message fields @@ -536,6 +544,7 @@ async def stream_response( status_code=_openai_compat_status_code(e), cause=e, is_context_overflow=_openai_compat_context_overflow(e), + is_retryable=_openai_compat_is_retryable(e), ) from e async def generate_response( @@ -592,6 +601,7 @@ async def generate_response( status_code=_openai_compat_status_code(e), cause=e, is_context_overflow=_openai_compat_context_overflow(e), + is_retryable=_openai_compat_is_retryable(e), ) from e async def health_check( diff --git a/services/ai/providers/types.py b/services/ai/providers/types.py index eeb0ee09..c75a79d8 100644 --- a/services/ai/providers/types.py +++ b/services/ai/providers/types.py @@ -36,6 +36,7 @@ def __init__( status_code: int | None = None, cause: BaseException | None = None, is_context_overflow: bool | None = None, + is_retryable: bool = False, ): super().__init__(message) self.message = message @@ -43,5 +44,10 @@ def __init__( self.model = model self.status_code = status_code self.is_context_overflow = bool(is_context_overflow) + # True only for transport-level failures (connection reset, mid-stream + # read error, remote disconnect) that a caller may safely retry once. + # Providers set it from the underlying exception type; a missing + # status_code alone does not imply transport failure. + self.is_retryable = is_retryable if cause is not None: self.__cause__ = cause diff --git a/services/ai/providers/vertex_ai.py b/services/ai/providers/vertex_ai.py index b857d5c8..90b007be 100644 --- a/services/ai/providers/vertex_ai.py +++ b/services/ai/providers/vertex_ai.py @@ -92,6 +92,7 @@ async def stream_response( status_code=e.status_code, cause=e, is_context_overflow=e.is_context_overflow, + is_retryable=e.is_retryable, ) from e async def generate_response( @@ -115,6 +116,7 @@ async def generate_response( status_code=e.status_code, cause=e, is_context_overflow=e.is_context_overflow, + is_retryable=e.is_retryable, ) from e async def health_check( diff --git a/services/ai/streaming/generate.py b/services/ai/streaming/generate.py index f2152c87..c51a6b45 100644 --- a/services/ai/streaming/generate.py +++ b/services/ai/streaming/generate.py @@ -89,6 +89,21 @@ class CompactionEnd: """Internal marker emitted after a compaction pass finishes.""" +_TRANSIENT_RETRY_BACKOFF_SECONDS = 1.0 + + +def _is_retryable_provider_error(error: ProviderError) -> bool: + """Whether a failed provider stream is worth one automatic retry. + + Only transport-level failures (connection reset, mid-stream read error, + remote disconnect) qualify; each provider flags them via + ``ProviderError.is_retryable``. HTTP status errors are intentionally out of + scope: the provider SDKs already retry those at request-creation time, so + retrying here as well would multiply the request count. + """ + return error.is_retryable + + async def event_stream_with_context_retry( turn_tools: list[dict], conversation_messages: list[MessageParam], @@ -102,12 +117,32 @@ async def event_stream_with_context_retry( model_record_id: str | None = None, model_name: str | None = None, ) -> AsyncIterator[MessageStreamEvent | CompactionStart | CompactionEnd]: - """Stream events from the LLM provider with one automatic compaction retry - on context-overflow errors. + """Stream events from the LLM provider with one-shot automatic recovery. + + A single provider call is retried when it fails before any substantive + content has been streamed: + + - a context-overflow error triggers a forced compaction, then one retry; + - a transport-level failure (connection drop, mid-stream read error) + flagged by the provider triggers one retry after a short backoff. + + The failed attempt's usage is persisted before the retry, since the + provider may bill for both attempts. + + A retry is only safe while the client has not seen a content block, so + once any non-``message_start`` event has been emitted the error + propagates instead. A retried stream re-emits the ``message_start`` + envelope; the duplicate is dropped here because the caller already + buffered the original and content blocks carry no message id, so the + retained envelope stays valid. ``conversation_messages`` is mutated in-place when a compaction retry replaces the full history with a compacted version. """ + # message_start envelope already delivered to the caller. Tracked across + # attempts so a retried stream's duplicate envelope can be dropped. + emitted_message_start = False + for llm_attempt in range(2): tracker = UsageTracker( UsageRepository(), @@ -145,15 +180,24 @@ async def event_stream_with_context_retry( processed_stream ) - emitted_event = False + emitted_non_envelope = False try: async for wrapped_event in processed_stream: - emitted_event = True + event_type = getattr(wrapped_event, "type", None) + if event_type == "message_start": + if emitted_message_start: + # Retried stream re-emits the envelope the caller + # already buffered from the failed attempt. + continue + emitted_message_start = True + else: + emitted_non_envelope = True yield wrapped_event tracker.save() return except ProviderError as e: - if e.is_context_overflow and llm_attempt == 0 and not emitted_event: + can_retry = llm_attempt == 0 and not emitted_non_envelope + if can_retry and e.is_context_overflow: logger.warning( "Chat %s hit provider context limit; retrying once after forced compaction", chat_id, @@ -173,6 +217,14 @@ async def event_stream_with_context_retry( if should_emit_progress: yield CompactionEnd() continue + if can_retry and _is_retryable_provider_error(e): + logger.warning( + "Chat %s hit a transient transport error; retrying once", + chat_id, + ) + tracker.save() + await asyncio.sleep(_TRANSIENT_RETRY_BACKOFF_SECONDS) + continue raise diff --git a/services/ai/tests/integration/test_chat_stream_lifecycle.py b/services/ai/tests/integration/test_chat_stream_lifecycle.py index 49012560..77bad2f1 100644 --- a/services/ai/tests/integration/test_chat_stream_lifecycle.py +++ b/services/ai/tests/integration/test_chat_stream_lifecycle.py @@ -182,6 +182,9 @@ async def seeded_chat( finally: async with db_pool.acquire() as conn: await conn.execute("DELETE FROM chat_messages WHERE chat_id = $1", chat_id) + # Usage upserts are fire-and-forget; clear them before the chat + # (and model) rows they reference are removed. + await conn.execute("DELETE FROM model_usage WHERE chat_id = $1", chat_id) await conn.execute("DELETE FROM chats WHERE id = $1", chat_id) async with db_pool.acquire() as conn: await conn.execute("DELETE FROM models WHERE id = $1", model_id) @@ -1043,6 +1046,68 @@ async def stream_response(self, **kwargs): ) assert text == "Partial answer", f"Partial content lost: {text!r}" + @pytest.mark.asyncio + async def test_transient_mid_stream_drop_retries_and_persists_single_row( + self, seeded_chat, redis_client, redis_keys + ): + """A provider-flagged transport failure after message_start but before + any content is retried once transparently: the client sees exactly one + message_start, the recovered text lands on a single assistant row, and + no stream_error is emitted.""" + import httpx + + from providers.types import ProviderError, ProviderType + from tests.helpers import message_start_event, text_response_events + + chat_id, _user_id, model_id = seeded_chat + + class DropThenRecoverLLM(GatedRecordingLLM): + async def stream_response(self, **kwargs): + self.calls.append({"kwargs": kwargs}) + if len(self.calls) == 1: + yield message_start_event() + raise ProviderError( + "Connection reset mid-stream", + provider_type=ProviderType.ANTHROPIC, + model=model_id, + is_retryable=True, + cause=httpx.ReadError("stream reset"), + ) + for event in text_response_events("Recovered answer"): + yield event + + llm = DropThenRecoverLLM([("text", "unused")], model_id) + app = _build_chat_app(llm, redis_client, model_id) + async with _client(app) as client: + events = await collect_sse_events(client, chat_id) + + message_starts = [ + data + for event_type, data, _sid in events + if event_type == "message" + and json.loads(data).get("type") == "message_start" + ] + assert ( + len(message_starts) == 1 + ), f"Expected exactly one message_start, got {len(message_starts)}" + assert not any( + event_type == "stream_error" for event_type, _d, _sid in events + ), f"Expected no stream_error, got events: {events}" + assert len(llm.calls) == 2, f"Expected one retry, got {len(llm.calls)} calls" + + db_msgs = await MessagesRepository().get_active_path(chat_id) + assert ( + len(db_msgs) == 2 + ), f"Expected user + one assistant row, got {len(db_msgs)} rows" + assistant = db_msgs[-1] + assert assistant.error is None + text = " ".join( + b.get("text", "") + for b in assistant.message.get("content", []) + if b.get("type") == "text" + ) + assert text == "Recovered answer", f"Recovered content lost: {text!r}" + @pytest.mark.asyncio async def test_tool_loop_error_uses_last_persisted_parent( self, seeded_chat, redis_client, redis_keys diff --git a/services/ai/tests/unit/test_provider_error_flags.py b/services/ai/tests/unit/test_provider_error_flags.py new file mode 100644 index 00000000..530cde7f --- /dev/null +++ b/services/ai/tests/unit/test_provider_error_flags.py @@ -0,0 +1,100 @@ +"""Unit tests for per-provider transport-error classification. + +Each provider sets ``ProviderError.is_retryable`` from its own SDK's +transport exception types. These tests pin that classification: transport +failures must be flagged, deterministic API/local errors must not. +""" + +import httpx +import pytest +from anthropic import ( + APIConnectionError as AnthropicConnectionError, +) +from anthropic import ( + APIStatusError as AnthropicStatusError, +) +from botocore.exceptions import ( + ClientError, + ConnectionClosedError, + NoCredentialsError, + ReadTimeoutError, +) +from google.genai.errors import APIError as GenaiAPIError +from openai import APIConnectionError as OpenAIConnectionError +from openai import APIStatusError as OpenAIStatusError + +from providers.anthropic import _anthropic_is_retryable +from providers.bedrock import _bedrock_is_retryable +from providers.gemini import _gemini_is_retryable +from providers.openai import _openai_is_retryable +from providers.openai_compatible import _openai_compat_is_retryable + +_REQUEST = httpx.Request("POST", "https://provider.test/v1/chat") +_RESPONSE_429 = httpx.Response(429, request=_REQUEST) + + +def _throttling_client_error() -> ClientError: + return ClientError( + { + "Error": {"Code": "ThrottlingException", "Message": "Slow down"}, + "ResponseMetadata": {"HTTPStatusCode": 429}, + }, + "InvokeModel", + ) + + +@pytest.mark.unit +class TestTransportErrorClassification: + def test_httpx_transport_errors_are_retryable_everywhere(self): + read_error = httpx.ReadError("stream reset", request=_REQUEST) + connect_error = httpx.ConnectError("connection refused", request=_REQUEST) + timeout = httpx.ReadTimeout("timed out", request=_REQUEST) + for classifier in ( + _anthropic_is_retryable, + _openai_is_retryable, + _openai_compat_is_retryable, + _gemini_is_retryable, + _bedrock_is_retryable, + ): + assert classifier(read_error) + assert classifier(connect_error) + assert classifier(timeout) + + def test_anthropic_sdk_connection_error_is_retryable(self): + assert _anthropic_is_retryable(AnthropicConnectionError(request=_REQUEST)) + + def test_openai_sdk_connection_error_is_retryable(self): + assert _openai_is_retryable(OpenAIConnectionError(request=_REQUEST)) + + def test_bedrock_transport_errors_are_retryable(self): + assert _bedrock_is_retryable( + ConnectionClosedError(endpoint_url="https://bedrock.test") + ) + assert _bedrock_is_retryable( + ReadTimeoutError(endpoint_url="https://bedrock.test") + ) + assert _bedrock_is_retryable(httpx.ReadError("stream reset", request=_REQUEST)) + assert _bedrock_is_retryable(AnthropicConnectionError(request=_REQUEST)) + + def test_anthropic_deterministic_errors_are_not_retryable(self): + assert not _anthropic_is_retryable( + AnthropicStatusError("bad key", response=_RESPONSE_429, body=None) + ) + assert not _anthropic_is_retryable(ValueError("malformed tool message")) + + def test_openai_deterministic_errors_are_not_retryable(self): + status_error = OpenAIStatusError( + "rate limited", response=_RESPONSE_429, body=None + ) + assert not _openai_is_retryable(status_error) + assert not _openai_compat_is_retryable(status_error) + assert not _openai_is_retryable(ValueError("malformed tool message")) + + def test_gemini_status_error_is_not_retryable(self): + assert not _gemini_is_retryable(GenaiAPIError(429, None)) + assert not _gemini_is_retryable(ValueError("malformed request")) + + def test_bedrock_deterministic_errors_are_not_retryable(self): + assert not _bedrock_is_retryable(_throttling_client_error()) + assert not _bedrock_is_retryable(NoCredentialsError()) + assert not _bedrock_is_retryable(ValueError("unknown model family")) diff --git a/services/ai/tests/unit/test_stream_retry.py b/services/ai/tests/unit/test_stream_retry.py new file mode 100644 index 00000000..78721627 --- /dev/null +++ b/services/ai/tests/unit/test_stream_retry.py @@ -0,0 +1,259 @@ +"""Unit tests for provider stream recovery in +``streaming.generate.event_stream_with_context_retry``. + +Covers the one-shot retry for transport-level failures flagged by the provider +(e.g. an ``httpx.ReadError`` dropping the SSE connection before any content is +sent), the no-retry guards, and the pre-existing context-overflow compaction +retry. +""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock, patch + +import httpx +import pytest + +from providers import ProviderError, ProviderType +from streaming.generate import event_stream_with_context_retry + + +def _evt(type_: str): + # Only ``.type`` is inspected by the code under test. + return SimpleNamespace(type=type_) + + +class _PassthroughTracker: + """UsageTracker stand-in that passes events through and counts saves.""" + + save_count = 0 + + def __init__(self, *args, **kwargs): + pass + + async def wrap_stream(self, stream): + async for event in stream: + yield event + + def save(self): + _PassthroughTracker.save_count += 1 + + +class _ScriptedProvider: + """LLMProvider whose stream_response plays a scripted sequence per call.""" + + supports_citations = True + provider_type = ProviderType.OPENAI_COMPATIBLE + model_record_id = "model-1" + model_name = "test-model" + PERSISTED_BLOCK_EXTRAS = () + + def __init__(self, scripts): + self._scripts = list(scripts) + self.calls = 0 + + async def stream_response(self, **kwargs): + script = self._scripts[self.calls] + self.calls += 1 + async for event in script(): + yield event + + +def _provider_error( + message: str, + status_code: int | None, + *, + is_retryable: bool = False, + is_context_overflow: bool = False, + cause: BaseException | None = None, +) -> ProviderError: + return ProviderError( + message, + provider_type=ProviderType.OPENAI_COMPATIBLE, + model="test-model", + status_code=status_code, + is_retryable=is_retryable, + is_context_overflow=is_context_overflow, + cause=cause, + ) + + +# --- scripted provider streams -------------------------------------------- + + +async def _drop_mid_stream(): + """Emit the envelope, then the connection dies before any content. + + Shaped like what ``OpenAICompatibleProvider`` actually raises: a + status-less ProviderError with a transport cause and the retryable flag. + """ + yield _evt("message_start") + raise _provider_error( + "connection lost", + status_code=None, + is_retryable=True, + cause=httpx.ReadError( + "stream reset", request=httpx.Request("POST", "http://x") + ), + ) + + +async def _local_error(): + # No status AND not flagged transport-level (e.g. a local validation + # ValueError wrapped by the provider). Must not be retried. + raise _provider_error( + "invalid tool message sequence", + status_code=None, + is_retryable=False, + cause=ValueError("tool_calls were never answered"), + ) + yield # unreachable; marks this as an async generator + + +async def _rate_limited(): + # Status errors are the provider SDK's retry territory; the outer layer + # must not retry them again. + raise _provider_error("rate limited", status_code=429, is_retryable=False) + yield # unreachable; marks this as an async generator + + +async def _context_overflow(): + raise _provider_error("prompt too long", status_code=413, is_context_overflow=True) + yield # unreachable; marks this as an async generator + + +async def _drop_after_content_block_start(): + yield _evt("message_start") + yield _evt("content_block_start") + raise _provider_error( + "connection lost", + status_code=None, + is_retryable=True, + cause=httpx.ReadError( + "stream reset", request=httpx.Request("POST", "http://x") + ), + ) + + +async def _ok_text_stream(): + yield _evt("message_start") + yield _evt("content_block_start") + yield _evt("content_block_delta") + yield _evt("content_block_stop") + yield _evt("message_stop") + + +_FULL_OK_TYPES = [ + "message_start", + "content_block_start", + "content_block_delta", + "content_block_stop", + "message_stop", +] + + +async def _drain(provider, conversation_messages=None): + if conversation_messages is None: + conversation_messages = [{"role": "user", "content": "hi"}] + compactor = SimpleNamespace( + select_legacy_compaction_split=lambda msgs: None, + compact_conversation=AsyncMock( + return_value=[{"role": "user", "content": "compacted"}] + ), + ) + sleep_mock = AsyncMock() + _PassthroughTracker.save_count = 0 + events = [] + with ( + patch("streaming.generate.UsageTracker", _PassthroughTracker), + patch("streaming.generate.UsageRepository", lambda: None), + patch("asyncio.sleep", sleep_mock), + ): + stream = event_stream_with_context_retry( + turn_tools=[], + conversation_messages=conversation_messages, + llm_provider=provider, + chat_user_id="user-1", + chat_id="chat-1", + system_prompt="sys", + compactor=compactor, + latest_compaction_summary=None, + summarizer_context_window_tokens=100_000, + ) + async for event in stream: + events.append(event) + return events, conversation_messages, compactor, sleep_mock + + +@pytest.mark.unit +class TestTransientStreamRetry: + @pytest.mark.asyncio + async def test_mid_stream_drop_retries_once_and_recovers(self): + provider = _ScriptedProvider([_drop_mid_stream, _ok_text_stream]) + events, _, _, sleep_mock = await _drain(provider) + + assert provider.calls == 2 + types = [e.type for e in events] + # The original envelope is retained, the retried stream's duplicate + # envelope is dropped, and the recovered content follows it. Exactly + # one message_start reaches the caller, so persistence creates exactly + # one assistant row. + assert types.count("message_start") == 1 + assert types == _FULL_OK_TYPES + # One backoff sleep before the retry, and the failed attempt's usage + # was persisted (both attempts count toward the upsert). + sleep_mock.assert_awaited_once_with(1.0) + assert _PassthroughTracker.save_count == 2 + + @pytest.mark.asyncio + async def test_non_retryable_statusless_error_is_not_retried(self): + # No HTTP status alone does not make an error transient: a local + # validation error wrapped by the provider must not be replayed. + provider = _ScriptedProvider([_local_error]) + with pytest.raises(ProviderError) as exc: + await _drain(provider) + + assert provider.calls == 1 + assert exc.value.status_code is None + assert exc.value.is_retryable is False + + @pytest.mark.asyncio + async def test_status_error_not_retried_by_outer_layer(self): + # 429/5xx are retried by the provider SDK at request-creation time; + # retrying them here too would multiply the request count. + provider = _ScriptedProvider([_rate_limited]) + with pytest.raises(ProviderError) as exc: + await _drain(provider) + + assert provider.calls == 1 + assert exc.value.status_code == 429 + + @pytest.mark.asyncio + async def test_no_retry_after_content_block_started(self): + provider = _ScriptedProvider([_drop_after_content_block_start]) + with pytest.raises(ProviderError): + await _drain(provider) + + # Once a content block has started, the run cannot be cleanly redone. + assert provider.calls == 1 + + @pytest.mark.asyncio + async def test_second_failure_is_not_retried_again(self): + provider = _ScriptedProvider([_drop_mid_stream, _drop_mid_stream]) + with pytest.raises(ProviderError): + await _drain(provider) + + assert provider.calls == 2 + + @pytest.mark.asyncio + async def test_context_overflow_still_compacts_and_retries(self): + provider = _ScriptedProvider([_context_overflow, _ok_text_stream]) + conversation_messages = [{"role": "user", "content": "hi"}] + events, conversation_messages, compactor, _ = await _drain( + provider, conversation_messages + ) + + assert provider.calls == 2 + assert compactor.compact_conversation.await_count == 1 + assert conversation_messages == [{"role": "user", "content": "compacted"}] + # Attempt one failed at request time, so attempt two's envelope is kept. + assert [e.type for e in events] == _FULL_OK_TYPES