@typespec/http-client-python 0.37.2 → 0.38.0

package/generator/pygen/codegen/serializers/builder_serializer.py

@@ -1150,22 +1150,36 @@ class _OperationSerializer(_BuilderBaseSerializer[OperationType]):
         return retval
 
     def handle_error_response(  # pylint: disable=too-many-statements, too-many-branches
-        self, builder: OperationType
+        self,
+        builder: OperationType,
+        *,
+        response_name: str = "response",
+        pipeline_response_name: str = "pipeline_response",
+        success_status_codes: Optional[list[int]] = None,
+        stream_value: Optional[Union[str, bool]] = None,
+        response_read_name: Optional[str] = None,
     ) -> list[str]:
         async_await = "await " if self.async_mode else ""
-        retval = [f"if response.status_code not in {str(builder.success_status_codes)}:"]
-        response_read = [
-            "    try:",
-            f"        {async_await}response.read()  # Load the body in memory and close the socket",
-            "    except (StreamConsumedError, StreamClosedError):",
-            "        pass",
-        ]
-        if builder.stream_value is True:  # _stream is True so no need to judge it
+        success_status_codes = builder.success_status_codes if success_status_codes is None else success_status_codes
+        stream_value = builder.stream_value if stream_value is None else stream_value
+        retval = [f"if {response_name}.status_code not in {str(success_status_codes)}:"]
+        if response_read_name:
+            response_read = [f"    {async_await}{response_read_name}({response_name})"]
+        else:
+            response_read = [
+                "    try:",
+                f"        {async_await}{response_name}.read()  # Load the body in memory and close the socket",
+                "    except (StreamConsumedError, StreamClosedError):",
+                "        pass",
+            ]
+        if stream_value is True:  # _stream is True so no need to judge it
             retval.extend(response_read)
-        elif isinstance(builder.stream_value, str):  # _stream is not sure, so we need to judge it
+        elif isinstance(stream_value, str):  # _stream is not sure, so we need to judge it
             retval.append("    if _stream:")
             retval.extend([f"    {l}" for l in response_read])
-        retval.append("    map_error(status_code=response.status_code, response=response, error_map=error_map)")
+        retval.append(
+            f"    map_error(status_code={response_name}.status_code, " f"response={response_name}, error_map=error_map)"
+        )
         error_model = ""
         if builder.non_default_errors and (  # pylint: disable=too-many-nested-blocks
             self.code_model.options["models-mode"] or self.code_model.generate_typeddict_only
@@ -1181,7 +1195,7 @@ class _OperationSerializer(_BuilderBaseSerializer[OperationType]):
                     pylint_disable = ""
                 if isinstance(e.status_codes[0], int):
                     for status_code in e.status_codes:
-                        retval.append(f"    {condition} response.status_code == {status_code}:")
+                        retval.append(f"    {condition} {response_name}.status_code == {status_code}:")
                         type_annotation = e.type.type_annotation(  # type: ignore
                             is_operation_file=True, skip_quote=True, serialize_namespace=self.serialize_namespace
                         )
@@ -1190,13 +1204,15 @@ class _OperationSerializer(_BuilderBaseSerializer[OperationType]):
                                 fn = "_failsafe_deserialize_xml"
                             else:
                                 fn = "_failsafe_deserialize"
-                            retval.append(f"        error = {fn}({type_annotation},{pylint_disable}\n  response)")
+                            retval.append(
+                                f"        error = {fn}({type_annotation},{pylint_disable}\n  {response_name})"
+                            )
                         else:
                             retval.extend(
                                 [
                                     "        error = self._deserialize.failsafe_deserialize(",
                                     f"            {type_annotation},{pylint_disable}",
-                                    "            pipeline_response,",
+                                    f"            {pipeline_response_name},",
                                     "        )",
                                 ]
                             )
@@ -1209,7 +1225,8 @@ class _OperationSerializer(_BuilderBaseSerializer[OperationType]):
                 # ranged status code only exist in typespec and will not have multiple status codes
                 else:
                     retval.append(
-                        f"    {condition} {e.status_codes[0][0]} <= response.status_code <= {e.status_codes[0][1]}:"
+                        f"    {condition} {e.status_codes[0][0]} <= "
+                        f"{response_name}.status_code <= {e.status_codes[0][1]}:"
                     )
                     type_annotation = e.type.type_annotation(  # type: ignore
                         is_operation_file=True, skip_quote=True, serialize_namespace=self.serialize_namespace
@@ -1218,19 +1235,19 @@ class _OperationSerializer(_BuilderBaseSerializer[OperationType]):
                         if xml_serializable(str(e.default_content_type)):
                             retval.append(
                                 "        error = _failsafe_deserialize_xml("
-                                f"{type_annotation},{pylint_disable}\n  response)"
+                                f"{type_annotation},{pylint_disable}\n  {response_name})"
                             )
                         else:
                             retval.append(
-                                "        error = _failsafe_deserialize("
-                                f"{type_annotation},{pylint_disable}\n  response)"
+                                f"        error = _failsafe_deserialize("
+                                f"{type_annotation},{pylint_disable}\n  {response_name})"
                             )
                     else:
                         retval.extend(
                             [
                                 "        error = self._deserialize.failsafe_deserialize(",
                                 f"            {type_annotation},{pylint_disable}",
-                                "            pipeline_response,",
+                                f"            {pipeline_response_name},",
                                 "        )",
                             ]
                         )
@@ -1254,7 +1271,7 @@ class _OperationSerializer(_BuilderBaseSerializer[OperationType]):
                     [
                         f"{indent}error = {fn}(",
                         f"{indent}    {default_error_deserialization}",
-                        f"{indent}    response,",
+                        f"{indent}    {response_name},",
                         f"{indent})",
                     ]
                 )
@@ -1263,12 +1280,13 @@ class _OperationSerializer(_BuilderBaseSerializer[OperationType]):
                     [
                         f"{indent}error = self._deserialize.failsafe_deserialize(",
                         f"{indent}    {default_error_deserialization}",
-                        f"{indent}    pipeline_response,",
+                        f"{indent}    {pipeline_response_name},",
                         f"{indent})",
                     ]
                 )
         retval.append(
-            "    raise HttpResponseError(response=response{}{})".format(
+            "    raise HttpResponseError(response={}{}{})".format(
+                response_name,
                 error_model,
                 (", error_format=ARMErrorFormat" if self.code_model.options["azure-arm"] else ""),
             )
@@ -1378,6 +1396,41 @@ class _OperationSerializer(_BuilderBaseSerializer[OperationType]):
             stream_kwargs.append(f"terminal_event={terminal_event!r}")
         if terminal_event_names:
             stream_kwargs.append(f"terminal_event_names={terminal_event_names!r}")
+        if response.streaming_kind == "sse":  # type: ignore[attr-defined]
+            retval.append("")
+            retval.append(
+                "async def _reconnect(_last_event_id, _reconnect_delay):"
+                if self.async_mode
+                else "def _reconnect(_last_event_id, _reconnect_delay):"
+            )
+            retval.append("    _transport: Any = pipeline_response.context.transport")
+            retval.append(f"    {'await ' if self.async_mode else ''}_transport.sleep(_reconnect_delay)")
+            retval.append("    _update_sse_request_headers(_request, _last_event_id)")
+            retval.append(
+                f"    _reconnect_response = {self._call_method}self._client.send_request("
+                "_request, stream=True, **kwargs)"
+            )
+            reconnect_status_codes = list(builder.success_status_codes)
+            if 204 not in reconnect_status_codes:
+                reconnect_status_codes.append(204)
+            retval.extend(
+                [
+                    f"    {line}"
+                    for line in self.handle_error_response(
+                        builder,
+                        response_name="_reconnect_response",
+                        pipeline_response_name=(
+                            "PipelineResponse(_request, _reconnect_response, pipeline_response.context)"
+                        ),
+                        success_status_codes=reconnect_status_codes,
+                        stream_value=True,
+                        response_read_name=("_read_sse_response_async" if self.async_mode else "_read_sse_response"),
+                    )
+                ]
+            )
+            retval.append("    return _reconnect_response")
+            stream_kwargs.append("last_event_id=_last_event_id")
+            stream_kwargs.append("reconnect_callback=_reconnect")
         unnamed_terminal_values = (
             [discriminator_value for discriminator_value, event in unnamed_discriminator[1] if event.is_terminal]
             if unnamed_discriminator is not None
@@ -1843,7 +1896,7 @@ class _LROOperationSerializer(_OperationSerializer[LROOperationType]):
             if builder.lro_response.headers:
                 retval.append("    response_headers = {}")
             if (
-                (not self.code_model.options["models-mode"] and not self.code_model.generate_typeddict_only)
+                not self.code_model.options["models-mode"]
                 or self.code_model.options["models-mode"] == "dpg"
                 or builder.lro_response.headers
             ):

package/generator/pygen/codegen/serializers/unions_serializer.py

@@ -3,7 +3,7 @@
 # Licensed under the MIT License. See License.txt in the project root for
 # license information.
 # --------------------------------------------------------------------------
-from ..models import CodeModel
+from ..models import CodeModel, CombinedType
 from ..models.imports import FileImport, ImportType
 from ..models.utils import NamespaceType
 from .import_serializer import FileImportSerializer
@@ -18,15 +18,30 @@ class UnionsSerializer(BaseSerializer):
     ):
         super().__init__(code_model=code_model, env=env)
 
+    @property
+    def named_unions(self) -> list[CombinedType]:
+        # The same named union can reach codegen as multiple objects, so collapse
+        # by emitted alias name (keeping insertion order) to avoid duplicate aliases.
+        deduped: dict[str, CombinedType] = {}
+        for union in self.code_model.named_unions:
+            if union.name:
+                deduped.setdefault(union.name, union)
+        return list(deduped.values())
+
     def imports(self) -> FileImport:
         file_import = FileImport(self.code_model)
-        if self.code_model.named_unions:
+        if self.named_unions:
+            file_import.add_submodule_import(
+                "typing",
+                "TypeAlias",
+                ImportType.STDLIB,
+            )
             file_import.add_submodule_import(
                 "typing",
                 "Union",
                 ImportType.STDLIB,
             )
-        for nu in self.code_model.named_unions:
+        for nu in self.named_unions:
             file_import.merge(
                 nu.imports(
                     serialize_namespace=self.serialize_namespace,
@@ -39,6 +54,7 @@ class UnionsSerializer(BaseSerializer):
         template = self.env.get_template("unions.py.jinja2")
         return template.render(
             code_model=self.code_model,
+            named_unions=self.named_unions,
             imports=FileImportSerializer(self.imports()),
             serializer=self,
         )

package/generator/pygen/codegen/templates/streaming_base.py.jinja2

@@ -33,7 +33,8 @@ from typing import (
 
 from typing_extensions import Self
 
-from {{ code_model.core_library }}.rest import AsyncHttpResponse, HttpResponse
+from {{ code_model.core_library }}.exceptions import StreamClosedError, StreamConsumedError
+from {{ code_model.core_library }}.rest import AsyncHttpResponse, HttpRequest, HttpResponse
 
 DecodedType = TypeVar("DecodedType")
 ReturnType_co = TypeVar("ReturnType_co", covariant=True)
@@ -435,13 +436,56 @@ async def _aiter_sse_lines(iter_bytes: AsyncIterator[bytes]) -> AsyncGenerator[s
         yield line
 
 
+def _read_sse_response(response: HttpResponse) -> None:
+    """Read an SSE response body, ignoring responses that are already consumed or closed.
+
+    :param response: The response to read.
+    :type response: HttpResponse
+    """
+    try:
+        response.read()
+    except (StreamConsumedError, StreamClosedError):
+        pass
+
+
+async def _read_sse_response_async(response: AsyncHttpResponse) -> None:
+    """Read an SSE response body, ignoring responses that are already consumed or closed.
+
+    :param response: The response to read.
+    :type response: AsyncHttpResponse
+    """
+    try:
+        await response.read()
+    except (StreamConsumedError, StreamClosedError):
+        pass
+
+
+def _update_sse_request_headers(request: HttpRequest, last_event_id: Optional[str]) -> None:
+    """Update the last event ID header for an SSE reconnection request.
+
+    :param request: The reconnection request to update.
+    :type request: HttpRequest
+    :param last_event_id: The current value of the SSE last event ID buffer.
+    :type last_event_id: str or None
+    """
+    if last_event_id:
+        request.headers["Last-Event-ID"] = last_event_id
+    else:
+        request.headers.pop("Last-Event-ID", None)
+
+
 class _SSEEventBuilder:
-    """Accumulates SSE field lines and builds :class:`ServerSentEvent` instances."""
+    """Accumulates SSE field lines and builds :class:`ServerSentEvent` instances.
 
-    def __init__(self) -> None:
+    :param last_event_id: The initial value of the SSE last event ID buffer.
+    :type last_event_id: str
+    """
+
+    def __init__(self, last_event_id: str = "") -> None:
         self._data: List[str] = []
         self._event_type = ""
-        self._last_id = ""
+        self._last_id = last_event_id
+        self._pending_last_id = last_event_id
         self._retry: Optional[int] = None
 
     def add_line(self, line: str) -> Optional[ServerSentEvent]:
@@ -469,7 +513,7 @@ class _SSEEventBuilder:
             self._data.append(value)
         elif field == "id":
             if "\x00" not in value:
-                self._last_id = value
+                self._pending_last_id = value
         elif field == "retry":
             if value.isascii() and value.isdigit():
                 try:
@@ -482,6 +526,7 @@ class _SSEEventBuilder:
         return None
 
     def _dispatch(self) -> Optional[ServerSentEvent]:
+        self._last_id = self._pending_last_id
         if not self._data:
             # No data accumulated: reset and dispatch nothing.
             self._event_type = ""
@@ -496,13 +541,55 @@ class _SSEEventBuilder:
         self._event_type = ""
         return event
 
+    @property
+    def last_event_id(self) -> str:
+        """Return the current value of the last event ID buffer.
+
+        :return: The current value of the last event ID buffer.
+        :rtype: str
+        """
+        return self._last_id
+
+    @property
+    def retry(self) -> Optional[int]:
+        """Return the most recently received valid reconnection time.
+
+        :return: The reconnection time in milliseconds, if one was provided.
+        :rtype: int or None
+        """
+        return self._retry
+
 
 class SSEDecoder:
     """Decoder for Server-Sent Events (SSE).
 
     https://html.spec.whatwg.org/multipage/server-sent-events.html
+
+    :keyword last_event_id: The initial value of the SSE last event ID buffer.
+    :paramtype last_event_id: str or None
     """
 
+    def __init__(self, *, last_event_id: Optional[str] = None) -> None:
+        self._builder = _SSEEventBuilder(last_event_id or "")
+
+    @property
+    def last_event_id(self) -> str:
+        """Return the current value of the last event ID buffer.
+
+        :return: The current value of the last event ID buffer.
+        :rtype: str
+        """
+        return self._builder.last_event_id
+
+    @property
+    def retry(self) -> Optional[int]:
+        """Return the most recently received valid reconnection time.
+
+        :return: The reconnection time in milliseconds, if one was provided.
+        :rtype: int or None
+        """
+        return self._builder.retry
+
     def iter_events(self, iter_bytes: Iterator[bytes]) -> Iterator[ServerSentEvent]:
         """Iterate over SSE events from a byte iterator.
 
@@ -511,9 +598,8 @@ class SSEDecoder:
         :rtype: Iterator[ServerSentEvent]
         :return: An iterator of server-sent events.
         """
-        builder = _SSEEventBuilder()
         for line in _iter_sse_lines(iter_bytes):
-            event = builder.add_line(line)
+            event = self._builder.add_line(line)
             if event is not None:
                 yield event
 
@@ -522,8 +608,32 @@ class AsyncSSEDecoder:
     """Asynchronous decoder for Server-Sent Events (SSE).
 
     https://html.spec.whatwg.org/multipage/server-sent-events.html
+
+    :keyword last_event_id: The initial value of the SSE last event ID buffer.
+    :paramtype last_event_id: str or None
     """
 
+    def __init__(self, *, last_event_id: Optional[str] = None) -> None:
+        self._builder = _SSEEventBuilder(last_event_id or "")
+
+    @property
+    def last_event_id(self) -> str:
+        """Return the current value of the last event ID buffer.
+
+        :return: The current value of the last event ID buffer.
+        :rtype: str
+        """
+        return self._builder.last_event_id
+
+    @property
+    def retry(self) -> Optional[int]:
+        """Return the most recently received valid reconnection time.
+
+        :return: The reconnection time in milliseconds, if one was provided.
+        :rtype: int or None
+        """
+        return self._builder.retry
+
     # pylint: disable=invalid-overridden-method
     async def aiter_events(self, iter_bytes: AsyncIterator[bytes]) -> AsyncIterator[ServerSentEvent]:
         """Asynchronously iterate over SSE events from a byte iterator.
@@ -533,10 +643,9 @@ class AsyncSSEDecoder:
         :rtype: AsyncIterator[ServerSentEvent]
         :return: An asynchronous iterator of server-sent events.
         """
-        builder = _SSEEventBuilder()
         async with aclosing(_aiter_sse_lines(iter_bytes)) as lines:
             async for line in lines:
-                event = builder.add_line(line)
+                event = self._builder.add_line(line)
                 if event is not None:
                     yield event
 
@@ -563,6 +672,14 @@ class Stream(Iterator[ReturnType_co]):
     :keyword terminal_event_predicate: Optional predicate that identifies a payload-bearing
         terminal event. The event is yielded before iteration stops.
     :paramtype terminal_event_predicate: Callable[[Any], bool] or None
+    :keyword last_event_id: Optional event ID from which this stream is resuming.
+    :paramtype last_event_id: str or None
+    :keyword reconnect_callback: Optional callback used to re-establish the underlying response
+        when the stream disconnects before a terminal event is seen. Reconnection uses a
+        three-second delay unless the server provides an SSE ``retry`` value. The callback is
+        called with the last-seen SSE event ID (or ``None``) and the reconnect delay in seconds,
+        and must return a new response.
+    :paramtype reconnect_callback: Callable[[str or None, float], ~{{ code_model.core_library }}.rest.HttpResponse] or None
     """
 
     def __init__(
@@ -574,17 +691,26 @@ class Stream(Iterator[ReturnType_co]):
         terminal_event: Optional[str] = None,
         terminal_event_names: Optional[Sequence[str]] = None,
         terminal_event_predicate: Optional[Callable[[DecodedType], bool]] = None,
+        last_event_id: Optional[str] = None,
+        reconnect_callback: Optional[Callable[[Optional[str], float], HttpResponse]] = None,
     ) -> None:
         self._response = response
         content_type = response.headers.get("Content-Type", "").split(";", 1)[0].strip().lower()
         self._decoder: StreamDecoder[Any] = (
-            decoder if decoder is not None else (SSEDecoder() if content_type == "text/event-stream" else JSONLDecoder())
+            decoder
+            if decoder is not None
+            else (
+                SSEDecoder(last_event_id=last_event_id)
+                if content_type == "text/event-stream"
+                else JSONLDecoder()
+            )
         )
         self._deserialization_callback = deserialization_callback
         self._terminal_event = terminal_event
         self._terminal_event_names = frozenset(terminal_event_names or ())
         self._terminal_event_predicate = terminal_event_predicate
-        self._last_event_id: Optional[str] = None
+        self._reconnect_callback = reconnect_callback
+        self._last_event_id = last_event_id
         self._retry: Optional[int] = None
         self._iterator = self._iter_results()
 
@@ -605,26 +731,46 @@ class Stream(Iterator[ReturnType_co]):
         return self
 
     def _iter_results(self) -> Generator[ReturnType_co, None, None]:
+        terminal = False
         try:
-            for event in self._decoder.iter_events(self._response.iter_bytes()):
-                event_id = getattr(event, "id", None)
-                if event_id is not None:
-                    self._last_event_id = event_id
-                event_retry = getattr(event, "retry", None)
-                if event_retry is not None:
-                    self._retry = event_retry
-                if self._terminal_event is not None and getattr(event, "data", None) == self._terminal_event:
+            while True:
+                if getattr(self._response, "status_code", None) == 204:
                     break
-                result = self._deserialization_callback(self._response, event)
-                yield result
-                if (
-                    self._terminal_event_names
-                    and getattr(event, "event", None) in self._terminal_event_names
-                ) or (
-                    self._terminal_event_predicate is not None
-                    and self._terminal_event_predicate(event)
-                ):
+                for event in self._decoder.iter_events(self._response.iter_bytes()):
+                    event_id = getattr(event, "id", None)
+                    if event_id is not None:
+                        self._last_event_id = event_id
+                    event_retry = getattr(event, "retry", None)
+                    if event_retry is not None:
+                        self._retry = event_retry
+                    if self._terminal_event is not None and getattr(event, "data", None) == self._terminal_event:
+                        terminal = True
+                        break
+                    result = self._deserialization_callback(self._response, cast(Any, event))
+                    yield result
+                    if (
+                        self._terminal_event_names
+                        and getattr(event, "event", None) in self._terminal_event_names
+                    ) or (
+                        self._terminal_event_predicate is not None
+                        and self._terminal_event_predicate(cast(Any, event))
+                    ):
+                        terminal = True
+                        break
+                if terminal or self._reconnect_callback is None:
                     break
+                if isinstance(self._decoder, SSEDecoder):
+                    self._last_event_id = self._decoder.last_event_id
+                    if self._decoder.retry is not None:
+                        self._retry = self._decoder.retry
+                self._response.close()
+                reconnect_delay = (self._retry if self._retry is not None else 3000) / 1000
+                self._response = self._reconnect_callback(self._last_event_id, reconnect_delay)
+                self._decoder = (
+                    SSEDecoder(last_event_id=self._last_event_id)
+                    if isinstance(self._decoder, SSEDecoder)
+                    else type(self._decoder)()
+                )
         finally:
             self._response.close()
 
@@ -668,6 +814,14 @@ class AsyncStream(AsyncIterator[ReturnType_co]):
     :keyword terminal_event_predicate: Optional predicate that identifies a payload-bearing
         terminal event. The event is yielded before iteration stops.
     :paramtype terminal_event_predicate: Callable[[Any], bool] or None
+    :keyword last_event_id: Optional event ID from which this stream is resuming.
+    :paramtype last_event_id: str or None
+    :keyword reconnect_callback: Optional callback used to re-establish the underlying response
+        when the stream disconnects before a terminal event is seen. Reconnection uses a
+        three-second delay unless the server provides an SSE ``retry`` value. The callback is
+        called with the last-seen SSE event ID (or ``None``) and the reconnect delay in seconds,
+        and must resolve to a new response.
+    :paramtype reconnect_callback: Callable[[str or None, float], Any] or None
     """
 
     def __init__(
@@ -679,19 +833,26 @@ class AsyncStream(AsyncIterator[ReturnType_co]):
         terminal_event: Optional[str] = None,
         terminal_event_names: Optional[Sequence[str]] = None,
         terminal_event_predicate: Optional[Callable[[DecodedType], bool]] = None,
+        last_event_id: Optional[str] = None,
+        reconnect_callback: Optional[Callable[[Optional[str], float], Any]] = None,
     ) -> None:
         self._response = response
         content_type = response.headers.get("Content-Type", "").split(";", 1)[0].strip().lower()
         self._decoder: AsyncStreamDecoder[Any] = (
             decoder
             if decoder is not None
-            else (AsyncSSEDecoder() if content_type == "text/event-stream" else AsyncJSONLDecoder())
+            else (
+                AsyncSSEDecoder(last_event_id=last_event_id)
+                if content_type == "text/event-stream"
+                else AsyncJSONLDecoder()
+            )
         )
         self._deserialization_callback = deserialization_callback
         self._terminal_event = terminal_event
         self._terminal_event_names = frozenset(terminal_event_names or ())
         self._terminal_event_predicate = terminal_event_predicate
-        self._last_event_id: Optional[str] = None
+        self._reconnect_callback = reconnect_callback
+        self._last_event_id = last_event_id
         self._retry: Optional[int] = None
         self._iterator = self._iter_results()
 
@@ -712,34 +873,54 @@ class AsyncStream(AsyncIterator[ReturnType_co]):
         return self
 
     async def _iter_results(self) -> AsyncGenerator[ReturnType_co, None]:
-        events = self._decoder.aiter_events(self._response.iter_bytes())
+        terminal = False
         try:
-            async for event in events:
-                event_id = getattr(event, "id", None)
-                if event_id is not None:
-                    self._last_event_id = event_id
-                event_retry = getattr(event, "retry", None)
-                if event_retry is not None:
-                    self._retry = event_retry
-                if self._terminal_event is not None and getattr(event, "data", None) == self._terminal_event:
+            while True:
+                if getattr(self._response, "status_code", None) == 204:
                     break
-                result = self._deserialization_callback(self._response, event)
-                yield result
-                if (
-                    self._terminal_event_names
-                    and getattr(event, "event", None) in self._terminal_event_names
-                ) or (
-                    self._terminal_event_predicate is not None
-                    and self._terminal_event_predicate(event)
-                ):
+                events = self._decoder.aiter_events(self._response.iter_bytes())
+                try:
+                    async for event in events:
+                        event_id = getattr(event, "id", None)
+                        if event_id is not None:
+                            self._last_event_id = event_id
+                        event_retry = getattr(event, "retry", None)
+                        if event_retry is not None:
+                            self._retry = event_retry
+                        if self._terminal_event is not None and getattr(event, "data", None) == self._terminal_event:
+                            terminal = True
+                            break
+                        result = self._deserialization_callback(self._response, cast(Any, event))
+                        yield result
+                        if (
+                            self._terminal_event_names
+                            and getattr(event, "event", None) in self._terminal_event_names
+                        ) or (
+                            self._terminal_event_predicate is not None
+                            and self._terminal_event_predicate(cast(Any, event))
+                        ):
+                            terminal = True
+                            break
+                finally:
+                    aclose = getattr(events, "aclose", None)
+                    if aclose is not None:
+                        await aclose()
+                if terminal or self._reconnect_callback is None:
                     break
-        finally:
-            try:
-                aclose = getattr(events, "aclose", None)
-                if aclose is not None:
-                    await aclose()
-            finally:
+                if isinstance(self._decoder, AsyncSSEDecoder):
+                    self._last_event_id = self._decoder.last_event_id
+                    if self._decoder.retry is not None:
+                        self._retry = self._decoder.retry
                 await self._response.close()
+                reconnect_delay = (self._retry if self._retry is not None else 3000) / 1000
+                self._response = await self._reconnect_callback(self._last_event_id, reconnect_delay)
+                self._decoder = (
+                    AsyncSSEDecoder(last_event_id=self._last_event_id)
+                    if isinstance(self._decoder, AsyncSSEDecoder)
+                    else type(self._decoder)()
+                )
+        finally:
+            await self._response.close()
 
     async def __aexit__(
         self,

package/generator/pygen/codegen/templates/unions.py.jinja2

@@ -4,6 +4,6 @@
 {% endif %}
 
 {{ imports }}
-{% for nu in code_model.named_unions %}
-{{nu.name}} = {{nu.type_definition()}}
+{% for nu in named_unions %}
+{{nu.name}}: TypeAlias = {{nu.type_definition()}}
 {% endfor %}