Skip to main content

monoloop_interpreter/
openai_chat.rs

1//! OpenAI Chat Completions v1 streaming SSE → canonical fragments.
2//!
3//! Frames `data: <json>` / `data: [DONE]` across arbitrary byte boundaries.
4//! Emits complete text units via the shared segmenter path and
5//! `ToolRequestReady` only when tool arguments are complete valid JSON.
6//! OpenAI Responses and non-streaming JSON are unsupported.
7
8use crate::acp::{AcpFragment, ToolSignal};
9use monoloop_contracts::{InterpreterError, InterpreterErrorKind, TextChannel, ToolActionId};
10use serde_json::Value;
11use std::collections::HashMap;
12
13/// Default accepted choice index for the first product path.
14pub const DEFAULT_CHOICE_INDEX: u32 = 0;
15
16/// Incremental SSE assembler for one interpretation.
17#[derive(Debug, Default)]
18pub struct OpenAiSseState {
19    /// Incomplete trailing line bytes.
20    line_carry: Vec<u8>,
21    /// Data lines of the current SSE event.
22    event_data_lines: Vec<String>,
23    /// Partial tool calls keyed by provider tool-call index within the choice.
24    tools: HashMap<u32, PartialToolCall>,
25    /// Whether `data: [DONE]` was observed.
26    saw_done: bool,
27    /// Selected choice index.
28    choice_index: u32,
29    /// Maximum SSE line bytes.
30    max_line_bytes: usize,
31    /// Maximum single event data bytes.
32    max_event_bytes: usize,
33    /// Maximum tool argument accumulation bytes.
34    max_tool_arg_bytes: usize,
35}
36
37#[derive(Debug, Default)]
38struct PartialToolCall {
39    id: String,
40    name: String,
41    arguments: String,
42    emitted_waiting: bool,
43    ready: bool,
44}
45
46impl OpenAiSseState {
47    /// Construct with bounds from interpretation limits.
48    pub fn new(max_line_bytes: usize, max_event_bytes: usize, max_tool_arg_bytes: usize) -> Self {
49        Self {
50            choice_index: DEFAULT_CHOICE_INDEX,
51            max_line_bytes: max_line_bytes.max(64),
52            max_event_bytes: max_event_bytes.max(64),
53            max_tool_arg_bytes: max_tool_arg_bytes.max(64),
54            ..Default::default()
55        }
56    }
57
58    /// Whether a terminal `[DONE]` marker was observed.
59    pub fn saw_done(&self) -> bool {
60        self.saw_done
61    }
62
63    /// Ingest a raw transport chunk; return fragments to apply.
64    pub fn push_bytes(&mut self, chunk: &[u8]) -> Result<Vec<AcpFragment>, InterpreterError> {
65        if self.saw_done {
66            return Ok(Vec::new());
67        }
68        let mut out = Vec::new();
69        for &b in chunk {
70            if b == b'\n' {
71                let mut line = std::mem::take(&mut self.line_carry);
72                // Strip CR for CRLF.
73                if line.last() == Some(&b'\r') {
74                    line.pop();
75                }
76                self.handle_line(&line, &mut out)?;
77            } else {
78                if self.line_carry.len() >= self.max_line_bytes {
79                    return Err(InterpreterError::new(
80                        InterpreterErrorKind::FrameLimitExceeded,
81                        "SSE line exceeds bound",
82                    ));
83                }
84                self.line_carry.push(b);
85            }
86        }
87        Ok(out)
88    }
89
90    /// Flush on clean end: incomplete trailing event without a `DONE` marker fails closed.
91    pub fn seal_clean(&mut self) -> Result<Vec<AcpFragment>, InterpreterError> {
92        if !self.line_carry.is_empty() {
93            let line = std::mem::take(&mut self.line_carry);
94            let mut out = Vec::new();
95            self.handle_line(&line, &mut out)?;
96            if !out.is_empty() {
97                // fall through
98            }
99        }
100        if !self.event_data_lines.is_empty() {
101            // Incomplete event at EOF.
102            return Err(InterpreterError::new(
103                InterpreterErrorKind::MalformedFrame,
104                "incomplete SSE event at end of stream",
105            ));
106        }
107        if !self.saw_done {
108            return Err(InterpreterError::new(
109                InterpreterErrorKind::MalformedFrame,
110                "missing [DONE] terminator",
111            ));
112        }
113        // Incomplete tool args never become Ready.
114        let mut frags = Vec::new();
115        for partial in self.tools.values() {
116            if !partial.ready && !partial.id.is_empty() {
117                frags.push(AcpFragment::Tool {
118                    action_id: ToolActionId::new(partial.id.clone()),
119                    signal: ToolSignal::Waiting {
120                        tool_name: if partial.name.is_empty() {
121                            None
122                        } else {
123                            Some(partial.name.clone())
124                        },
125                        waiting_for: "incomplete tool arguments at stream end".into(),
126                    },
127                    source_time_ms: None,
128                    source_step: None,
129                });
130            }
131        }
132        Ok(frags)
133    }
134
135    fn handle_line(
136        &mut self,
137        line: &[u8],
138        out: &mut Vec<AcpFragment>,
139    ) -> Result<(), InterpreterError> {
140        if line.is_empty() {
141            // Dispatch event.
142            if self.event_data_lines.is_empty() {
143                return Ok(());
144            }
145            let data = self.event_data_lines.join("\n");
146            self.event_data_lines.clear();
147            if data.len() > self.max_event_bytes {
148                return Err(InterpreterError::new(
149                    InterpreterErrorKind::FrameLimitExceeded,
150                    "SSE event exceeds bound",
151                ));
152            }
153            if data.trim() == "[DONE]" {
154                self.saw_done = true;
155                return Ok(());
156            }
157            self.map_data_json(&data, out)?;
158            return Ok(());
159        }
160
161        // Comment / id / event / retry ignored; only data: matters for Chat Completions.
162        let line_str = std::str::from_utf8(line).map_err(|_| {
163            InterpreterError::new(
164                InterpreterErrorKind::MalformedFrame,
165                "SSE line is not valid UTF-8",
166            )
167        })?;
168        if let Some(rest) = line_str.strip_prefix("data:") {
169            let payload = rest.strip_prefix(' ').unwrap_or(rest);
170            self.event_data_lines.push(payload.to_string());
171        }
172        Ok(())
173    }
174
175    fn map_data_json(
176        &mut self,
177        data: &str,
178        out: &mut Vec<AcpFragment>,
179    ) -> Result<(), InterpreterError> {
180        let value: Value = serde_json::from_str(data).map_err(|_| {
181            InterpreterError::new(
182                InterpreterErrorKind::MalformedSemanticPayload,
183                "SSE data is not valid JSON",
184            )
185        })?;
186        let Some(choices) = value.get("choices").and_then(|c| c.as_array()) else {
187            return Ok(());
188        };
189        for choice in choices {
190            let index = choice.get("index").and_then(|i| i.as_u64()).unwrap_or(0) as u32;
191            if index != self.choice_index {
192                // Unsupported extra choices are ignored (not executed).
193                continue;
194            }
195            if let Some(delta) = choice.get("delta") {
196                if let Some(content) = delta.get("content").and_then(|c| c.as_str()) {
197                    if !content.is_empty() {
198                        out.push(AcpFragment::TextDelta {
199                            channel: TextChannel::PublicResponse,
200                            text: content.to_string(),
201                            source_time_ms: None,
202                            source_step: None,
203                        });
204                    }
205                }
206                if let Some(calls) = delta.get("tool_calls").and_then(|t| t.as_array()) {
207                    for call in calls {
208                        self.ingest_tool_delta(call, out)?;
209                    }
210                }
211            }
212            if let Some(reason) = choice.get("finish_reason").and_then(|r| r.as_str()) {
213                self.on_finish_reason(reason, out)?;
214            }
215        }
216        Ok(())
217    }
218
219    fn ingest_tool_delta(
220        &mut self,
221        call: &Value,
222        out: &mut Vec<AcpFragment>,
223    ) -> Result<(), InterpreterError> {
224        let idx = call.get("index").and_then(|i| i.as_u64()).unwrap_or(0) as u32;
225        let entry = self.tools.entry(idx).or_default();
226        if let Some(id) = call.get("id").and_then(|v| v.as_str()) {
227            if !id.is_empty() {
228                entry.id = id.to_string();
229            }
230        }
231        if let Some(func) = call.get("function") {
232            if let Some(name) = func.get("name").and_then(|n| n.as_str()) {
233                if !name.is_empty() {
234                    entry.name.push_str(name);
235                }
236            }
237            if let Some(args) = func.get("arguments").and_then(|a| a.as_str()) {
238                if entry.arguments.len().saturating_add(args.len()) > self.max_tool_arg_bytes {
239                    return Err(InterpreterError::new(
240                        InterpreterErrorKind::ToolLimitExceeded,
241                        "tool arguments exceed bound",
242                    ));
243                }
244                entry.arguments.push_str(args);
245            }
246        }
247        // D-016: accumulate deltas only. Waiting once we have an id; Ready only
248        // on qualified finish_reason == "tool_calls" (never mid-stream).
249        if !entry.id.is_empty() && !entry.emitted_waiting && !entry.ready {
250            entry.emitted_waiting = true;
251            out.push(AcpFragment::Tool {
252                action_id: ToolActionId::new(entry.id.clone()),
253                signal: ToolSignal::Waiting {
254                    tool_name: if entry.name.is_empty() {
255                        None
256                    } else {
257                        Some(entry.name.clone())
258                    },
259                    waiting_for: "tool arguments".into(),
260                },
261                source_time_ms: None,
262                source_step: None,
263            });
264        }
265        Ok(())
266    }
267
268    fn on_finish_reason(
269        &mut self,
270        reason: &str,
271        out: &mut Vec<AcpFragment>,
272    ) -> Result<(), InterpreterError> {
273        match reason {
274            // D-016: only tool_calls finish may promote Ready; length/content_filter
275            // must not execute incomplete argument fragments.
276            "tool_calls" => {
277                let keys: Vec<u32> = self.tools.keys().copied().collect();
278                for k in keys {
279                    let Some(entry) = self.tools.get_mut(&k) else {
280                        continue;
281                    };
282                    if entry.ready {
283                        continue;
284                    }
285                    if entry.id.is_empty()
286                        || entry.name.is_empty()
287                        || !is_complete_json_value(&entry.arguments)
288                    {
289                        return Err(InterpreterError::new(
290                            InterpreterErrorKind::MalformedSemanticPayload,
291                            "incomplete tool call at tool_calls finish",
292                        ));
293                    }
294                    entry.ready = true;
295                    out.push(AcpFragment::Tool {
296                        action_id: ToolActionId::new(entry.id.clone()),
297                        signal: ToolSignal::RequestReady {
298                            tool_name: entry.name.clone(),
299                            arguments_json: entry.arguments.clone(),
300                        },
301                        source_time_ms: None,
302                        source_step: None,
303                    });
304                }
305                Ok(())
306            }
307            "stop" | "length" | "content_filter" | "null" => {
308                // Do not promote incomplete tools (D-016).
309                Ok(())
310            }
311            other => Err(InterpreterError::new(
312                InterpreterErrorKind::MalformedSemanticPayload,
313                format!("unsupported finish_reason: {other}"),
314            )),
315        }
316    }
317}
318
319/// True when `s` parses as a complete JSON value (object/array/primitive).
320fn is_complete_json_value(s: &str) -> bool {
321    let t = s.trim();
322    if t.is_empty() {
323        return false;
324    }
325    serde_json::from_str::<Value>(t).is_ok()
326}
327
328#[cfg(test)]
329mod tests {
330    use super::*;
331
332    fn st() -> OpenAiSseState {
333        OpenAiSseState::new(4096, 64 * 1024, 64 * 1024)
334    }
335
336    #[test]
337    fn fragmented_sse_text() {
338        let mut s = st();
339        let mut all = Vec::new();
340        for chunk in [
341            b"data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"Hel".as_slice(),
342            b"lo\"}}]}\n\n".as_slice(),
343            b"data: [DONE]\n\n".as_slice(),
344        ] {
345            all.extend(s.push_bytes(chunk).unwrap());
346        }
347        assert!(s.saw_done());
348        let text: String = all
349            .iter()
350            .filter_map(|f| match f {
351                AcpFragment::TextDelta { text, .. } => Some(text.as_str()),
352                _ => None,
353            })
354            .collect();
355        assert_eq!(text, "Hello");
356    }
357
358    #[test]
359    fn tool_args_fragmented_only_ready_when_complete() {
360        let mut s = st();
361        let c1 = br#"data: {"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"search","arguments":"{\"q\":"}}]}}]}"#;
362        let mut fr1 = s.push_bytes(c1).unwrap();
363        fr1.extend(s.push_bytes(b"\n\n").unwrap());
364        assert!(fr1.iter().any(|f| matches!(
365            f,
366            AcpFragment::Tool {
367                signal: ToolSignal::Waiting { .. },
368                ..
369            }
370        )));
371        assert!(!fr1.iter().any(|f| matches!(
372            f,
373            AcpFragment::Tool {
374                signal: ToolSignal::RequestReady { .. },
375                ..
376            }
377        )));
378
379        let c2 = br#"data: {"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"hi\"}"}}]},"finish_reason":"tool_calls"}]}"#;
380        let mut fr2 = s.push_bytes(c2).unwrap();
381        fr2.extend(s.push_bytes(b"\n\n").unwrap());
382        assert!(fr2.iter().any(|f| matches!(
383            f,
384            AcpFragment::Tool {
385                signal: ToolSignal::RequestReady {
386                    tool_name,
387                    arguments_json,
388                },
389                ..
390            } if tool_name == "search" && arguments_json.contains("hi")
391        )));
392        let _ = s.push_bytes(b"data: [DONE]\n\n").unwrap();
393        assert!(s.saw_done());
394    }
395
396    #[test]
397    fn invalid_json_args_never_ready() {
398        let mut s = st();
399        let c = br#"data: {"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"c1","function":{"name":"x","arguments":"{not-json"}}]},"finish_reason":"tool_calls"}]}"#;
400        let fr = s.push_bytes(c).unwrap();
401        assert!(!fr.iter().any(|f| matches!(
402            f,
403            AcpFragment::Tool {
404                signal: ToolSignal::RequestReady { .. },
405                ..
406            }
407        )));
408        let err = s
409            .push_bytes(b"\n\n")
410            .expect_err("invalid tool JSON must fail closed at finish");
411        assert_eq!(err.kind, InterpreterErrorKind::MalformedSemanticPayload);
412        assert!(
413            !s.saw_done(),
414            "malformed tool args must not reach a successful done state"
415        );
416    }
417
418    #[test]
419    fn missing_done_fails_seal() {
420        let mut s = st();
421        let _ = s
422            .push_bytes(br#"data: {"choices":[{"index":0,"delta":{"content":"x"}}]}"#)
423            .unwrap();
424        let _ = s.push_bytes(b"\n\n").unwrap();
425        assert!(s.seal_clean().is_err());
426    }
427
428    #[test]
429    fn other_choice_index_ignored() {
430        let mut s = st();
431        let c = br#"data: {"choices":[{"index":1,"delta":{"content":"nope"}}]}"#;
432        let fr = s.push_bytes(c).unwrap();
433        s.push_bytes(b"\n\n").unwrap();
434        assert!(fr.is_empty());
435    }
436}