agent_engine/runtime/openai/
wire.rs

1//! SSE wire decoder for OpenAI-compatible streams. Ported from
2//! `openai-runtime::stream`, with `StreamEvent` renamed to `OaiEvent`.
3
4use super::types::{FunctionCall, OaiEvent, ToolCall};
5use serde::Deserialize;
6use std::collections::HashMap;
7
8// ─── Wire types ──────────────────────────────────────────────────────────────
9
10#[derive(Debug, Deserialize)]
11pub struct RawChunk {
12    #[serde(default)]
13    pub choices: Vec<RawChoice>,
14    #[serde(default)]
15    pub usage: Option<RawUsage>,
16}
17
18#[derive(Debug, Deserialize)]
19pub struct RawChoice {
20    #[serde(default)]
21    pub delta: Option<RawDelta>,
22    #[serde(default)]
23    pub finish_reason: Option<String>,
24}
25
26#[derive(Debug, Deserialize, Default)]
27pub struct RawDelta {
28    #[serde(default)]
29    pub role: Option<String>,
30    #[serde(default)]
31    pub content: Option<String>,
32    #[serde(default)]
33    pub tool_calls: Option<Vec<RawToolCallDelta>>,
34}
35
36#[derive(Debug, Deserialize)]
37pub struct RawToolCallDelta {
38    #[serde(default)]
39    pub index: Option<u32>,
40    #[serde(default)]
41    pub id: Option<String>,
42    #[serde(rename = "type", default)]
43    pub kind: Option<String>,
44    #[serde(default)]
45    pub function: Option<RawFunctionDelta>,
46}
47
48#[derive(Debug, Deserialize)]
49pub struct RawFunctionDelta {
50    #[serde(default)]
51    pub name: Option<String>,
52    #[serde(default)]
53    pub arguments: Option<String>,
54}
55
56#[derive(Debug, Deserialize)]
57pub struct RawUsage {
58    #[serde(default)]
59    pub prompt_tokens: Option<u32>,
60    #[serde(default)]
61    pub completion_tokens: Option<u32>,
62    #[serde(default)]
63    pub prompt_tokens_details: Option<RawPromptTokensDetails>,
64}
65
66#[derive(Debug, Deserialize)]
67pub struct RawPromptTokensDetails {
68    #[serde(default)]
69    pub cached_tokens: Option<u32>,
70}
71
72/// Legacy text-only SSE line parser. Kept for simple use-cases; the main
73/// decoder path is `StreamDecoder`.
74pub fn parse_sse_line(line: &str) -> Option<OaiEvent> {
75    let line = line.trim_end_matches('\r');
76    if line.is_empty() || line.starts_with(':') {
77        return None;
78    }
79    let data = line.strip_prefix("data:")?.trim_start();
80    if data == "[DONE]" {
81        return Some(OaiEvent::Done);
82    }
83    let chunk: RawChunk = serde_json::from_str(data).ok()?;
84    let delta = chunk.choices.into_iter().next()?.delta?;
85    if let Some(role) = delta.role {
86        return Some(OaiEvent::RoleStart(role));
87    }
88    if let Some(content) = delta.content {
89        if !content.is_empty() {
90            return Some(OaiEvent::TextDelta(content));
91        }
92    }
93    None
94}
95
96// ─── Accumulator ─────────────────────────────────────────────────────────────
97
98#[derive(Debug, Default)]
99pub struct ToolCallAccumulator {
100    pub id: String,
101    pub name: String,
102    pub arguments: String,
103    started: bool,
104}
105
106#[derive(Debug)]
107pub struct StreamDecoder {
108    pub calls: HashMap<u32, ToolCallAccumulator>,
109    pub truncated: bool,
110    pub completed: bool,
111    /// First `finish_reason` string observed on the wire (set-once). Raw
112    /// provider value — read it only through a normalizing mapper (trace
113    /// records store the normalized enum, never this string).
114    pub finish_reason: Option<String>,
115    role_emitted: bool,
116    done_emitted: bool,
117}
118
119impl Default for StreamDecoder {
120    fn default() -> Self {
121        Self::new()
122    }
123}
124
125impl StreamDecoder {
126    pub fn new() -> Self {
127        Self {
128            calls: HashMap::new(),
129            truncated: false,
130            completed: false,
131            finish_reason: None,
132            role_emitted: false,
133            done_emitted: false,
134        }
135    }
136
137    pub fn push_line<E: Extend<OaiEvent>>(&mut self, line: &str, sink: &mut E) {
138        let line = line.trim_end_matches('\r');
139        if line.is_empty() || line.starts_with(':') {
140            return;
141        }
142        let payload = match line.strip_prefix("data:").map(str::trim_start) {
143            Some(p) => p,
144            None => return,
145        };
146        if payload == "[DONE]" {
147            self.finish(sink);
148            return;
149        }
150        match serde_json::from_str::<RawChunk>(payload) {
151            Ok(chunk) => self.push_chunk(chunk, sink),
152            // Privacy (spec §5.1): `payload` is provider-controlled and may
153            // echo the request — never include its bytes in a warning that
154            // feeds the log sink. serde's error carries position info only.
155            Err(e) => sink.extend(Some(OaiEvent::Warning(format!(
156                "sse parse error: {e} (payload {} bytes, not logged)",
157                payload.len()
158            )))),
159        }
160    }
161
162    fn push_chunk<E: Extend<OaiEvent>>(&mut self, chunk: RawChunk, sink: &mut E) {
163        for choice in chunk.choices {
164            let is_finish = choice.finish_reason.is_some();
165            if let Some(delta) = choice.delta {
166                if let Some(role) = delta.role {
167                    if !self.role_emitted {
168                        self.role_emitted = true;
169                        sink.extend(Some(OaiEvent::RoleStart(role)));
170                    }
171                }
172                if let Some(text) = delta.content {
173                    if !text.is_empty() {
174                        sink.extend(Some(OaiEvent::TextDelta(text)));
175                    }
176                }
177                // Process tool_calls — but de-dup finish-chunk re-sends.
178                // Some providers re-send the full tool_calls on the finish frame.
179                // We only skip if the chunk has finish_reason AND the tool_call
180                // has an id (indicating a full re-send, not a final argument delta).
181                if let Some(tcs) = delta.tool_calls {
182                    for tc in tcs {
183                        let is_resend =
184                            is_finish && tc.id.as_ref().is_some_and(|id| !id.is_empty());
185                        if !is_resend {
186                            self.apply_tool_call_delta(tc, sink);
187                        }
188                    }
189                }
190            }
191            if let Some(reason) = choice.finish_reason {
192                if self.finish_reason.is_none() {
193                    self.finish_reason = Some(reason.clone());
194                }
195                match reason.as_str() {
196                    "tool_calls" => self.flush_complete(sink),
197                    "length" => {
198                        if !self.calls.is_empty() {
199                            self.truncated = true;
200                            self.flush_complete(sink);
201                        }
202                    }
203                    "stop" | "content_filter" => {}
204                    // Privacy (spec §5.1): the value is provider-controlled;
205                    // keep it out of warnings that feed the log sink.
206                    _ => sink.extend(Some(OaiEvent::Warning(
207                        "unknown finish_reason (value not logged)".to_string(),
208                    ))),
209                }
210            }
211        }
212        if let Some(u) = chunk.usage {
213            let cached = u
214                .prompt_tokens_details
215                .and_then(|d| d.cached_tokens)
216                .unwrap_or(0);
217            sink.extend(Some(OaiEvent::Usage {
218                prompt_tokens: u.prompt_tokens.unwrap_or(0),
219                completion_tokens: u.completion_tokens.unwrap_or(0),
220                cached_tokens: cached,
221            }));
222        }
223    }
224
225    fn apply_tool_call_delta<E: Extend<OaiEvent>>(&mut self, tc: RawToolCallDelta, sink: &mut E) {
226        let idx = tc.index.unwrap_or(0);
227        let acc = self.calls.entry(idx).or_default();
228
229        if let Some(id) = tc.id {
230            if !id.is_empty() {
231                acc.id = id;
232            }
233        }
234        if let Some(f) = tc.function {
235            if let Some(n) = f.name {
236                if !n.is_empty() {
237                    acc.name = n;
238                }
239            }
240            if !acc.started && !acc.id.is_empty() && !acc.name.is_empty() {
241                acc.started = true;
242                sink.extend(Some(OaiEvent::ToolCallStart {
243                    index: idx,
244                    id: acc.id.clone(),
245                    name: acc.name.clone(),
246                }));
247            }
248            if let Some(args) = f.arguments {
249                if !args.is_empty() {
250                    acc.arguments.push_str(&args);
251                    sink.extend(Some(OaiEvent::ToolCallArgumentsDelta {
252                        index: idx,
253                        id: acc.id.clone(),
254                        delta: args,
255                    }));
256                }
257            }
258        } else if !acc.started && !acc.id.is_empty() && !acc.name.is_empty() {
259            acc.started = true;
260            sink.extend(Some(OaiEvent::ToolCallStart {
261                index: idx,
262                id: acc.id.clone(),
263                name: acc.name.clone(),
264            }));
265        }
266    }
267
268    pub fn finish<E: Extend<OaiEvent>>(&mut self, sink: &mut E) {
269        self.flush_complete(sink);
270        if !self.done_emitted {
271            self.done_emitted = true;
272            sink.extend(Some(OaiEvent::Done));
273        }
274    }
275
276    fn flush_complete<E: Extend<OaiEvent>>(&mut self, sink: &mut E) {
277        if self.completed || self.calls.is_empty() {
278            return;
279        }
280        self.completed = true;
281        let mut entries: Vec<(u32, ToolCallAccumulator)> = self.calls.drain().collect();
282        entries.sort_by_key(|(k, _)| *k);
283        let calls: Vec<ToolCall> = entries
284            .into_iter()
285            .map(|(_, acc)| ToolCall {
286                id: acc.id,
287                kind: "function".to_string(),
288                function: FunctionCall {
289                    name: acc.name,
290                    arguments: acc.arguments,
291                },
292            })
293            .collect();
294        sink.extend(Some(OaiEvent::ToolCallsComplete {
295            calls,
296            truncated: self.truncated,
297        }));
298    }
299}
300
301#[cfg(test)]
302mod tests {
303    //! Phase 1 privacy (spec §5.1): SSE payloads are provider-controlled and
304    //! may echo the request. Decoder warnings feed `tracing::warn!` (see
305    //! `translate::oai_event_to_stream_event`), so no provider-controlled
306    //! bytes may ride along in a warning message.
307
308    use super::*;
309
310    #[test]
311    fn parse_error_warning_never_contains_payload_bytes() {
312        const SENTINEL: &str = "PH1-WIRE-SENTINEL-4c1a-RAW";
313        let mut decoder = StreamDecoder::new();
314        let mut sink: Vec<OaiEvent> = Vec::new();
315        decoder.push_line(&format!("data: {{\"broken\": \"{SENTINEL}"), &mut sink);
316        let [OaiEvent::Warning(w)] = sink.as_slice() else {
317            panic!("malformed payload must produce exactly one warning, got {sink:?}");
318        };
319        assert!(w.contains("sse parse error"), "{w}");
320        assert!(
321            !w.contains(SENTINEL),
322            "provider payload bytes leaked into a warning destined for logs: {w}"
323        );
324    }
325
326    #[test]
327    fn unknown_finish_reason_warning_never_contains_provider_value() {
328        const SENTINEL: &str = "PH1-FINISH-SENTINEL-7e2b-RAW";
329        let mut decoder = StreamDecoder::new();
330        let mut sink: Vec<OaiEvent> = Vec::new();
331        decoder.push_line(
332            &format!("data: {{\"choices\":[{{\"delta\":{{}},\"finish_reason\":\"{SENTINEL}\"}}]}}"),
333            &mut sink,
334        );
335        let warning = sink.iter().find_map(|e| match e {
336            OaiEvent::Warning(w) => Some(w.clone()),
337            _ => None,
338        });
339        let w = warning.expect("unknown finish_reason must still warn");
340        assert!(
341            !w.contains(SENTINEL),
342            "provider-controlled finish_reason leaked into a warning destined for logs: {w}"
343        );
344    }
345}
agent_engine/runtime/openai/wire.rs

agent_engine/runtime/openai/
wire.rs
