feat: add asynchronous reconnecting SSE connection

OmarAlJarrah · claude · OmarAlJarrah · commit a5c3200c0a8d · 2026-06-09T17:47:36.000+03:00
Co-Authored-By: Claude Opus 4.8 &lt;noreply@anthropic.com&gt;
diff --git a/packages/dexpace-sdk-core/src/dexpace/sdk/core/http/sse/connection.py b/packages/dexpace-sdk-core/src/dexpace/sdk/core/http/sse/connection.py
@@ -22,19 +22,23 @@
 
 from ...errors import HttpResponseError, ServiceRequestError, ServiceResponseError
 from ...pipeline.dispatch import (
+    AsyncPipelineLike,
+    SendAsync,
     SendSync,
     SyncPipelineLike,
 )
-from ...util.clock import SYSTEM_CLOCK, Clock
+from ...util.clock import ASYNC_SYSTEM_CLOCK, SYSTEM_CLOCK, AsyncClock, Clock
 from ..context.dispatch_context import DispatchContext
-from .parser import SseParser
+from ..response.async_response_body import _shielded_cleanup
+from .parser import SseParser, parse_async_events
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Generator, Iterator
+    from collections.abc import AsyncIterator, Callable, Generator, Iterator
     from types import TracebackType
     from typing import Self
 
     from ..request.request import Request
+    from ..response.async_response import AsyncResponse
     from ..response.response import Response
     from .parser import SseEvent
 
@@ -245,4 +249,145 @@ def _stream(self, response: Response) -> Generator[SseEvent, None, bool]:
         return progressed
 
 
-__all__ = ["SseConnection"]
+class AsyncSseConnection:
+    """Asynchronous twin of :class:`SseConnection`.
+
+    Mirrors the sync client exactly with ``async`` iteration semantics. Each
+    (re)connection is dispatched through ``source`` — an ``AsyncPipeline`` (run
+    with a fresh dispatch context per connection) or an async
+    ``Request -> AsyncResponse`` callable. The response is closed through the
+    shielded-cleanup convention, so a cancelled consumer still releases the
+    transport handle before the cancellation continues to propagate.
+
+    Args:
+        source: An async pipeline or an async send-callable.
+        initial_request: The request opening the stream; reused for every
+            reconnection with an updated ``Last-Event-ID``.
+        last_event_id: Seed id to resume a previously-interrupted stream.
+        default_retry: Backoff base (seconds) used until the server sends a
+            ``retry:`` value.
+        max_backoff: Ceiling on any single reconnect delay.
+        max_reconnects: Maximum consecutive failed reconnections before
+            iteration raises. ``None`` reconnects indefinitely.
+        jitter: Upward jitter fraction applied to the backoff.
+        clock: Async time source for backoff sleeps (injected for tests).
+        rand: RNG for jitter (injected for tests).
+        dispatch_factory: Builds the dispatch context per connection when
+            ``source`` is a pipeline. Defaults to ``DispatchContext.noop``.
+    """
+
+    __slots__ = (
+        "_clock",
+        "_dispatch_factory",
+        "_initial",
+        "_jitter",
+        "_last_event_id",
+        "_max_backoff",
+        "_max_reconnects",
+        "_rand",
+        "_response",
+        "_retry_base",
+        "_send",
+    )
+
+    def __init__(
+        self,
+        source: AsyncPipelineLike | SendAsync,
+        initial_request: Request,
+        *,
+        last_event_id: str | None = None,
+        default_retry: float = _DEFAULT_RETRY_SECONDS,
+        max_backoff: float = _DEFAULT_MAX_BACKOFF,
+        max_reconnects: int | None = None,
+        jitter: float = _DEFAULT_JITTER,
+        clock: AsyncClock = ASYNC_SYSTEM_CLOCK,
+        rand: random.Random | None = None,
+        dispatch_factory: Callable[[], DispatchContext] | None = None,
+    ) -> None:
+        self._initial = initial_request
+        self._last_event_id = last_event_id
+        self._retry_base = default_retry
+        self._max_backoff = max_backoff
+        self._max_reconnects = max_reconnects
+        self._jitter = jitter
+        self._clock = clock
+        self._rand = rand if rand is not None else random.Random()
+        self._dispatch_factory = dispatch_factory or DispatchContext.noop
+        self._response: AsyncResponse | None = None
+        self._send = self._normalise(source)
+
+    def _normalise(self, source: AsyncPipelineLike | SendAsync) -> SendAsync:
+        if isinstance(source, AsyncPipelineLike):
+            pipeline = source
+
+            async def send(request: Request) -> AsyncResponse:
+                return await pipeline.run(request, self._dispatch_factory())
+
+            return send
+        return source
+
+    async def aclose(self) -> None:
+        """Close the current response, if any. Idempotent and cancel-safe."""
+        response = self._response
+        if response is None:
+            return
+        self._response = None
+        await _shielded_cleanup(response.close())
+
+    async def __aenter__(self) -> Self:
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        await self.aclose()
+
+    def __aiter__(self) -> AsyncIterator[SseEvent]:
+        return self._events()
+
+    async def _events(self) -> AsyncIterator[SseEvent]:
+        failures = 0
+        try:
+            while True:
+                response = await self._connect(_resume_request(self._initial, self._last_event_id))
+                progressed = False
+                body = response.body
+                if body is not None:
+                    stream = parse_async_events(body.aiter_bytes())
+                    try:
+                        async with stream:
+                            async for event in stream:
+                                self._last_event_id = _last_event_id_of(event, self._last_event_id)
+                                progressed = True
+                                yield event
+                    except _TRANSIENT:
+                        pass
+                    if stream.retry is not None:
+                        self._retry_base = stream.retry / 1000.0
+                await self.aclose()
+                if progressed:
+                    failures = 0
+                if self._max_reconnects is not None and failures >= self._max_reconnects:
+                    raise ServiceResponseError("SSE reconnect budget exhausted")
+                await self._clock.sleep(
+                    _next_backoff(
+                        self._retry_base, failures, self._max_backoff, self._jitter, self._rand
+                    )
+                )
+                failures += 1
+        finally:
+            await self.aclose()
+
+    async def _connect(self, request: Request) -> AsyncResponse:
+        response = await self._send(request)
+        self._response = response
+        if not response.status.is_success:
+            await self.aclose()
+            raise HttpResponseError(response=response)
+        return response
+
+
+__all__ = ["AsyncSseConnection", "SseConnection"]
diff --git a/packages/dexpace-sdk-core/src/dexpace/sdk/core/http/sse/parser.py b/packages/dexpace-sdk-core/src/dexpace/sdk/core/http/sse/parser.py
@@ -331,6 +331,14 @@ async def __anext__(self) -> SseEvent:
             self._parser.feed(chunk)
             self._pending = self._parser.drain()
 
+    @property
+    def retry(self) -> int | None:
+        """The most recent ``retry:`` value seen by the underlying parser, in
+        milliseconds, or ``None``. Mirrors :attr:`SseParser.retry` so a
+        reconnecting client can read the server's sticky hint after iteration.
+        """
+        return self._parser.retry
+
     async def aclose(self) -> None:
         """Release the upstream byte stream. Idempotent.
 
diff --git a/packages/dexpace-sdk-core/tests/sse/test_async_connection.py b/packages/dexpace-sdk-core/tests/sse/test_async_connection.py
@@ -0,0 +1,147 @@
+# Copyright (c) 2026 dexpace and Omar Aljarrah.
+# Licensed under the MIT License. See LICENSE.md in the repository root for details.
+
+"""Tests for the asynchronous reconnecting SSE client."""
+
+from __future__ import annotations
+
+import asyncio
+import random
+from collections.abc import AsyncIterator
+
+import pytest
+
+from dexpace.sdk.core.errors import HttpResponseError, ServiceResponseError
+from dexpace.sdk.core.http.common import MediaType, Protocol, Url
+from dexpace.sdk.core.http.request import Method, Request
+from dexpace.sdk.core.http.response import AsyncResponse, AsyncResponseBody, Status
+from dexpace.sdk.core.http.sse.connection import AsyncSseConnection
+
+
+def _request() -> Request:
+    return Request(method=Method.GET, url=Url.parse("https://api.example.com/stream"))
+
+
+class _AsyncDropBody(AsyncResponseBody):
+    """Yields scripted chunks, optionally then raises, optionally blocks."""
+
+    def __init__(
+        self,
+        chunks: list[bytes],
+        *,
+        error: BaseException | None = None,
+        block: bool = False,
+    ) -> None:
+        self._chunks = chunks
+        self._error = error
+        self._block = block
+        self.closed = False
+
+    def media_type(self) -> MediaType | None:
+        return MediaType.parse("text/event-stream")
+
+    def content_length(self) -> int:
+        return -1
+
+    async def aiter_bytes(self, chunk_size: int = 64 * 1024) -> AsyncIterator[bytes]:
+        for chunk in self._chunks:
+            yield chunk
+        if self._block:
+            await asyncio.Event().wait()  # park until cancelled
+        if self._error is not None:
+            raise self._error
+
+    async def close(self) -> None:
+        self.closed = True
+
+
+def _response(body: AsyncResponseBody | None, *, status: Status = Status.OK) -> AsyncResponse:
+    return AsyncResponse(
+        request=_request(),
+        protocol=Protocol.HTTP_1_1,
+        status=status,
+        body=body,
+    )
+
+
+class _AsyncScript:
+    def __init__(self, responses: list[AsyncResponse]) -> None:
+        self._responses = list(responses)
+        self.requests: list[Request] = []
+
+    async def __call__(self, request: Request) -> AsyncResponse:
+        self.requests.append(request)
+        if not self._responses:
+            raise AssertionError("send called more times than scripted")
+        return self._responses.pop(0)
+
+
+class _RecordingAsyncClock:
+    def __init__(self) -> None:
+        self.sleeps: list[float] = []
+
+    def now(self) -> float:
+        return 0.0
+
+    def monotonic(self) -> float:
+        return 0.0
+
+    async def sleep(self, duration: float) -> None:
+        self.sleeps.append(duration)
+
+
+class _LowJitter(random.Random):
+    def uniform(self, a: float, b: float) -> float:
+        return a
+
+
+async def test_async_yields_events_and_reconnects_with_replay() -> None:
+    dropped = _AsyncDropBody([b"id: 9\ndata: one\n\n"], error=ServiceResponseError("drop"))
+    resumed = _AsyncDropBody([b"data: two\n\n"])
+    script = _AsyncScript([_response(dropped), _response(resumed)])
+    conn = AsyncSseConnection(script, _request(), clock=_RecordingAsyncClock(), rand=_LowJitter())
+
+    received: list[str] = []
+    async with conn as events:
+        async for event in events:
+            received.append(event.data)
+            if len(received) == 2:
+                break
+
+    assert received == ["one", "two"]
+    assert script.requests[1].headers.get("last-event-id") == "9"
+
+
+async def test_async_non_success_status_raises() -> None:
+    script = _AsyncScript([_response(None, status=Status.BAD_GATEWAY)])
+    conn = AsyncSseConnection(script, _request(), clock=_RecordingAsyncClock(), rand=_LowJitter())
+
+    with pytest.raises(HttpResponseError):
+        async for _ in conn:
+            pass
+    assert len(script.requests) == 1
+
+
+async def test_async_cancellation_propagates_and_closes() -> None:
+    body = _AsyncDropBody([b"data: hi\n\n"], block=True)
+    script = _AsyncScript([_response(body)])
+    conn = AsyncSseConnection(script, _request(), clock=_RecordingAsyncClock(), rand=_LowJitter())
+
+    seen: list[str] = []
+
+    async def consume() -> None:
+        async with conn as events:
+            async for event in events:
+                seen.append(event.data)
+
+    task = asyncio.ensure_future(consume())
+    for _ in range(10):
+        await asyncio.sleep(0)
+        if seen:
+            break
+    task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await task
+
+    assert seen == ["hi"]
+    assert body.closed is True
