Reconnect stream on timeout

alb-rl · alb-rl · commit ed970e17e6bb · 2025-09-08T17:02:03.000-07:00
diff --git a/src/runloop_api_client/_streaming.py b/src/runloop_api_client/_streaming.py
@@ -4,12 +4,31 @@
 import json
 import inspect
 from types import TracebackType
-from typing import TYPE_CHECKING, Any, Generic, TypeVar, Iterator, AsyncIterator, cast
-from typing_extensions import Self, Protocol, TypeGuard, override, get_origin, runtime_checkable
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Generic,
+    TypeVar,
+    Callable,
+    Iterator,
+    Optional,
+    Awaitable,
+    AsyncIterator,
+    cast,
+)
+from typing_extensions import (
+    Self,
+    Protocol,
+    TypeGuard,
+    override,
+    get_origin,
+    runtime_checkable,
+)
 
 import httpx
 
 from ._utils import extract_type_var_from_base
+from ._exceptions import APIStatusError, APITimeoutError
 
 if TYPE_CHECKING:
     from ._client import Runloop, AsyncRunloop
@@ -55,6 +74,17 @@ def __stream__(self) -> Iterator[_T]:
         iterator = self._iter_events()
 
         for sse in iterator:
+            # Surface server-sent error events as API errors to allow callers to handle/retry
+            if sse.event == "error":
+                try:
+                    error_obj = json.loads(sse.data)
+                    status_code = int(error_obj.get("code", 500))
+                    # Build a synthetic response to mirror normal error handling
+                    fake_resp = httpx.Response(status_code, request=response.request, content=sse.data)
+                except Exception:
+                    fake_resp = httpx.Response(500, request=response.request, content=sse.data)
+                raise self._client._make_status_error_from_response(fake_resp)
+
             yield process_data(data=sse.json(), cast_to=cast_to, response=response)
 
         # Ensure the entire stream is consumed
@@ -119,6 +149,17 @@ async def __stream__(self) -> AsyncIterator[_T]:
         iterator = self._iter_events()
 
         async for sse in iterator:
+            # Surface server-sent error events as API errors to allow callers to handle/retry
+            if sse.event == "error":
+                try:
+                    error_obj = json.loads(sse.data)
+                    status_code = int(error_obj.get("code", 500))
+                    # Build a synthetic response to mirror normal error handling
+                    fake_resp = httpx.Response(status_code, request=response.request, content=sse.data)
+                except Exception:
+                    fake_resp = httpx.Response(500, request=response.request, content=sse.data)
+                raise self._client._make_status_error_from_response(fake_resp)
+
             yield process_data(data=sse.json(), cast_to=cast_to, response=response)
 
         # Ensure the entire stream is consumed
@@ -331,3 +372,149 @@ class MyStream(Stream[bytes]):
         generic_bases=cast("tuple[type, ...]", (Stream, AsyncStream)),
         failure_message=failure_message,
     )
+
+
+class ReconnectingStream(Generic[_T]):
+    """Wraps a Stream with automatic reconnection on timeout (HTTP 408) or read timeouts.
+
+    The reconnection uses the last observed offset from each item, as provided by
+    the given `get_offset` callback. The `stream_creator` will be called with the
+    last known offset to resume the stream.
+    """
+
+    def __init__(
+        self,
+        *,
+        current_stream: Stream[_T],
+        stream_creator: Callable[[Optional[str]], Stream[_T]],
+        get_offset: Callable[[_T], Optional[str]],
+    ) -> None:
+        self._current_stream = current_stream
+        self._stream_creator = stream_creator
+        self._get_offset = get_offset
+        self._last_offset: Optional[str] = None
+        self._iterator = self.__stream__()
+
+    @property
+    def response(self) -> httpx.Response:
+        return self._current_stream.response
+
+    def __next__(self) -> _T:
+        return self._iterator.__next__()
+
+    def __iter__(self) -> Iterator[_T]:
+        for item in self._iterator:
+            yield item
+
+    def __enter__(self) -> "ReconnectingStream[_T]":
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        self.close()
+
+    def close(self) -> None:
+        self._current_stream.close()
+
+    def __stream__(self) -> Iterator[_T]:
+        while True:
+            try:
+                for item in self._current_stream:
+                    offset = self._get_offset(item)
+                    if offset is not None:
+                        self._last_offset = offset
+                    yield item
+                return
+            except Exception as e:
+                # Reconnect on timeouts
+                should_reconnect = False
+                if isinstance(e, APITimeoutError):
+                    should_reconnect = True
+                elif isinstance(e, APIStatusError) and getattr(e, "status_code", None) == 408:
+                    should_reconnect = True
+                elif isinstance(e, httpx.TimeoutException):
+                    should_reconnect = True
+
+                if should_reconnect:
+                    # Close existing response before reconnecting
+                    try:
+                        self._current_stream.close()
+                    except Exception:
+                        pass
+                    self._current_stream = self._stream_creator(self._last_offset)
+                    continue
+                raise
+
+
+class AsyncReconnectingStream(Generic[_T]):
+    """Async variant of ReconnectingStream supporting auto-reconnect on timeouts."""
+
+    def __init__(
+        self,
+        *,
+        current_stream: AsyncStream[_T],
+        stream_creator: Callable[[Optional[str]], Awaitable[AsyncStream[_T]]],
+        get_offset: Callable[[_T], Optional[str]],
+    ) -> None:
+        self._current_stream = current_stream
+        self._stream_creator = stream_creator
+        self._get_offset = get_offset
+        self._last_offset: Optional[str] = None
+        self._iterator = self.__stream__()
+
+    @property
+    def response(self) -> httpx.Response:
+        return self._current_stream.response
+
+    async def __anext__(self) -> _T:
+        return await self._iterator.__anext__()
+
+    async def __aiter__(self) -> AsyncIterator[_T]:
+        async for item in self._iterator:
+            yield item
+
+    async def __aenter__(self) -> "AsyncReconnectingStream[_T]":
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        await self.close()
+
+    async def close(self) -> None:
+        await self._current_stream.close()
+
+    async def __stream__(self) -> AsyncIterator[_T]:
+        while True:
+            try:
+                async for item in self._current_stream:
+                    offset = self._get_offset(item)
+                    if offset is not None:
+                        self._last_offset = offset
+                    yield item
+                return
+            except Exception as e:
+                # Reconnect on timeouts
+                should_reconnect = False
+                if isinstance(e, APITimeoutError):
+                    should_reconnect = True
+                elif isinstance(e, APIStatusError) and getattr(e, "status_code", None) == 408:
+                    should_reconnect = True
+                elif isinstance(e, httpx.TimeoutException):
+                    should_reconnect = True
+
+                if should_reconnect:
+                    try:
+                        await self._current_stream.close()
+                    except Exception:
+                        pass
+                    self._current_stream = await self._stream_creator(self._last_offset)
+                    continue
+                raise
diff --git a/src/runloop_api_client/resources/devboxes/executions.py b/src/runloop_api_client/resources/devboxes/executions.py
@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import Optional
+from typing import Optional, cast
 
 import httpx
 
@@ -16,10 +16,10 @@
     async_to_raw_response_wrapper,
     async_to_streamed_response_wrapper,
 )
-from ..._constants import DEFAULT_TIMEOUT
+from ..._constants import DEFAULT_TIMEOUT, RAW_RESPONSE_HEADER
+from ..._streaming import Stream, AsyncStream, ReconnectingStream, AsyncReconnectingStream
 from ..._exceptions import APIStatusError, APITimeoutError
 from ...lib.polling import PollingConfig, poll_until
-from ..._streaming import Stream, AsyncStream
 from ..._base_client import make_request_options
 from ...types.devboxes import (
     execution_kill_params,
@@ -357,18 +357,54 @@ def stream_updates(
             raise ValueError(f"Expected a non-empty value for `devbox_id` but received {devbox_id!r}")
         if not execution_id:
             raise ValueError(f"Expected a non-empty value for `execution_id` but received {execution_id!r}")
-        return self._get(
-            f"/v1/devboxes/{devbox_id}/executions/{execution_id}/stream_updates",
-            options=make_request_options(
-                extra_headers=extra_headers,
-                extra_query=extra_query,
-                extra_body=extra_body,
-                timeout=timeout,
-                query=maybe_transform({"offset": offset}, execution_stream_updates_params.ExecutionStreamUpdatesParams),
-            ),
-            cast_to=DevboxAsyncExecutionDetailView,
-            stream=True,
-            stream_cls=Stream[ExecutionUpdateChunk],
+        # If caller requested a raw or streaming response wrapper, return the underlying stream as-is
+        if extra_headers and extra_headers.get(RAW_RESPONSE_HEADER):
+            return self._get(
+                f"/v1/devboxes/{devbox_id}/executions/{execution_id}/stream_updates",
+                options=make_request_options(
+                    extra_headers=extra_headers,
+                    extra_query=extra_query,
+                    extra_body=extra_body,
+                    timeout=timeout,
+                    query=maybe_transform(
+                        {"offset": offset}, execution_stream_updates_params.ExecutionStreamUpdatesParams
+                    ),
+                ),
+                cast_to=DevboxAsyncExecutionDetailView,
+                stream=True,
+                stream_cls=Stream[ExecutionUpdateChunk],
+            )
+
+        # Otherwise, wrap with auto-reconnect using last seen offset
+        def create_stream(last_offset: str | None) -> Stream[ExecutionUpdateChunk]:
+            new_offset = last_offset if last_offset is not None else (None if isinstance(offset, NotGiven) else offset)
+            return self._get(
+                f"/v1/devboxes/{devbox_id}/executions/{execution_id}/stream_updates",
+                options=make_request_options(
+                    extra_headers=extra_headers,
+                    extra_query=extra_query,
+                    extra_body=extra_body,
+                    timeout=timeout,
+                    query=maybe_transform(
+                        {"offset": new_offset}, execution_stream_updates_params.ExecutionStreamUpdatesParams
+                    ),
+                ),
+                cast_to=DevboxAsyncExecutionDetailView,
+                stream=True,
+                stream_cls=Stream[ExecutionUpdateChunk],
+            )
+
+        initial_stream = create_stream(None)
+
+        def get_offset(item: ExecutionUpdateChunk) -> str | None:
+            value = getattr(item, "offset", None)
+            if value is None:
+                return None
+            return str(value)
+
+        return cast(
+            Stream[ExecutionUpdateChunk],
+            ReconnectingStream(current_stream=initial_stream, stream_creator=create_stream, get_offset=get_offset),
         )
 
 
@@ -683,20 +719,53 @@ async def stream_updates(
             raise ValueError(f"Expected a non-empty value for `devbox_id` but received {devbox_id!r}")
         if not execution_id:
             raise ValueError(f"Expected a non-empty value for `execution_id` but received {execution_id!r}")
-        return await self._get(
-            f"/v1/devboxes/{devbox_id}/executions/{execution_id}/stream_updates",
-            options=make_request_options(
-                extra_headers=extra_headers,
-                extra_query=extra_query,
-                extra_body=extra_body,
-                timeout=timeout,
-                query=await async_maybe_transform(
-                    {"offset": offset}, execution_stream_updates_params.ExecutionStreamUpdatesParams
+        # If caller requested a raw or streaming response wrapper, return the underlying stream as-is
+        if extra_headers and extra_headers.get(RAW_RESPONSE_HEADER):
+            return await self._get(
+                f"/v1/devboxes/{devbox_id}/executions/{execution_id}/stream_updates",
+                options=make_request_options(
+                    extra_headers=extra_headers,
+                    extra_query=extra_query,
+                    extra_body=extra_body,
+                    timeout=timeout,
+                    query=await async_maybe_transform(
+                        {"offset": offset}, execution_stream_updates_params.ExecutionStreamUpdatesParams
+                    ),
                 ),
-            ),
-            cast_to=DevboxAsyncExecutionDetailView,
-            stream=True,
-            stream_cls=AsyncStream[ExecutionUpdateChunk],
+                cast_to=DevboxAsyncExecutionDetailView,
+                stream=True,
+                stream_cls=AsyncStream[ExecutionUpdateChunk],
+            )
+
+        async def create_stream(last_offset: str | None) -> AsyncStream[ExecutionUpdateChunk]:
+            new_offset = last_offset if last_offset is not None else (None if isinstance(offset, NotGiven) else offset)
+            return await self._get(
+                f"/v1/devboxes/{devbox_id}/executions/{execution_id}/stream_updates",
+                options=make_request_options(
+                    extra_headers=extra_headers,
+                    extra_query=extra_query,
+                    extra_body=extra_body,
+                    timeout=timeout,
+                    query=await async_maybe_transform(
+                        {"offset": new_offset}, execution_stream_updates_params.ExecutionStreamUpdatesParams
+                    ),
+                ),
+                cast_to=DevboxAsyncExecutionDetailView,
+                stream=True,
+                stream_cls=AsyncStream[ExecutionUpdateChunk],
+            )
+
+        initial_stream = await create_stream(None)
+
+        def get_offset(item: ExecutionUpdateChunk) -> str | None:
+            value = getattr(item, "offset", None)
+            if value is None:
+                return None
+            return str(value)
+
+        return cast(
+            AsyncStream[ExecutionUpdateChunk],
+            AsyncReconnectingStream(current_stream=initial_stream, stream_creator=create_stream, get_offset=get_offset),
         )
 
 
diff --git a/tests/api_resources/devboxes/test_executions.py b/tests/api_resources/devboxes/test_executions.py
