# This file was auto-generated by Fern from our API Definition. import codecs import re import time from contextlib import asynccontextmanager, contextmanager from typing import ( Any, AsyncContextManager, AsyncGenerator, AsyncIterator, Callable, ContextManager, Iterator, Optional, ) import anyio import httpx from ._decoders import SSEDecoder from ._exceptions import SSEError from ._models import ServerSentEvent MAX_LINE_SIZE: int = 1_048_576 # 1 MiB # Reconnection defaults, mirroring the TypeScript SDK's Stream implementation. DEFAULT_MAX_RECONNECTION_ATTEMPTS: int = 5 DEFAULT_RECONNECT_DELAY_MS: int = 1_000 MAX_RECONNECT_DELAY_MS: int = 30_000 # A reconnect callback re-issues the original request (with a ``Last-Event-ID`` # header set to the supplied event id) and returns a *context manager* yielding # a fresh streaming ``httpx.Response``. Sync clients supply a sync context # manager; async clients supply an async one. class EventSource: def __init__( self, response: httpx.Response, *, resumable: bool = False, stream_reconnection_enabled: bool = True, max_stream_reconnection_attempts: Optional[int] = None, stream_terminator: Optional[str] = None, reconnect: Optional[Callable[[str], Any]] = None, ) -> None: self._response = response self._resumable = resumable self._stream_reconnection_enabled = stream_reconnection_enabled self._max_stream_reconnection_attempts = max_stream_reconnection_attempts self._stream_terminator = stream_terminator self._reconnect = reconnect @staticmethod def _is_event_stream(response: httpx.Response) -> bool: content_type = response.headers.get("content-type", "").partition(";")[0] return "text/event-stream" in content_type def _check_content_type(self) -> None: if not self._is_event_stream(self._response): content_type = self._response.headers.get("content-type", "").partition(";")[0] raise SSEError( f"Expected response header Content-Type to contain 'text/event-stream', got {content_type!r}" ) def _is_reconnect_response_usable(self, response: httpx.Response) -> bool: """Whether a reconnected response can be resumed as an SSE stream. ``httpx.stream`` does not raise on non-success status, so a resume that returns an error page (e.g. ``200 text/html`` or a ``500`` body) would otherwise be parsed as SSE and yield garbage/zero events. Such a response is treated as a failed attempt (back off and retry) instead. """ return response.status_code < 400 and self._is_event_stream(response) def _get_charset(self, response: Optional[httpx.Response] = None) -> str: """Extract charset from Content-Type header, fallback to UTF-8.""" resolved = response if response is not None else self._response content_type = resolved.headers.get("content-type", "") # Parse charset parameter using regex charset_match = re.search(r"charset=([^;\s]+)", content_type, re.IGNORECASE) if charset_match: charset = charset_match.group(1).strip("\"'") # Validate that it's a known encoding try: # Test if the charset is valid by trying to encode/decode "test".encode(charset).decode(charset) return charset except (LookupError, UnicodeError): # If charset is invalid, fall back to UTF-8 pass # Default to UTF-8 if no charset specified or invalid charset return "utf-8" @property def response(self) -> httpx.Response: return self._response @staticmethod def _normalize_sse_line_endings(buf: str) -> str: """Normalize line endings per the SSE spec (\\r\\n → \\n, bare \\r → \\n). A trailing \\r is preserved because it may pair with a leading \\n in the next chunk to form a single \\r\\n terminator. """ buf = buf.replace("\r\n", "\n") if buf.endswith("\r"): return buf[:-1].replace("\r", "\n") + "\r" return buf.replace("\r", "\n") def _new_text_decoder(self, response: Optional[httpx.Response] = None) -> "codecs.IncrementalDecoder": return codecs.getincrementaldecoder(self._get_charset(response))(errors="replace") def _reconnect_applicable(self) -> bool: """Whether reconnection is configured for this stream at all. This is the terminator-gating half of the reconnect decision, kept separate from :meth:`_should_reconnect` (which additionally requires a last *dispatched* id and an unexhausted attempt budget). The split lets a mid-stream transport error terminate consistently: - a stream that can never reconnect (non-resumable, no terminator, disabled, or no callback) must re-raise the error to the caller, so a truncated stream is not mistaken for a clean completion; - a resumable stream that has merely run out of attempts (or has no id to resume from) ends cleanly — the same way an exhausted empty/error -body resume already does, matching the TypeScript ``return``. """ return ( self._resumable and self._stream_terminator is not None and self._stream_reconnection_enabled and self._reconnect is not None ) def _should_reconnect(self, last_dispatched_id: Optional[str], reconnect_attempts: int) -> bool: """Decide whether a prematurely-ended stream should be reconnected. Mirrors the TypeScript ``shouldReconnect`` gating: - only resumable SSE endpoints with a configured terminator, reconnect enabled, and a reconnect callback are eligible (see :meth:`_reconnect_applicable`); - a last *dispatched* event id must exist to resume from; - the consecutive-failed-attempt cap must not be exceeded. """ if not self._reconnect_applicable(): return False if not last_dispatched_id: return False max_attempts = ( self._max_stream_reconnection_attempts if self._max_stream_reconnection_attempts is not None else DEFAULT_MAX_RECONNECTION_ATTEMPTS ) if reconnect_attempts >= max_attempts: return False return True def _reconnect_delay_seconds(self, last_retry: Optional[int]) -> float: """Backoff before a reconnect. Uses the server's most recent ``retry:`` directive (milliseconds) when present, otherwise a default of ``DEFAULT_RECONNECT_DELAY_MS``, clamped to ``MAX_RECONNECT_DELAY_MS``. """ base_ms = last_retry if (last_retry is not None and last_retry > 0) else DEFAULT_RECONNECT_DELAY_MS return min(base_ms, MAX_RECONNECT_DELAY_MS) / 1000.0 def _sleep_before_reconnect(self, last_retry: Optional[int]) -> None: # ``time.sleep`` blocks the calling thread but remains interruptible by # signals (e.g. ``KeyboardInterrupt``), which propagate out and abort # the reconnect without issuing another request. time.sleep(self._reconnect_delay_seconds(last_retry)) async def _asleep_before_reconnect(self, last_retry: Optional[int]) -> None: # ``anyio.sleep`` is cancellation-aware: if the consumer cancels the task # or closes the async generator mid-delay, this raises (and no further # request is issued) instead of blocking for the whole interval. await anyio.sleep(self._reconnect_delay_seconds(last_retry)) def _decode_response( self, response: httpx.Response, decoder: SSEDecoder, text_decoder: "codecs.IncrementalDecoder", ) -> Iterator[ServerSentEvent]: buf = "" for chunk in response.iter_bytes(): buf += text_decoder.decode(chunk) buf = self._normalize_sse_line_endings(buf) while "\n" in buf: line, buf = buf.split("\n", 1) sse = decoder.decode(line) if sse is not None: yield sse if len(buf) > MAX_LINE_SIZE: raise SSEError( f"SSE line exceeded maximum size of {MAX_LINE_SIZE} characters without encountering a newline" ) yield from self._flush_decoder(buf, decoder, text_decoder) async def _adecode_response( self, response: httpx.Response, decoder: SSEDecoder, text_decoder: "codecs.IncrementalDecoder", ) -> AsyncGenerator[ServerSentEvent, None]: buf = "" async for chunk in response.aiter_bytes(): buf += text_decoder.decode(chunk) buf = self._normalize_sse_line_endings(buf) while "\n" in buf: line, buf = buf.split("\n", 1) sse = decoder.decode(line) if sse is not None: yield sse if len(buf) > MAX_LINE_SIZE: raise SSEError( f"SSE line exceeded maximum size of {MAX_LINE_SIZE} characters without encountering a newline" ) for sse in self._flush_decoder(buf, decoder, text_decoder): yield sse def _flush_decoder( self, buf: str, decoder: SSEDecoder, text_decoder: "codecs.IncrementalDecoder", ) -> Iterator[ServerSentEvent]: # Flush any remaining bytes from the incremental decoder buf += text_decoder.decode(b"", final=True) buf = buf.replace("\r\n", "\n").replace("\r", "\n") if len(buf) > MAX_LINE_SIZE: raise SSEError( f"SSE line exceeded maximum size of {MAX_LINE_SIZE} characters without encountering a newline" ) while "\n" in buf: line, buf = buf.split("\n", 1) sse = decoder.decode(line) if sse is not None: yield sse if buf.strip(): sse = decoder.decode(buf) if sse is not None: yield sse def iter_sse(self) -> Iterator[ServerSentEvent]: self._check_content_type() decoder = SSEDecoder() text_decoder = self._new_text_decoder() last_dispatched_id: Optional[str] = None last_retry: Optional[int] = None # Consecutive failed reconnection attempts. Reset to 0 whenever an event # is successfully dispatched (reset-on-progress) — matching browser # `EventSource` semantics: a server that emits >=1 event then drops on # every connection can reconnect indefinitely. reconnect_attempts = 0 # ``None`` means there is no live stream to read this iteration (e.g. a # failed reconnect); the loop then re-evaluates the reconnect decision # without re-reading an exhausted response. response: Optional[httpx.Response] = self._response # Context manager for a response we opened ourselves and must close. # The initial response is owned by the caller, so it starts as None. owned_cm: Optional[ContextManager[httpx.Response]] = None try: while True: if response is not None: events = self._decode_response(response, decoder, text_decoder) while True: try: sse = next(events) except StopIteration: break except SSEError: # A protocol violation (e.g. an oversized line) is a # genuine error, not a dropped connection; propagate it. # Listed first because ``SSEError`` subclasses # ``httpx.TransportError``. raise except httpx.TransportError: # A transport error mid-stream (e.g. the server dropped # the connection: ``ReadError``/``RemoteProtocolError``) # is a premature end. Only swallow it when reconnection # is configured for this stream; otherwise re-raise so a # non-resumable stream still surfaces the error to the # caller instead of looking like a clean completion. # When reconnection is applicable but the attempt budget # is exhausted, we ``break`` and end cleanly below — the # same way an exhausted empty/error-body resume does, so # give-up is consistent regardless of failure shape. # ``next`` is used rather than ``for`` so this cannot # swallow a ``GeneratorExit`` raised at a ``yield``. if not self._reconnect_applicable(): raise break yield sse if sse.id: last_dispatched_id = sse.id if sse.retry is not None: last_retry = sse.retry reconnect_attempts = 0 if not self._should_reconnect(last_dispatched_id, reconnect_attempts): return reconnect_attempts += 1 self._sleep_before_reconnect(last_retry) # Close the previously-opened reconnect response before opening # a new one so we never hold more than one extra connection. if owned_cm is not None: owned_cm.__exit__(None, None, None) owned_cm = None assert self._reconnect is not None # guaranteed by _should_reconnect try: cm: ContextManager[httpx.Response] = self._reconnect(last_dispatched_id or "") new_response = cm.__enter__() except Exception: # A failed reconnect consumes an attempt; back off and retry. response = None continue owned_cm = cm if new_response is None or not self._is_reconnect_response_usable(new_response): # Null/empty body or a non-SSE/error response (e.g. 204/304, # a 500, or an HTML error page): treat as a failed attempt. response = None continue response = new_response # Drop any partial event left over from the dropped stream, but # keep the last event id (per the SSE spec) and start a fresh # incremental text decoder for the new connection. decoder.reset_in_progress_event() text_decoder = self._new_text_decoder(new_response) finally: if owned_cm is not None: owned_cm.__exit__(None, None, None) async def aiter_sse(self) -> AsyncGenerator[ServerSentEvent, None]: self._check_content_type() decoder = SSEDecoder() text_decoder = self._new_text_decoder() last_dispatched_id: Optional[str] = None last_retry: Optional[int] = None reconnect_attempts = 0 response: Optional[httpx.Response] = self._response owned_cm: Optional[AsyncContextManager[httpx.Response]] = None try: while True: if response is not None: events = self._adecode_response(response, decoder, text_decoder) while True: try: sse = await events.__anext__() except StopAsyncIteration: break except SSEError: # A protocol violation (e.g. an oversized line) is a # genuine error, not a dropped connection; propagate it. # Listed first because ``SSEError`` subclasses # ``httpx.TransportError``. raise except httpx.TransportError: # A transport error mid-stream (e.g. the server dropped # the connection: ``ReadError``/``RemoteProtocolError``) # is a premature end. Only swallow it when reconnection # is configured for this stream; otherwise re-raise so a # non-resumable stream still surfaces the error to the # caller instead of looking like a clean completion. # When reconnection is applicable but the attempt budget # is exhausted, we ``break`` and end cleanly below — the # same way an exhausted empty/error-body resume does, so # give-up is consistent regardless of failure shape. if not self._reconnect_applicable(): raise break yield sse if sse.id: last_dispatched_id = sse.id if sse.retry is not None: last_retry = sse.retry reconnect_attempts = 0 if not self._should_reconnect(last_dispatched_id, reconnect_attempts): return reconnect_attempts += 1 await self._asleep_before_reconnect(last_retry) if owned_cm is not None: await owned_cm.__aexit__(None, None, None) owned_cm = None assert self._reconnect is not None # guaranteed by _should_reconnect try: cm: AsyncContextManager[httpx.Response] = self._reconnect(last_dispatched_id or "") new_response = await cm.__aenter__() except Exception: response = None continue owned_cm = cm if new_response is None or not self._is_reconnect_response_usable(new_response): response = None continue response = new_response decoder.reset_in_progress_event() text_decoder = self._new_text_decoder(new_response) finally: if owned_cm is not None: # Shield the close so a cancellation delivered while reading a # reconnected response still fully tears the connection down # instead of leaking it until the client is closed. with anyio.CancelScope(shield=True): await owned_cm.__aexit__(None, None, None) @contextmanager def connect_sse(client: httpx.Client, method: str, url: str, **kwargs: Any) -> Iterator[EventSource]: headers = kwargs.pop("headers", {}) headers["Accept"] = "text/event-stream" headers["Cache-Control"] = "no-store" with client.stream(method, url, headers=headers, **kwargs) as response: yield EventSource(response) @asynccontextmanager async def aconnect_sse( client: httpx.AsyncClient, method: str, url: str, **kwargs: Any, ) -> AsyncIterator[EventSource]: headers = kwargs.pop("headers", {}) headers["Accept"] = "text/event-stream" headers["Cache-Control"] = "no-store" async with client.stream(method, url, headers=headers, **kwargs) as response: yield EventSource(response)