gproxy_protocol/adapt/responses_ws/
turn.rs

1use super::*;
2use crate::{
3    codec,
4    connection::WsFrame,
5    transform::generate::stream::responses::ResponsesStreamCollector,
6    wire::{
7        DeclaredFields,
8        openai::responses::websocket::{
9            ClientEvent, ErrorMessage, RequestMessage, ServerEvent, ServerMessage,
10        },
11    },
12};
13use futures_core::Stream;
14use std::{
15    pin::Pin,
16    task::{Context, Poll},
17};
18
19pub struct ResponsesWsTurn<'a> {
20    session: &'a mut ResponsesWsSession,
21    socket: Option<WebSocket>,
22    send: SendState,
23    collector: Option<ResponsesStreamCollector>,
24    received_bytes: usize,
25    received_frames: usize,
26    sent_bytes: usize,
27    sent_frames: usize,
28    terminal: bool,
29    lane: Option<String>,
30    failed: bool,
31}
32
33impl Unpin for ResponsesWsTurn<'_> {}
34enum SendState {
35    Ready(WsFrame),
36    Flush,
37    Done,
38}
39
40impl<'a> ResponsesWsTurn<'a> {
41    pub(super) fn new(
42        session: &'a mut ResponsesWsSession,
43        mut request: RequestMessage,
44    ) -> Result<Self, TransformError> {
45        let ClientEvent::ResponseCreate(body) = &mut request.event else {
46            return Err(TransformError::unsupported(
47                "responses.websocket.control",
48                "steering and injection require a duplex continuation driver, not a single response turn",
49            ));
50        };
51        if body.background.flatten() == Some(true) {
52            return Err(TransformError::unsupported(
53                "responses.websocket.background",
54                "WebSocket mode has no background generation control",
55            ));
56        }
57        body.stream = None;
58        body.background = None;
59        let event = request.into_declared();
60        let lane = event.stream_id.clone();
61        let bytes = codec::encode_json(
62            &event,
63            codec_limits(session.bounds.send_event.min(session.bounds.send_bytes)),
64        )
65        .map_err(|e| codec_error(e, true))?;
66        let len = bytes.len();
67        let frame =
68            WsFrame::Text(String::from_utf8(bytes.to_vec()).expect("JSON serialization is UTF-8"));
69        let collector = ResponsesStreamCollector::new(session.bounds.collector);
70        let socket = session
71            .socket
72            .take()
73            .ok_or_else(|| invalid("socket unavailable"))?;
74        // Closed-before-poll: even mem::forget(turn) cannot reopen this session.
75        session.busy = true;
76        Ok(Self {
77            session,
78            socket: Some(socket),
79            send: SendState::Ready(frame),
80            collector: Some(collector),
81            received_bytes: 0,
82            received_frames: 0,
83            sent_bytes: len,
84            sent_frames: 1,
85            terminal: false,
86            lane,
87            failed: false,
88        })
89    }
90    pub fn received_frames(&self) -> usize {
91        self.received_frames
92    }
93    pub fn received_bytes(&self) -> usize {
94        self.received_bytes
95    }
96    /// Reserved sends, including an uncertain pending readiness/flush operation.
97    pub fn reserved_send_frames(&self) -> usize {
98        self.sent_frames
99    }
100    pub fn reserved_send_bytes(&self) -> usize {
101        self.sent_bytes
102    }
103    pub fn failure_event(&self) -> Option<&StreamEvent> {
104        self.session.failure_event()
105    }
106    pub fn websocket_failure(&self) -> Option<&ErrorMessage> {
107        self.session.websocket_failure()
108    }
109    fn fail(&mut self, error: TransformError) -> TransformError {
110        self.failed = true;
111        self.socket.take();
112        self.session.poison();
113        error
114    }
115    fn account(&mut self, frame: &WsFrame) -> Result<(), TransformError> {
116        self.received_frames = self
117            .received_frames
118            .checked_add(1)
119            .ok_or_else(|| limit("responses.websocket.frames", "receive frame count overflow"))?;
120        let bytes = match frame {
121            WsFrame::Text(v) => v.len(),
122            WsFrame::Binary(v) | WsFrame::Ping(v) | WsFrame::Pong(v) => v.len(),
123            WsFrame::Close(Some(v)) => v
124                .reason
125                .len()
126                .checked_add(2)
127                .ok_or_else(|| limit("responses.websocket.frame", "close payload overflow"))?,
128            WsFrame::Close(None) => 0,
129        };
130        if bytes > self.session.bounds.receive_frame {
131            return Err(limit(
132                "responses.websocket.frame",
133                "receive frame byte cap exceeded",
134            ));
135        }
136        self.received_bytes = self
137            .received_bytes
138            .checked_add(bytes)
139            .filter(|v| *v <= self.session.bounds.receive_bytes)
140            .ok_or_else(|| {
141                limit(
142                    "responses.websocket.bytes",
143                    "aggregate receive byte cap exceeded",
144                )
145            })?;
146        if matches!(
147            frame,
148            WsFrame::Ping(_) | WsFrame::Pong(_) | WsFrame::Close(_)
149        ) && bytes > 125
150        {
151            return Err(invalid("invalid oversized WebSocket control frame"));
152        }
153        Ok(())
154    }
155    fn pong(&mut self, payload: bytes::Bytes) -> Result<(), TransformError> {
156        if payload.len() > self.session.bounds.send_frame {
157            return Err(limit(
158                "responses.websocket.pong",
159                "Pong exceeds host send frame cap",
160            ));
161        }
162        let bytes = self
163            .sent_bytes
164            .checked_add(payload.len())
165            .filter(|v| *v <= self.session.bounds.send_bytes)
166            .ok_or_else(|| {
167                limit(
168                    "responses.websocket.pong",
169                    "Pong exceeds aggregate send byte cap",
170                )
171            })?;
172        let frames = self
173            .sent_frames
174            .checked_add(1)
175            .ok_or_else(|| limit("responses.websocket.pong", "send frame count overflow"))?;
176        self.sent_bytes = bytes;
177        self.sent_frames = frames;
178        self.send = SendState::Ready(WsFrame::Pong(payload));
179        Ok(())
180    }
181    fn decode(&mut self, bytes: &[u8]) -> Result<StreamEvent, TransformError> {
182        if bytes.len() > self.session.bounds.receive_event {
183            return Err(limit(
184                "responses.websocket.event",
185                "JSON event cap exceeded",
186            ));
187        }
188        let limits = codec_limits(self.session.bounds.receive_event);
189        // Known nested WS errors must not fall back to the SSE Error variant
190        // and discard a malformed error object as an unknown extension.
191        let probe: serde_json::Value =
192            codec::decode_json(bytes, limits).map_err(|e| codec_error(e, false))?;
193        if probe.get("type").and_then(serde_json::Value::as_str) == Some("error")
194            && probe.get("error").is_some()
195        {
196            let error = codec::decode_json::<ErrorMessage>(bytes, limits)
197                .map_err(|e| codec_error(e, false))?
198                .into_declared();
199            let lane_matches = error.stream_id.is_none() || error.stream_id == self.lane;
200            self.session.websocket_failure = Some(Box::new(error.clone()));
201            return Err(if lane_matches {
202                TransformError::with_source(
203                    crate::transform::TransformErrorKind::InvalidResult,
204                    "responses.websocket",
205                    "native WebSocket error; see websocket_failure receipt",
206                    crate::transform::UpstreamError(error),
207                )
208            } else {
209                invalid("WebSocket error belongs to another lane")
210            });
211        }
212        let message = codec::decode_json::<ServerMessage>(bytes, limits)
213            .map_err(|e| codec_error(e, false))?
214            .into_declared();
215        if message.stream_id != self.lane {
216            return Err(invalid("WebSocket event belongs to another lane"));
217        }
218        let ServerEvent::Response(event) = message.event else {
219            return Err(TransformError::unsupported(
220                "responses.websocket.control",
221                "steering and injection events require a duplex continuation driver",
222            ));
223        };
224        if matches!(event, StreamEvent::Failed(_) | StreamEvent::Error(_)) {
225            self.session.failure = Some(Box::new(event.clone()));
226        }
227        let terminal = matches!(
228            event,
229            StreamEvent::Completed(_) | StreamEvent::Incomplete(_)
230        );
231        self.collector
232            .as_mut()
233            .ok_or_else(|| invalid("collector already consumed"))?
234            .push(event.clone())?;
235        if terminal {
236            let complete = self
237                .collector
238                .take()
239                .ok_or_else(|| invalid("collector already consumed"))?
240                .finish()?;
241            let socket = self
242                .socket
243                .take()
244                .ok_or_else(|| invalid("terminal socket unavailable"))?;
245            self.session.last_response_id = Some(complete.value.id);
246            self.session.last_status = complete.value.status;
247            self.session.last_report = complete.report;
248            self.session.socket = Some(socket);
249            self.session.busy = false;
250            self.terminal = true;
251        }
252        Ok(event)
253    }
254}
255
256impl Stream for ResponsesWsTurn<'_> {
257    type Item = Result<StreamEvent, TransformError>;
258    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
259        let this = self.get_mut();
260        if this.terminal || this.failed {
261            return Poll::Ready(None);
262        }
263        let Some(socket) = this.socket.as_mut() else {
264            return Poll::Ready(Some(Err(this.fail(invalid("socket closed")))));
265        };
266        match poll_send(socket, &mut this.send, cx) {
267            Poll::Pending => return Poll::Pending,
268            Poll::Ready(Err(e)) => return Poll::Ready(Some(Err(this.fail(e)))),
269            Poll::Ready(Ok(())) => {}
270        }
271        if this.received_bytes >= this.session.bounds.receive_bytes {
272            return Poll::Ready(Some(Err(this.fail(limit(
273                "responses.websocket.bytes",
274                "receive byte budget exhausted before terminal",
275            )))));
276        }
277        let frame = match socket.incoming.as_mut().poll_next(cx) {
278            Poll::Pending => return Poll::Pending,
279            Poll::Ready(None) => {
280                return Poll::Ready(Some(Err(
281                    this.fail(invalid("socket ended before terminal response"))
282                )));
283            }
284            Poll::Ready(Some(Err(e))) => return Poll::Ready(Some(Err(this.fail(host_error(e))))),
285            Poll::Ready(Some(Ok(frame))) => frame,
286        };
287        if let Err(e) = this.account(&frame) {
288            return Poll::Ready(Some(Err(this.fail(e))));
289        }
290        match frame {
291            WsFrame::Ping(payload) => {
292                if let Err(e) = this.pong(payload) {
293                    return Poll::Ready(Some(Err(this.fail(e))));
294                }
295                cx.waker().wake_by_ref();
296                Poll::Pending
297            }
298            WsFrame::Pong(_) => {
299                cx.waker().wake_by_ref();
300                Poll::Pending
301            }
302            WsFrame::Close(close) => {
303                this.session.peer_close = close;
304                Poll::Ready(Some(Err(
305                    this.fail(invalid("socket closed before terminal response"))
306                )))
307            }
308            WsFrame::Binary(_) if !this.session.bounds.allow_binary => Poll::Ready(Some(Err(this
309                .fail(TransformError::unsupported(
310                    "responses.websocket.frame",
311                    "binary JSON server messages are disabled",
312                ))))),
313            WsFrame::Text(text) => {
314                Poll::Ready(Some(this.decode(text.as_bytes()).map_err(|e| this.fail(e))))
315            }
316            WsFrame::Binary(bytes) => {
317                Poll::Ready(Some(this.decode(&bytes).map_err(|e| this.fail(e))))
318            }
319        }
320    }
321}
322
323impl Drop for ResponsesWsTurn<'_> {
324    fn drop(&mut self) {
325        if !self.terminal {
326            self.socket.take();
327            self.session.poison();
328        }
329    }
330}
331
332fn poll_send(
333    socket: &mut WebSocket,
334    state: &mut SendState,
335    cx: &mut Context<'_>,
336) -> Poll<Result<(), TransformError>> {
337    loop {
338        match state {
339            SendState::Ready(_) => match socket.outgoing.as_mut().poll_ready(cx) {
340                Poll::Pending => return Poll::Pending,
341                Poll::Ready(Err(e)) => return Poll::Ready(Err(host_error(e))),
342                Poll::Ready(Ok(())) => {
343                    // Move this exact pending frame once. Readiness/flush Pending
344                    // never recreates response.create or replaces a queued Pong.
345                    let SendState::Ready(frame) = std::mem::replace(state, SendState::Flush) else {
346                        // The enclosing arm already matched SendState::Ready.
347                        unreachable!()
348                    };
349                    if let Err(e) = socket.outgoing.as_mut().start_send(frame) {
350                        return Poll::Ready(Err(host_error(e)));
351                    }
352                }
353            },
354            SendState::Flush => match socket.outgoing.as_mut().poll_flush(cx) {
355                Poll::Pending => return Poll::Pending,
356                Poll::Ready(Err(e)) => return Poll::Ready(Err(host_error(e))),
357                Poll::Ready(Ok(())) => *state = SendState::Done,
358            },
359            SendState::Done => return Poll::Ready(Ok(())),
360        }
361    }
362}
gproxy_protocol/adapt/responses_ws/turn.rs

gproxy_protocol/adapt/responses_ws/
turn.rs
