lago_api/sse/
vercel.rs

1use lago_core::EventEnvelope;
2use lago_core::event::EventPayload;
3use serde_json::json;
4
5use super::format::{SseFormat, SseFrame};
6
7/// Vercel AI SDK compatible SSE format.
8///
9/// Formats events using the Vercel AI SDK UI message stream protocol with
10/// lifecycle frames (`start-step`, `text-start`, `text-delta`, `text-end`,
11/// `finish-step`) and tool streaming. Adds the
12/// `x-vercel-ai-ui-message-stream: v1` header.
13pub struct VercelFormat;
14
15impl SseFormat for VercelFormat {
16    fn format(&self, event: &EventEnvelope) -> Vec<SseFrame> {
17        let id = Some(event.seq.to_string());
18
19        match &event.payload {
20            EventPayload::Message { content, .. } => {
21                // Full message: emit lifecycle frames
22                vec![
23                    make_frame("start-step", json!({}), &id),
24                    make_frame("text-start", json!({}), &id),
25                    make_frame(
26                        "text-delta",
27                        json!({
28                            "id": event.event_id.to_string(),
29                            "delta": content,
30                        }),
31                        &id,
32                    ),
33                    make_frame("text-end", json!({}), &id),
34                    make_frame("finish-step", json!({}), &id),
35                ]
36            }
37
38            EventPayload::MessageDelta { delta, .. } => {
39                vec![make_frame(
40                    "text-delta",
41                    json!({
42                        "id": event.event_id.to_string(),
43                        "delta": delta,
44                    }),
45                    &id,
46                )]
47            }
48
49            EventPayload::ToolInvoke {
50                call_id,
51                tool_name,
52                arguments,
53                ..
54            } => {
55                let args_str = arguments.to_string();
56                vec![
57                    make_frame(
58                        "tool-input-start",
59                        json!({
60                            "toolCallId": call_id,
61                            "toolName": tool_name,
62                        }),
63                        &id,
64                    ),
65                    make_frame(
66                        "tool-input-delta",
67                        json!({
68                            "toolCallId": call_id,
69                            "delta": args_str,
70                        }),
71                        &id,
72                    ),
73                    make_frame(
74                        "tool-input-available",
75                        json!({
76                            "toolCallId": call_id,
77                            "toolName": tool_name,
78                            "input": arguments,
79                        }),
80                        &id,
81                    ),
82                ]
83            }
84
85            EventPayload::ToolResult {
86                call_id,
87                tool_name,
88                result,
89                ..
90            } => {
91                vec![make_frame(
92                    "tool-output-available",
93                    json!({
94                        "toolCallId": call_id,
95                        "toolName": tool_name,
96                        "output": result,
97                    }),
98                    &id,
99                )]
100            }
101
102            // Non-message/tool events are filtered out in Vercel format
103            _ => Vec::new(),
104        }
105    }
106
107    fn done_frame(&self) -> Option<SseFrame> {
108        let done = json!({
109            "type": "finish",
110            "finishReason": "stop",
111        });
112        Some(SseFrame {
113            event: None,
114            data: done.to_string(),
115            id: None,
116        })
117    }
118
119    fn extra_headers(&self) -> Vec<(String, String)> {
120        vec![(
121            "x-vercel-ai-ui-message-stream".to_string(),
122            "v1".to_string(),
123        )]
124    }
125
126    fn name(&self) -> &str {
127        "vercel"
128    }
129}
130
131/// Helper to create a typed SSE frame.
132fn make_frame(frame_type: &str, mut data: serde_json::Value, id: &Option<String>) -> SseFrame {
133    data["type"] = json!(frame_type);
134    SseFrame {
135        event: None,
136        data: data.to_string(),
137        id: id.clone(),
138    }
139}
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144    use lago_core::event::SpanStatus;
145    use lago_core::id::*;
146    use std::collections::HashMap;
147
148    fn make_envelope(payload: EventPayload, seq: u64) -> EventEnvelope {
149        EventEnvelope {
150            event_id: EventId::from_string("EVT001"),
151            session_id: SessionId::from_string("SESS001"),
152            branch_id: BranchId::from_string("main"),
153            run_id: None,
154            seq,
155            timestamp: 1_700_000_000_000_000,
156            parent_id: None,
157            payload,
158            metadata: HashMap::new(),
159        }
160    }
161
162    fn parse_frame(frame: &SseFrame) -> serde_json::Value {
163        serde_json::from_str(&frame.data).unwrap()
164    }
165
166    #[test]
167    fn message_produces_lifecycle_frames() {
168        let fmt = VercelFormat;
169        let event = make_envelope(
170            EventPayload::Message {
171                role: "assistant".into(),
172                content: "Hello!".into(),
173                model: None,
174                token_usage: None,
175            },
176            3,
177        );
178        let frames = fmt.format(&event);
179        assert_eq!(frames.len(), 5);
180
181        assert_eq!(parse_frame(&frames[0])["type"], "start-step");
182        assert_eq!(parse_frame(&frames[1])["type"], "text-start");
183        assert_eq!(parse_frame(&frames[2])["type"], "text-delta");
184        assert_eq!(parse_frame(&frames[2])["delta"], "Hello!");
185        assert_eq!(parse_frame(&frames[3])["type"], "text-end");
186        assert_eq!(parse_frame(&frames[4])["type"], "finish-step");
187
188        // All frames share the same id
189        for frame in &frames {
190            assert_eq!(frame.id.as_deref(), Some("3"));
191        }
192    }
193
194    #[test]
195    fn message_delta_produces_text_delta_frame() {
196        let fmt = VercelFormat;
197        let event = make_envelope(
198            EventPayload::MessageDelta {
199                role: "assistant".into(),
200                delta: "chunk".into(),
201                index: 0,
202            },
203            7,
204        );
205        let frames = fmt.format(&event);
206        assert_eq!(frames.len(), 1);
207        let data = parse_frame(&frames[0]);
208        assert_eq!(data["type"], "text-delta");
209        assert_eq!(data["delta"], "chunk");
210    }
211
212    #[test]
213    fn tool_invoke_produces_tool_input_frames() {
214        let fmt = VercelFormat;
215        let event = make_envelope(
216            EventPayload::ToolInvoke {
217                call_id: "call-1".into(),
218                tool_name: "read_file".into(),
219                arguments: serde_json::json!({"path": "/etc/hosts"}),
220                category: None,
221            },
222            10,
223        );
224        let frames = fmt.format(&event);
225        assert_eq!(frames.len(), 3);
226
227        let f0 = parse_frame(&frames[0]);
228        assert_eq!(f0["type"], "tool-input-start");
229        assert_eq!(f0["toolCallId"], "call-1");
230        assert_eq!(f0["toolName"], "read_file");
231
232        let f1 = parse_frame(&frames[1]);
233        assert_eq!(f1["type"], "tool-input-delta");
234        assert_eq!(f1["toolCallId"], "call-1");
235        assert!(f1["delta"].as_str().unwrap().contains("/etc/hosts"));
236
237        let f2 = parse_frame(&frames[2]);
238        assert_eq!(f2["type"], "tool-input-available");
239        assert_eq!(f2["toolCallId"], "call-1");
240        assert_eq!(f2["input"]["path"], "/etc/hosts");
241    }
242
243    #[test]
244    fn tool_result_produces_tool_output_frame() {
245        let fmt = VercelFormat;
246        let event = make_envelope(
247            EventPayload::ToolResult {
248                call_id: "call-1".into(),
249                tool_name: "read_file".into(),
250                result: serde_json::json!({"content": "data"}),
251                duration_ms: 42,
252                status: SpanStatus::Ok,
253            },
254            11,
255        );
256        let frames = fmt.format(&event);
257        assert_eq!(frames.len(), 1);
258
259        let data = parse_frame(&frames[0]);
260        assert_eq!(data["type"], "tool-output-available");
261        assert_eq!(data["toolCallId"], "call-1");
262        assert_eq!(data["output"]["content"], "data");
263    }
264
265    #[test]
266    fn non_message_events_filtered() {
267        let fmt = VercelFormat;
268        let event = make_envelope(
269            EventPayload::FileWrite {
270                path: "/a".into(),
271                blob_hash: BlobHash::from_hex("abc"),
272                size_bytes: 10,
273                content_type: None,
274            },
275            1,
276        );
277        assert!(fmt.format(&event).is_empty());
278    }
279
280    #[test]
281    fn done_frame_is_finish() {
282        let fmt = VercelFormat;
283        let done = fmt.done_frame().unwrap();
284        let data = parse_frame(&done);
285        assert_eq!(data["type"], "finish");
286        assert_eq!(data["finishReason"], "stop");
287    }
288
289    #[test]
290    fn extra_headers_include_ui_message_stream_header() {
291        let fmt = VercelFormat;
292        let headers = fmt.extra_headers();
293        assert_eq!(headers.len(), 1);
294        assert_eq!(headers[0].0, "x-vercel-ai-ui-message-stream");
295        assert_eq!(headers[0].1, "v1");
296    }
297
298    #[test]
299    fn name_is_vercel() {
300        assert_eq!(VercelFormat.name(), "vercel");
301    }
302
303    #[test]
304    fn sandbox_events_filtered() {
305        let fmt = VercelFormat;
306        let event = make_envelope(
307            EventPayload::SandboxCreated {
308                sandbox_id: "sbx-001".into(),
309                tier: lago_core::sandbox::SandboxTier::Container,
310                config: lago_core::sandbox::SandboxConfig {
311                    tier: lago_core::sandbox::SandboxTier::Container,
312                    allowed_paths: vec![],
313                    allowed_commands: vec![],
314                    network_access: false,
315                    max_memory_mb: None,
316                    max_cpu_seconds: None,
317                },
318            },
319            20,
320        );
321        assert!(fmt.format(&event).is_empty());
322    }
323}
lago_api/sse/vercel.rs

lago_api/sse/
vercel.rs
