gproxy_transform/transform/
stream_adapter.rs

1//! Runtime SSE adaptation for cross-protocol content-generation streams.
2
3mod buffered;
4mod responses;
5mod synthesize;
6
7use super::common::sse::{SseDecoder, SseFrame};
8use super::{TransformContext, TransformPair, dispatch};
9use crate::protocol::ContentGenerationKind;
10use crate::protocol::openai::ResponseStreamEvent;
11
12use responses::ResponsesStreamState;
13
14pub use buffered::{aggregate_buffered, convert_buffered};
15pub use responses::ResponsesStreamNormalizer;
16pub use synthesize::synthesize_sse;
17
18pub struct SseTransformer {
19    decoder: SseDecoder,
20    /// Reverse pair: upstream kind to inbound kind.
21    pair: TransformPair,
22    ctx: TransformContext,
23    inbound: ContentGenerationKind,
24    responses: Option<ResponsesStreamState>,
25    skipped: u64,
26}
27
28impl SseTransformer {
29    pub fn new(pair: TransformPair, ctx: TransformContext, inbound: ContentGenerationKind) -> Self {
30        Self {
31            decoder: SseDecoder::new(),
32            pair,
33            ctx,
34            inbound,
35            responses: matches!(
36                inbound,
37                ContentGenerationKind::OpenAiResponses
38                    | ContentGenerationKind::OpenAiResponsesWebSocket
39            )
40            .then(ResponsesStreamState::default),
41            skipped: 0,
42        }
43    }
44
45    /// Feed one upstream chunk; returns encoded inbound bytes (possibly empty).
46    pub fn push(&mut self, chunk: &[u8]) -> Vec<u8> {
47        let mut out = Vec::new();
48        for frame in self.decoder.push(chunk) {
49            self.convert_into(frame, &mut out);
50        }
51        out
52    }
53
54    /// Flush the trailing frame and emit the inbound terminator.
55    pub fn finish(&mut self) -> Vec<u8> {
56        let mut out = Vec::new();
57        if let Some(frame) = self.decoder.finish() {
58            self.convert_into(frame, &mut out);
59        }
60        if let Some(responses) = self.responses.as_mut() {
61            for event in responses.finish() {
62                encode_responses_event(&event, &mut out);
63            }
64        }
65        if self.inbound == ContentGenerationKind::OpenAiChatCompletions {
66            out.extend_from_slice(b"data: [DONE]\n\n");
67        }
68        if self.skipped > 0 {
69            tracing::warn!(
70                skipped = self.skipped,
71                "stream transform skipped unconvertible frames"
72            );
73        }
74        out
75    }
76
77    fn convert_into(&mut self, frame: SseFrame, out: &mut Vec<u8>) {
78        if frame.data.trim() == "[DONE]" {
79            return;
80        }
81        match dispatch::stream_event(self.pair, &self.ctx, &frame.data) {
82            Ok(dispatch::StreamEventOut::Encoded { event, data }) => {
83                encode_frame(self.inbound, event.as_deref(), &data, out);
84            }
85            Ok(dispatch::StreamEventOut::Responses(event)) => {
86                if let Some(responses) = self.responses.as_mut() {
87                    for event in responses.push(*event) {
88                        encode_responses_event(&event, out);
89                    }
90                } else {
91                    // Defensive: Responses events only occur with a Responses
92                    // inbound, where the aggregation state is always present.
93                    encode_responses_event(&event, out);
94                }
95            }
96            Err(_) => self.skipped += 1,
97        }
98    }
99}
100
101/// Encode one converted event in the inbound wire format. Claude and Responses
102/// inbound streams carry named SSE events (missing names fall back to
103/// "message", as before the typed path); chat and Gemini are data-only.
104fn encode_frame(kind: ContentGenerationKind, event: Option<&str>, data: &str, out: &mut Vec<u8>) {
105    use ContentGenerationKind as K;
106    let frame = match kind {
107        K::ClaudeMessages | K::OpenAiResponses | K::OpenAiResponsesWebSocket => {
108            SseFrame::event(event.unwrap_or("message"), data)
109        }
110        K::OpenAiChatCompletions | K::GeminiGenerateContent => SseFrame::data(data),
111    };
112    out.extend_from_slice(frame.encode().as_bytes());
113}
114
115/// Serialize + encode one Responses event as a named SSE frame.
116fn encode_responses_event(event: &ResponseStreamEvent, out: &mut Vec<u8>) {
117    let Ok(data) = serde_json::to_string(event) else {
118        return;
119    };
120    let frame = SseFrame::event(event.event_name().unwrap_or("message"), data);
121    out.extend_from_slice(frame.encode().as_bytes());
122}
123
124#[cfg(test)]
125mod tests;
gproxy_transform/transform/stream_adapter.rs

gproxy_transform/transform/
stream_adapter.rs
