Coverage for src / lexigram / ai / relay / gateway / stream.py: 94%
145 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-08-08 23:08 +0800
1"""Upstream streaming framing, cancellation, and session lifecycle.
3The gateway parses framing only and holds NO accumulated text or tool
4arguments: each ``UpstreamChunk`` is decoded into a source DTO and
5forwarded to the stateful ``RelayStreamSessionProtocol``, which owns any
6partial content. The terminal flag is set on the final event of each
7finalize batch, and upstream cancellation plus session finalization each
8happen at most once.
9"""
11from __future__ import annotations
13import asyncio
14from collections.abc import AsyncIterator
15from dataclasses import dataclass
16from typing import Any, cast
18from lexigram.contracts.ai.exceptions import RelayError
19from lexigram.contracts.ai.relay import (
20 RelayFormat,
21 RelayGatewayError,
22 RelayStreamSessionProtocol,
23 RelayUpstreamProtocol,
24 RelayWireEvent,
25 UpstreamChunk,
26 UpstreamRequest,
27)
28from lexigram.contracts.ai.relay.dto import (
29 ClaudeStreamEvent,
30 GeminiResponse,
31 OpenAIChatStreamChunk,
32 ResponsesEvent,
33)
34from lexigram.contracts.ai.relay.gateway import RelayGatewayErrorCode
35from lexigram.logging import get_logger
36from lexigram.serialization import loads
38__all__ = ["UpstreamEventParser", "relay_stream"]
40logger = get_logger(__name__)
43@dataclass(frozen=True, slots=True)
44class _Parsed:
45 """Outcome of framing one upstream chunk."""
47 events: tuple[Any, ...]
48 terminal: bool = False
49 error: bool = False
52class UpstreamEventParser:
53 """Frame upstream chunks into target session events.
55 The parser decodes one chunk at a time, classifies it (keepalive,
56 delta, terminal, or error), and forwards source DTOs to the injected
57 session. It never accumulates text or tool arguments.
58 """
60 def __init__(
61 self,
62 session: RelayStreamSessionProtocol,
63 source: RelayFormat,
64 *,
65 request_id: str,
66 ) -> None:
67 """Bind the parser to a session and source wire format.
69 Args:
70 session: Stateful stream session accepting source DTOs and
71 emitting target events.
72 source: Wire format of the upstream stream.
73 request_id: Id stamped on the malformed-stream errors this
74 parser raises.
75 """
76 self._session = session
77 self._source = source
78 self._request_id = request_id
79 self._finalized = False
80 self._finalize_cache: tuple[Any, ...] = ()
81 self.finalized = False
82 self.truncated = False
83 self.cancelled = False
85 def parse(self, chunk: UpstreamChunk) -> _Parsed:
86 """Frame one upstream chunk into target session events.
88 Args:
89 chunk: One raw upstream frame.
91 Returns:
92 The framed outcome: emitted target events plus terminal and
93 error classification. A transport-level ``terminal=True``
94 chunk short-circuits decoding.
96 Raises:
97 RelayGatewayError: With code ``UPSTREAM_MALFORMED`` (502,
98 never retryable) when the payload is malformed JSON,
99 fails DTO validation, or is rejected by the session.
100 """
101 if chunk.terminal:
102 return _Parsed(())
103 if self._source == RelayFormat.OPENAI_CHAT:
104 return self._parse_openai_chat(chunk)
105 if self._source == RelayFormat.OPENAI_RESPONSES:
106 return self._parse_openai_responses(chunk)
107 if self._source == RelayFormat.CLAUDE:
108 return self._parse_claude(chunk)
109 return self._parse_gemini(chunk)
111 def finalize(self) -> tuple[Any, ...]:
112 """Close the session deterministically exactly once.
114 The first call runs the session finalize and caches its events;
115 subsequent calls return the cached result without touching the
116 session again.
118 Returns:
119 The session's terminal events.
120 """
121 if self._finalized:
122 return self._finalize_cache
123 self._finalized = True
124 self._finalize_cache = self._session.finalize()
125 return self._finalize_cache
127 def _parse_openai_chat(self, chunk: UpstreamChunk) -> _Parsed:
128 """Frame an OpenAI Chat chunk, honoring keepalives and ``[DONE]``."""
129 data = chunk.data.strip()
130 if not data:
131 return _Parsed(())
132 if data == "[DONE]":
133 return _Parsed((), terminal=True)
134 return self._accept(self._decode(OpenAIChatStreamChunk, data))
136 def _parse_openai_responses(self, chunk: UpstreamChunk) -> _Parsed:
137 """Frame an OpenAI Responses chunk by its ``type`` discriminator."""
138 if not chunk.data.strip():
139 return _Parsed(())
140 dto = self._decode(ResponsesEvent, chunk.data)
141 if dto.type == "response.completed":
142 return _Parsed((), terminal=True)
143 if dto.type in {"response.error", "response.failed", "response.incomplete"}:
144 return _Parsed((), error=True)
145 return self._accept(dto)
147 def _parse_claude(self, chunk: UpstreamChunk) -> _Parsed:
148 """Frame a Claude chunk by its ``type`` discriminator."""
149 if not chunk.data.strip():
150 return _Parsed(())
151 dto = self._decode(ClaudeStreamEvent, chunk.data)
152 if dto.type == "ping":
153 return _Parsed(())
154 if dto.type == "message_stop":
155 return _Parsed((), terminal=True)
156 if dto.type == "error":
157 return _Parsed((), error=True)
158 return self._accept(dto)
160 def _parse_gemini(self, chunk: UpstreamChunk) -> _Parsed:
161 """Frame a Gemini NDJSON line; never terminal or error."""
162 if not chunk.data.strip():
163 return _Parsed(())
164 return self._accept(self._decode(GeminiResponse, chunk.data))
166 def _decode(self, dto_type: type[Any], data: str) -> Any:
167 """Decode a JSON string into a wire DTO, mapping errors safely."""
168 try:
169 decoded = loads(data)
170 except (TypeError, ValueError) as error:
171 raise self._malformed(None) from error
172 if not isinstance(decoded, dict):
173 raise self._malformed(None)
174 try:
175 return dto_type.from_dict(decoded)
176 except RelayError as error:
177 raise self._malformed(error) from error
179 def _accept(self, dto: Any) -> _Parsed:
180 """Forward a DTO to the session, mapping session errors safely."""
181 try:
182 events = self._session.accept(dto)
183 except RelayError as error:
184 raise self._malformed(error) from error
185 return _Parsed(events)
187 def _malformed(self, error: RelayError | None) -> RelayGatewayError:
188 """Build the malformed-stream error from a safe public message."""
189 message = error.message if error is not None else "malformed upstream chunk"
190 return RelayGatewayError(
191 code=RelayGatewayErrorCode.UPSTREAM_MALFORMED,
192 message=message,
193 status_code=502,
194 request_id=self._request_id,
195 retryable=False,
196 )
199def _wire_events(events: tuple[Any, ...], terminal: bool) -> tuple[RelayWireEvent, ...]:
200 """Frame target DTOs as wire events, flagging the last as terminal.
202 Args:
203 events: Target DTOs to frame.
204 terminal: Whether the final event should carry ``terminal=True``.
206 Returns:
207 A ``RelayWireEvent`` tuple re-emitted from the target DTOs.
208 """
209 return tuple(
210 RelayWireEvent(
211 event=getattr(event, "type", None),
212 data=event.to_dict(),
213 terminal=terminal and index == len(events) - 1,
214 )
215 for index, event in enumerate(events)
216 )
219async def relay_stream(
220 upstream: RelayUpstreamProtocol,
221 request: UpstreamRequest,
222 parser: UpstreamEventParser,
223 cancel_handle: asyncio.Event | None = None,
224) -> AsyncIterator[RelayWireEvent]:
225 """Relay one upstream stream with cancellation and session lifecycle.
227 The ``async for`` inside this generator consumes exactly one upstream
228 chunk per consumer ``__anext__``: backpressure is inherent and there
229 is no buffering or prefetch.
231 Lifecycle: terminal frames finalize the session with no cancellation;
232 error frames and consumer disconnects cancel upstream once and
233 finalize truncated; streams that end without a terminal marker (for
234 example Gemini or a cut SSE stream) finalize truncated without
235 cancelling. Upstream ``cancel`` and session ``finalize`` each run at
236 most once, even across nested exception paths.
238 Args:
239 upstream: The upstream transport implementing
240 ``RelayUpstreamProtocol``.
241 request: The fully-resolved upstream request.
242 parser: Stateful session parser whose bookkeeping attributes
243 (``finalized``, ``truncated``, ``cancelled``) track the stream.
244 cancel_handle: Optional operator cancel handle from the stream
245 registry. When set, the relay cancels upstream once and
246 finalizes truncated at the next chunk boundary.
248 Yields:
249 Normalized ``RelayWireEvent`` values; terminal flag on the last
250 event of each finalize batch.
252 Raises:
253 RelayGatewayError: Malformed upstream framing or a session
254 rejection (502, never retryable).
255 asyncio.CancelledError: Upstream or consumer task cancellation;
256 always re-raised.
257 GeneratorExit: The consumer closed the generator mid-stream.
258 """
259 cancel_guard = False
261 async def cancel_once() -> None:
262 """Request upstream cancellation exactly once ever."""
263 nonlocal cancel_guard
264 if not cancel_guard:
265 cancel_guard = True
266 parser.cancelled = True
267 await upstream.cancel(request.request_id)
268 logger.info(
269 "relay_gateway_stream_cancelled",
270 request_id=request.request_id,
271 )
273 def finalize_once(truncated: bool) -> tuple[Any, ...]:
274 """Finalize the session once, remembering the truncation choice."""
275 if not parser.finalized:
276 parser.finalized = True
277 parser.truncated = truncated
278 return parser.finalize()
279 return ()
281 saw_terminal = False
282 saw_error = False
283 try:
284 stream_iter = cast("AsyncIterator[UpstreamChunk]", upstream.stream(request))
285 async for chunk in stream_iter:
286 if cancel_handle is not None and cancel_handle.is_set():
287 await cancel_once()
288 finalize_once(truncated=True)
289 break
290 try:
291 peaked = parser.parse(chunk)
292 except (
293 RelayGatewayError,
294 RelayError,
295 asyncio.CancelledError,
296 GeneratorExit,
297 ):
298 await cancel_once()
299 finalize_once(truncated=True)
300 raise
301 saw_terminal = saw_terminal or peaked.terminal
302 saw_error = saw_error or peaked.error
303 for wire in _wire_events(peaked.events, False):
304 yield wire
305 except (asyncio.CancelledError, GeneratorExit):
306 await cancel_once()
307 finalize_once(truncated=True)
308 raise
309 if saw_terminal:
310 for wire in _wire_events(finalize_once(truncated=False), True):
311 yield wire
312 return
313 if saw_error:
314 await cancel_once()
315 for wire in _wire_events(finalize_once(truncated=True), True):
316 yield wire
317 return
318 for wire in _wire_events(finalize_once(truncated=True), True):
319 yield wire