gproxy_protocol/adapt/generate/stream/
websocket.rs

1//! Concrete Responses WebSocket transport using the same incremental pair states.
2//! The host binds authentication to its target and state scope at connection time.
3
4use super::super::GenerationStateAccess;
5use super::{StreamChunk, StreamInvocation, bridge::StreamBridge};
6use crate::{
7    HttpBody, WireRequest, WireResponse,
8    adapt::responses_ws as ws,
9    capability::{StateStore, Upstream},
10    transform::TransformError,
11    wire::{
12        DeclaredFields,
13        openai::responses::{
14            self as r,
15            stream::StreamEvent,
16            websocket::{ClientEvent, RequestMessage},
17        },
18    },
19};
20
21pub struct GenerationWsSession {
22    native: ws::ResponsesWsSession,
23    /// The state the connection was opened under; each turn must use it.
24    binding: super::binding::StateBinding,
25}
26
27#[allow(clippy::large_enum_variant)]
28pub enum GenerationWsConnect {
29    Connected {
30        handshake: WireResponse<()>,
31        session: GenerationWsSession,
32    },
33    Rejected(WireResponse<HttpBody>),
34}
35
36impl GenerationWsSession {
37    pub fn native(&self) -> &ws::ResponsesWsSession {
38        &self.native
39    }
40}
41
42pub async fn connect<U: Upstream, S: StateStore>(
43    upstream: &U,
44    target: &U::Target,
45    request: WireRequest<()>,
46    limits: ws::ResponsesWsLimits,
47    state: &GenerationStateAccess<'_, S>,
48) -> Result<GenerationWsConnect, ws::ResponsesWsConnectError> {
49    let binding = super::binding::StateBinding::new(state);
50    match ws::connect(upstream, target, request, limits).await? {
51        ws::ResponsesWsConnect::Connected { handshake, session } => {
52            Ok(GenerationWsConnect::Connected {
53                handshake,
54                session: GenerationWsSession {
55                    native: session,
56                    binding,
57                },
58            })
59        }
60        ws::ResponsesWsConnect::Rejected(response) => Ok(GenerationWsConnect::Rejected(response)),
61    }
62}
63
64/// Owns the native turn while borrowing its caller-owned conversion progress.
65/// Dropping a nonterminal native turn poisons the connection; it never resends.
66pub struct GenerationWsTurn<'a, B: StreamBridge<NativeEvent = StreamEvent>> {
67    invocation: &'a mut StreamInvocation<B>,
68    turn: Option<ws::ResponsesWsTurn<'a>>,
69    failure: Option<StreamEvent>,
70    websocket_failure: Option<r::websocket::ErrorMessage>,
71}
72
73impl<B: StreamBridge<NativeEvent = StreamEvent>> GenerationWsTurn<'_, B> {
74    pub async fn next<S: StateStore>(
75        &mut self,
76        state: &GenerationStateAccess<'_, S>,
77    ) -> Result<Option<StreamChunk<B::ClientEvent>>, TransformError> {
78        let Some(turn) = &mut self.turn else {
79            return Err(super::invalid("WebSocket generation turn failed"));
80        };
81        let result = self.invocation.next_inner(state, Some(turn)).await;
82        if result.is_err() && self.invocation.failed {
83            self.failure = turn.failure_event().cloned();
84            self.websocket_failure = turn.websocket_failure().cloned();
85            self.turn.take();
86        }
87        result
88    }
89    pub async fn collect<S: StateStore>(
90        &mut self,
91        state: &GenerationStateAccess<'_, S>,
92    ) -> Result<super::invoke::ClientFull<B>, TransformError> {
93        while self.next(state).await?.is_some() {}
94        let result = self
95            .invocation
96            .client_result()
97            .ok_or_else(|| super::missing("WebSocket client result is not durably committed"))?
98            .clone();
99        crate::codec::encode_json(&result, self.invocation.settings.codec)
100            .map_err(super::codec_error)?;
101        Ok(result)
102    }
103    pub fn invocation(&self) -> &StreamInvocation<B> {
104        self.invocation
105    }
106    pub fn failure_event(&self) -> Option<&StreamEvent> {
107        self.failure
108            .as_ref()
109            .or_else(|| self.turn.as_ref().and_then(|turn| turn.failure_event()))
110    }
111    pub fn websocket_failure(&self) -> Option<&r::websocket::ErrorMessage> {
112        self.websocket_failure
113            .as_ref()
114            .or_else(|| self.turn.as_ref().and_then(|turn| turn.websocket_failure()))
115    }
116}
117
118impl<B: StreamBridge<NativeEvent = StreamEvent, NativeRequest = r::GenerateContentRequestBody>>
119    StreamInvocation<B>
120{
121    /// Sends response.create at most once. The invocation and the connection
122    /// must both have been prepared against this state.
123    pub async fn start_websocket<'a, S: StateStore>(
124        &'a mut self,
125        session: &'a mut GenerationWsSession,
126        state: &GenerationStateAccess<'_, S>,
127    ) -> Result<GenerationWsTurn<'a, B>, TransformError> {
128        self.start_websocket_on_lane(session, None, state).await
129    }
130    pub async fn start_websocket_on_lane<'a, S: StateStore>(
131        &'a mut self,
132        session: &'a mut GenerationWsSession,
133        lane: Option<String>,
134        state: &GenerationStateAccess<'_, S>,
135    ) -> Result<GenerationWsTurn<'a, B>, TransformError> {
136        self.binding.check(state)?;
137        session.binding.check(state)?;
138        if self.sent || self.failed {
139            return Err(super::conflict("invocation already sent or failed"));
140        }
141        if self
142            .websocket_lane
143            .as_ref()
144            .is_some_and(|known| known != &lane)
145        {
146            return Err(super::conflict("WebSocket turn lane changed"));
147        }
148        let mut request = self.target.clone().into_declared();
149        // WS is intrinsically streaming; this is the transport control added by
150        // prepare_stream, not a user-requested change of generation semantics.
151        request.stream = None;
152        self.websocket_lane = Some(lane.clone());
153        let turn = session.native.turn_message(RequestMessage {
154            stream_id: lane,
155            generate: None,
156            event: ClientEvent::ResponseCreate(request),
157        })?;
158        self.sent = true;
159        Ok(GenerationWsTurn {
160            invocation: self,
161            turn: Some(turn),
162            failure: None,
163            websocket_failure: None,
164        })
165    }
166}
167
168impl<B: StreamBridge<ClientEvent = StreamEvent>> StreamInvocation<B> {
169    /// Select bare, individually bounded JSON server messages before HTTP start.
170    /// Each returned nonempty chunk is exactly one WebSocket Text payload.
171    pub fn responses_websocket_output(&mut self) -> Result<(), TransformError> {
172        self.responses_websocket_output_on_lane(None)
173    }
174    pub fn responses_websocket_output_on_lane(
175        &mut self,
176        lane: Option<String>,
177    ) -> Result<(), TransformError> {
178        if self.sent || self.failed {
179            return Err(super::conflict("select WebSocket output before send"));
180        }
181        self.encoder = super::output::Encoder::WebSocket {
182            total: 0,
183            events: 0,
184            limits: self.settings.codec,
185            lane,
186        };
187        Ok(())
188    }
189}
190
191/// Decode only the declared response.create shape. Unknown extensions are scrubbed.
192pub fn decode_request(
193    bytes: &[u8],
194    limits: crate::codec::CodecLimits,
195) -> Result<r::GenerateContentRequestBody, TransformError> {
196    let message = decode_message(bytes, limits)?;
197    if message.stream_id.is_some() {
198        return Err(TransformError::shape(
199            "responses.websocket.stream_id",
200            "use decode_message and preserve the explicit output lane",
201        ));
202    }
203    let ClientEvent::ResponseCreate(request) = message.event else {
204        return Err(TransformError::unsupported(
205            "responses.websocket.control",
206            "steering and injection have no equivalent cross-protocol generation operation",
207        ));
208    };
209    Ok(request)
210}
211
212/// Decode the WS envelope for cross-protocol generation. Warmup has no
213/// equivalent operation on Chat, Claude or Gemini and is rejected explicitly.
214pub fn decode_message(
215    bytes: &[u8],
216    limits: crate::codec::CodecLimits,
217) -> Result<RequestMessage, TransformError> {
218    let message = decode_session_message(bytes, limits)?;
219    if message.generate == Some(false) {
220        return Err(TransformError::unsupported(
221            "responses.websocket.generate",
222            "prefill-only warmup requires a Responses session driver",
223        ));
224    }
225    if !matches!(message.event, ClientEvent::ResponseCreate(_)) {
226        return Err(TransformError::unsupported(
227            "responses.websocket.control",
228            "control messages require a Responses session driver",
229        ));
230    }
231    Ok(message)
232}
233
234/// Decode controls for an actual duplex session. Single-response adapters
235/// keep using `decode_message` and cannot silently accept these controls.
236pub fn decode_session_message(
237    bytes: &[u8],
238    limits: crate::codec::CodecLimits,
239) -> Result<RequestMessage, TransformError> {
240    let mut message = crate::codec::decode_json::<RequestMessage>(bytes, limits)
241        .map_err(|error| {
242            let kind = if error.kind() == crate::codec::CodecErrorKind::Limit {
243                crate::transform::TransformErrorKind::Limit
244            } else {
245                crate::transform::TransformErrorKind::InvalidInput
246            };
247            TransformError::with_source(
248                kind,
249                "responses.websocket.request",
250                error.to_string(),
251                error,
252            )
253        })?
254        .into_declared();
255
256    if let Some(lane) = &message.stream_id
257        && (lane.is_empty()
258            || lane.len() > 256
259            || !lane
260                .bytes()
261                .all(|b| b.is_ascii_alphanumeric() || b"_.-".contains(&b)))
262    {
263        return Err(TransformError::shape(
264            "responses.websocket.stream_id",
265            "invalid lane",
266        ));
267    }
268    let ClientEvent::ResponseCreate(request) = &mut message.event else {
269        return Ok(message);
270    };
271    if request.background.flatten() == Some(true) {
272        return Err(TransformError::unsupported(
273            "responses.websocket.background",
274            "WebSocket mode has no background generation control",
275        ));
276    }
277    request.stream = None;
278    request.background = None;
279    Ok(message)
280}
281
282mod image_resources;
gproxy_protocol/adapt/generate/stream/websocket.rs

gproxy_protocol/adapt/generate/stream/
websocket.rs
