feat: add SSE auto-reconnect for Axon subscriptions

- Add AxonSubscribeSseParams with after_sequence parameter (internal use only)
- Wrap subscribe_sse() with ReconnectingStream/AsyncReconnectingStream
- Automatically resume from last received event on timeout using sequence-based offset
- Add opt-out via RAW_RESPONSE_HEADER for users who want raw streams
- after_sequence is handled internally by reconnector, not exposed in public API

Per code review feedback from @dines-rl on TypeScript PR #765:
'after_sequence should just be a variable in the reconnector for follow-up'

Part of porting TypeScript PR #765 features to Python SDK.

Author: tode-rl

src/runloop_api_client/resources/axons/axons.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import Optional
+from typing import Optional, cast
 from typing_extensions import Literal
 
 import httpx
@@ -16,6 +16,7 @@
     AsyncSqlResourceWithStreamingResponse,
 )
 from ...types import axon_list_params, axon_create_params, axon_publish_params
+from ...types.axons import axon_subscribe_sse_params
 from ..._types import Body, Omit, Query, Headers, NotGiven, omit, not_given
 from ..._utils import path_template, maybe_transform, async_maybe_transform
 from ..._compat import cached_property
@@ -26,7 +27,8 @@
     async_to_raw_response_wrapper,
     async_to_streamed_response_wrapper,
 )
-from ..._streaming import Stream, AsyncStream
+from ..._constants import RAW_RESPONSE_HEADER
+from ..._streaming import Stream, AsyncStream, ReconnectingStream, AsyncReconnectingStream
 from ...pagination import SyncAxonsCursorIDPage, AsyncAxonsCursorIDPage
 from ..._base_client import AsyncPaginator, make_request_options
 from ...types.axon_view import AxonView
@@ -269,6 +271,8 @@ def subscribe_sse(
         """
         [Beta] Subscribe to an axon event stream via server-sent events.
 
+        Automatically reconnects on timeout, resuming from last received event.
+
         Args:
           extra_headers: Send extra headers
 
@@ -282,14 +286,54 @@ def subscribe_sse(
             raise ValueError(f"Expected a non-empty value for `id` but received {id!r}")
         default_headers: Headers = {"Accept": "text/event-stream"}
         merged_headers = default_headers if extra_headers is None else {**default_headers, **extra_headers}
-        return self._get(
-            path_template("/v1/axons/{id}/subscribe/sse", id=id),
-            options=make_request_options(
-                extra_headers=merged_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
+
+        # Check if user wants raw response (opt-out of reconnection)
+        if extra_headers is not None and RAW_RESPONSE_HEADER in extra_headers:
+            return self._get(
+                path_template("/v1/axons/{id}/subscribe/sse", id=id),
+                options=make_request_options(
+                    extra_headers=merged_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
+                ),
+                cast_to=AxonEventView,
+                stream=True,
+                stream_cls=Stream[AxonEventView],
+            )
+
+        def create_stream(last_sequence: str | None) -> Stream[AxonEventView]:
+            # after_sequence is used internally for reconnection only
+            sequence_int = int(last_sequence) if last_sequence is not None else None
+            return self._get(
+                path_template("/v1/axons/{id}/subscribe/sse", id=id),
+                options=make_request_options(
+                    extra_headers=merged_headers,
+                    extra_query=extra_query,
+                    extra_body=extra_body,
+                    timeout=timeout,
+                    query=maybe_transform(
+                        {"after_sequence": sequence_int},
+                        axon_subscribe_sse_params.AxonSubscribeSseParams,
+                    ),
+                ),
+                cast_to=AxonEventView,
+                stream=True,
+                stream_cls=Stream[AxonEventView],
+            )
+
+        initial_stream = create_stream(None)
+
+        def get_sequence(item: AxonEventView) -> str | None:
+            value = getattr(item, "sequence", None)
+            if value is None:
+                return None
+            return str(value)
+
+        return cast(
+            Stream[AxonEventView],
+            ReconnectingStream(
+                current_stream=initial_stream,
+                stream_creator=create_stream,
+                get_offset=get_sequence,
             ),
-            cast_to=AxonEventView,
-            stream=True,
-            stream_cls=Stream[AxonEventView],
         )
 
 
@@ -526,6 +570,8 @@ async def subscribe_sse(
         """
         [Beta] Subscribe to an axon event stream via server-sent events.
 
+        Automatically reconnects on timeout, resuming from last received event.
+
         Args:
           extra_headers: Send extra headers
 
@@ -539,14 +585,54 @@ async def subscribe_sse(
             raise ValueError(f"Expected a non-empty value for `id` but received {id!r}")
         default_headers: Headers = {"Accept": "text/event-stream"}
         merged_headers = default_headers if extra_headers is None else {**default_headers, **extra_headers}
-        return await self._get(
-            path_template("/v1/axons/{id}/subscribe/sse", id=id),
-            options=make_request_options(
-                extra_headers=merged_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
+
+        # Check if user wants raw response (opt-out of reconnection)
+        if extra_headers is not None and RAW_RESPONSE_HEADER in extra_headers:
+            return await self._get(
+                path_template("/v1/axons/{id}/subscribe/sse", id=id),
+                options=make_request_options(
+                    extra_headers=merged_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
+                ),
+                cast_to=AxonEventView,
+                stream=True,
+                stream_cls=AsyncStream[AxonEventView],
+            )
+
+        async def create_stream(last_sequence: str | None) -> AsyncStream[AxonEventView]:
+            # after_sequence is used internally for reconnection only
+            sequence_int = int(last_sequence) if last_sequence is not None else None
+            return await self._get(
+                path_template("/v1/axons/{id}/subscribe/sse", id=id),
+                options=make_request_options(
+                    extra_headers=merged_headers,
+                    extra_query=extra_query,
+                    extra_body=extra_body,
+                    timeout=timeout,
+                    query=maybe_transform(
+                        {"after_sequence": sequence_int},
+                        axon_subscribe_sse_params.AxonSubscribeSseParams,
+                    ),
+                ),
+                cast_to=AxonEventView,
+                stream=True,
+                stream_cls=AsyncStream[AxonEventView],
+            )
+
+        initial_stream = await create_stream(None)
+
+        def get_sequence(item: AxonEventView) -> str | None:
+            value = getattr(item, "sequence", None)
+            if value is None:
+                return None
+            return str(value)
+
+        return cast(
+            AsyncStream[AxonEventView],
+            AsyncReconnectingStream(
+                current_stream=initial_stream,
+                stream_creator=create_stream,
+                get_offset=get_sequence,
             ),
-            cast_to=AxonEventView,
-            stream=True,
-            stream_cls=AsyncStream[AxonEventView],
         )
 
 

src/runloop_api_client/types/axons/__init__.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from .axon_subscribe_sse_params import AxonSubscribeSseParams as AxonSubscribeSseParams
 from .sql_batch_params import SqlBatchParams as SqlBatchParams
 from .sql_query_params import SqlQueryParams as SqlQueryParams
 from .sql_step_error_view import SqlStepErrorView as SqlStepErrorView

src/runloop_api_client/types/axons/axon_subscribe_sse_params.py

@@ -0,0 +1,14 @@
+# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
+
+from __future__ import annotations
+
+from typing_extensions import TypedDict
+
+from ..._types import NotGiven
+
+__all__ = ["AxonSubscribeSseParams"]
+
+
+class AxonSubscribeSseParams(TypedDict, total=False):
+    after_sequence: int | NotGiven
+    """Resume SSE stream from events after this sequence number (used internally for reconnection)"""