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::TextDelta { 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::ToolCallRequested {
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::ToolCallCompleted {
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            schema_version: 1,
160        }
161    }
162
163    fn parse_frame(frame: &SseFrame) -> serde_json::Value {
164        serde_json::from_str(&frame.data).unwrap()
165    }
166
167    #[test]
168    fn message_produces_lifecycle_frames() {
169        let fmt = VercelFormat;
170        let event = make_envelope(
171            EventPayload::Message {
172                role: "assistant".into(),
173                content: "Hello!".into(),
174                model: None,
175                token_usage: None,
176            },
177            3,
178        );
179        let frames = fmt.format(&event);
180        assert_eq!(frames.len(), 5);
181
182        assert_eq!(parse_frame(&frames[0])["type"], "start-step");
183        assert_eq!(parse_frame(&frames[1])["type"], "text-start");
184        assert_eq!(parse_frame(&frames[2])["type"], "text-delta");
185        assert_eq!(parse_frame(&frames[2])["delta"], "Hello!");
186        assert_eq!(parse_frame(&frames[3])["type"], "text-end");
187        assert_eq!(parse_frame(&frames[4])["type"], "finish-step");
188
189        // All frames share the same id
190        for frame in &frames {
191            assert_eq!(frame.id.as_deref(), Some("3"));
192        }
193    }
194
195    #[test]
196    fn message_delta_produces_text_delta_frame() {
197        let fmt = VercelFormat;
198        let event = make_envelope(
199            EventPayload::TextDelta {
200                delta: "chunk".into(),
201                index: Some(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::ToolCallRequested {
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::ToolCallCompleted {
248                tool_run_id: lago_core::protocol_bridge::aios_protocol::ToolRunId::default(),
249                call_id: Some("call-1".into()),
250                tool_name: "read_file".into(),
251                result: serde_json::json!({"content": "data"}),
252                duration_ms: 42,
253                status: SpanStatus::Ok,
254            },
255            11,
256        );
257        let frames = fmt.format(&event);
258        assert_eq!(frames.len(), 1);
259
260        let data = parse_frame(&frames[0]);
261        assert_eq!(data["type"], "tool-output-available");
262        assert_eq!(data["toolCallId"], "call-1");
263        assert_eq!(data["output"]["content"], "data");
264    }
265
266    #[test]
267    fn non_message_events_filtered() {
268        let fmt = VercelFormat;
269        let event = make_envelope(
270            EventPayload::FileWrite {
271                path: "/a".into(),
272                blob_hash: BlobHash::from_hex("abc").into(),
273                size_bytes: 10,
274                content_type: None,
275            },
276            1,
277        );
278        assert!(fmt.format(&event).is_empty());
279    }
280
281    #[test]
282    fn done_frame_is_finish() {
283        let fmt = VercelFormat;
284        let done = fmt.done_frame().unwrap();
285        let data = parse_frame(&done);
286        assert_eq!(data["type"], "finish");
287        assert_eq!(data["finishReason"], "stop");
288    }
289
290    #[test]
291    fn extra_headers_include_ui_message_stream_header() {
292        let fmt = VercelFormat;
293        let headers = fmt.extra_headers();
294        assert_eq!(headers.len(), 1);
295        assert_eq!(headers[0].0, "x-vercel-ai-ui-message-stream");
296        assert_eq!(headers[0].1, "v1");
297    }
298
299    #[test]
300    fn name_is_vercel() {
301        assert_eq!(VercelFormat.name(), "vercel");
302    }
303
304    #[test]
305    fn sandbox_events_filtered() {
306        let fmt = VercelFormat;
307        let event = make_envelope(
308            EventPayload::SandboxCreated {
309                sandbox_id: "sbx-001".into(),
310                tier: "container".into(),
311                config: serde_json::json!({
312                    "tier": "container",
313                    "allowed_paths": [],
314                    "allowed_commands": [],
315                    "network_access": false,
316                }),
317            },
318            20,
319        );
320        assert!(fmt.format(&event).is_empty());
321    }
322}
lago_api/sse/vercel.rs

lago_api/sse/
vercel.rs
