Use Fern with_metadata in the high-level turn stream

src/truefoundry_gateway_sdk/agents/prepared_turn.py

@@ -5,9 +5,8 @@
 from ..types.turn import Turn as RawTurn
 from ..types.turn_created_event import TurnCreatedEvent
 from ..types.turn_done_event import TurnDoneEvent
-from ._sse_helpers import aiter_sse_stream, iter_sse_stream
 from .turn import AsyncTurn, Turn
-from .turn_stream_data import TurnStreamData
+from .turn_stream_data import TurnStreamData, parse_sequence_number
 
 # this is used as the default value for optional parameters
 OMIT = typing.cast(typing.Any, ...)
@@ -326,18 +325,19 @@ def _start_and_wait(self, poll_interval_ms: int, request_options: typing.Optiona
 
     def _consume_stream(self, request_options: typing.Optional[RequestOptions]) -> typing.Iterator[TurnStreamData]:
         """Consume the create_turn SSE, adopting the inner Turn from the first turn.created."""
-        with self._client.agents.sessions.with_raw_response.create_turn(
+        with self._client.agents.sessions.create_turn(
             self._session_id,
             input=self._input_param,
             previous_turn_id=self._previous_turn_id,
             request_options=request_options,
-        ) as r:
-            for item in iter_sse_stream(r._response):
-                if isinstance(item.event, TurnCreatedEvent) and self._turn is None:
-                    self._adopt_turn(item.event)
-                elif self._turn is not None and isinstance(item.event, TurnDoneEvent):
-                    self._replace_turn_state(item.event.state)
-                yield item
+        ) as sse:
+            for event in sse.with_metadata():
+                sequence_number = parse_sequence_number(event.id)
+                if isinstance(event.data, TurnCreatedEvent) and self._turn is None:
+                    self._adopt_turn(event.data)
+                elif self._turn is not None and isinstance(event.data, TurnDoneEvent):
+                    self._replace_turn_state(event.data.state)
+                yield TurnStreamData(sequence_number=sequence_number, event=event.data)
 
     def _must_get_turn(self) -> Turn:
         if self._turn is None:
@@ -688,18 +688,19 @@ async def _start_and_wait(
     async def _consume_stream(
         self, request_options: typing.Optional[RequestOptions]
     ) -> typing.AsyncIterator[TurnStreamData]:
-        async with self._client.agents.sessions.with_raw_response.create_turn(
+        async with self._client.agents.sessions.create_turn(
             self._session_id,
             input=self._input_param,
             previous_turn_id=self._previous_turn_id,
             request_options=request_options,
-        ) as r:
-            async for item in aiter_sse_stream(r._response):
-                if isinstance(item.event, TurnCreatedEvent) and self._turn is None:
-                    self._adopt_turn(item.event)
-                elif self._turn is not None and isinstance(item.event, TurnDoneEvent):
-                    self._replace_turn_state(item.event.state)
-                yield item
+        ) as sse:
+            async for event in sse.with_metadata():
+                sequence_number = parse_sequence_number(event.id)
+                if isinstance(event.data, TurnCreatedEvent) and self._turn is None:
+                    self._adopt_turn(event.data)
+                elif self._turn is not None and isinstance(event.data, TurnDoneEvent):
+                    self._replace_turn_state(event.data.state)
+                yield TurnStreamData(sequence_number=sequence_number, event=event.data)
 
     def _must_get_turn(self) -> AsyncTurn:
         if self._turn is None:

src/truefoundry_gateway_sdk/agents/turn.py

@@ -8,8 +8,7 @@
 from ..types.turn_state_cancelled import TurnStateCancelled
 from ..types.turn_state_done import TurnStateDone
 from ..types.turn_state_error import TurnStateError
-from ._sse_helpers import aiter_sse_stream, iter_sse_stream
-from .turn_stream_data import TurnStreamData
+from .turn_stream_data import TurnStreamData, parse_sequence_number
 
 # this is used as the default value for optional parameters
 OMIT = typing.cast(typing.Any, ...)
@@ -217,15 +216,16 @@ def stream(
         TurnStreamData
             SSE stream items.
         """
-        with self._client.agents.sessions.with_raw_response.subscribe_to_turn(
+        with self._client.agents.sessions.subscribe_to_turn(
             self._session_id,
             self._id,
             after_sequence_number=after_sequence_number,
             request_options=request_options,
-        ) as r:
-            for item in iter_sse_stream(r._response):
-                self._apply_event(item.event)
-                yield item
+        ) as sse:
+            for event in sse.with_metadata():
+                sequence_number = parse_sequence_number(event.id)
+                self._apply_event(event.data)
+                yield TurnStreamData(sequence_number=sequence_number, event=event.data)
 
     def cancel(self, *, request_options: typing.Optional[RequestOptions] = None) -> None:
         """
@@ -465,15 +465,16 @@ async def stream(
         TurnStreamData
             SSE stream items.
         """
-        async with self._client.agents.sessions.with_raw_response.subscribe_to_turn(
+        async with self._client.agents.sessions.subscribe_to_turn(
             self._session_id,
             self._id,
             after_sequence_number=after_sequence_number,
             request_options=request_options,
-        ) as r:
-            async for item in aiter_sse_stream(r._response):
-                self._apply_event(item.event)
-                yield item
+        ) as sse:
+            async for event in sse.with_metadata():
+                sequence_number = parse_sequence_number(event.id)
+                self._apply_event(event.data)
+                yield TurnStreamData(sequence_number=sequence_number, event=event.data)
 
     async def cancel(self, *, request_options: typing.Optional[RequestOptions] = None) -> None:
         """

src/truefoundry_gateway_sdk/agents/turn_stream_data.py

@@ -20,3 +20,17 @@ class TurnStreamData:
 
     sequence_number: int
     event: TurnStreamingEvent
+
+
+def parse_sequence_number(sse_id: typing.Optional[str]) -> int:
+    """Parse the SSE ``id`` field as a sequence number.
+
+    Raises ``ValueError`` when the id is absent or not a valid integer —
+    mirroring the TypeScript ``parseSequenceNumber``.
+    """
+    if not sse_id:
+        raise ValueError("Missing SSE sequence number id.")
+    try:
+        return int(sse_id)
+    except (ValueError, TypeError) as exc:
+        raise ValueError(f"Invalid SSE sequence number id: {sse_id!r}.") from exc