1use crate::acp::{AcpFragment, ToolSignal};
9use monoloop_contracts::{InterpreterError, InterpreterErrorKind, TextChannel, ToolActionId};
10use serde_json::Value;
11use std::collections::HashMap;
12
13pub const DEFAULT_CHOICE_INDEX: u32 = 0;
15
16#[derive(Debug, Default)]
18pub struct OpenAiSseState {
19 line_carry: Vec<u8>,
21 event_data_lines: Vec<String>,
23 tools: HashMap<u32, PartialToolCall>,
25 saw_done: bool,
27 choice_index: u32,
29 max_line_bytes: usize,
31 max_event_bytes: usize,
33 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 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 pub fn saw_done(&self) -> bool {
60 self.saw_done
61 }
62
63 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 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 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 }
99 }
100 if !self.event_data_lines.is_empty() {
101 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 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 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 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 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 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 "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 Ok(())
310 }
311 other => Err(InterpreterError::new(
312 InterpreterErrorKind::MalformedSemanticPayload,
313 format!("unsupported finish_reason: {other}"),
314 )),
315 }
316 }
317}
318
319fn 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}