Skip to main content

lucy/
protocol.rs

1use std::io::{self, Write};
2
3use serde::Serialize;
4use serde_json::Value;
5
6#[derive(Debug, Serialize, Clone, PartialEq, Eq)]
7#[serde(tag = "type")]
8pub enum ProtocolEvent {
9    #[serde(rename = "session")]
10    Session { session_id: String, resumed: bool },
11    #[serde(rename = "assistant_delta")]
12    AssistantDelta { text: String },
13    #[serde(rename = "tool_call")]
14    ToolCall {
15        id: String,
16        name: String,
17        arguments: String,
18    },
19    #[serde(rename = "tool_result")]
20    ToolResult {
21        id: String,
22        name: String,
23        result: Value,
24    },
25    #[serde(rename = "turn_end")]
26    TurnEnd,
27    #[serde(rename = "turn_interrupted")]
28    TurnInterrupted { reason: String, phase: String },
29    #[serde(rename = "error")]
30    Error { message: String },
31}
32
33/// The normalized event boundary shared by the machine protocol and the TUI.
34pub trait EventSink {
35    fn emit_event(&mut self, event: &ProtocolEvent) -> io::Result<()>;
36}
37
38pub struct ProtocolWriter<W> {
39    writer: W,
40}
41
42impl<W: Write> ProtocolWriter<W> {
43    pub fn new(writer: W) -> Self {
44        Self { writer }
45    }
46
47    pub fn emit(&mut self, event: &ProtocolEvent) -> io::Result<()> {
48        self.emit_serializable(event)
49    }
50
51    pub fn emit_serializable<T: Serialize>(&mut self, record: &T) -> io::Result<()> {
52        serde_json::to_writer(&mut self.writer, record)
53            .map_err(|error| io::Error::other(format!("encode protocol event: {error}")))?;
54        self.writer.write_all(b"\n")?;
55        self.writer.flush()
56    }
57
58    pub fn session(&mut self, session_id: &str, resumed: bool) -> io::Result<()> {
59        self.emit(&ProtocolEvent::Session {
60            session_id: session_id.to_owned(),
61            resumed,
62        })
63    }
64
65    pub fn assistant_delta(&mut self, text: &str) -> io::Result<()> {
66        if text.is_empty() {
67            return Ok(());
68        }
69        self.emit(&ProtocolEvent::AssistantDelta {
70            text: text.to_owned(),
71        })
72    }
73
74    pub fn tool_call(&mut self, id: &str, name: &str, arguments: &str) -> io::Result<()> {
75        self.emit(&ProtocolEvent::ToolCall {
76            id: id.to_owned(),
77            name: name.to_owned(),
78            arguments: arguments.to_owned(),
79        })
80    }
81
82    pub fn tool_result(&mut self, id: &str, name: &str, result: Value) -> io::Result<()> {
83        self.emit(&ProtocolEvent::ToolResult {
84            id: id.to_owned(),
85            name: name.to_owned(),
86            result,
87        })
88    }
89
90    pub fn turn_end(&mut self) -> io::Result<()> {
91        self.emit(&ProtocolEvent::TurnEnd)
92    }
93
94    pub fn turn_interrupted(&mut self, reason: &str, phase: &str) -> io::Result<()> {
95        self.emit(&ProtocolEvent::TurnInterrupted {
96            reason: reason.to_owned(),
97            phase: phase.to_owned(),
98        })
99    }
100
101    pub fn error(&mut self, message: &str) -> io::Result<()> {
102        self.emit(&ProtocolEvent::Error {
103            message: message.to_owned(),
104        })
105    }
106}
107
108impl<W: Write> EventSink for ProtocolWriter<W> {
109    fn emit_event(&mut self, event: &ProtocolEvent) -> io::Result<()> {
110        self.emit(event)
111    }
112}
113
114#[cfg(test)]
115mod tests {
116    use super::*;
117
118    #[test]
119    fn protocol_writer_emits_only_single_line_json_records() {
120        let mut output = Vec::new();
121        {
122            let mut writer = ProtocolWriter::new(&mut output);
123            writer.assistant_delta("line one\nline two").expect("event");
124            writer
125                .tool_result(
126                    "call-1",
127                    "cmd",
128                    serde_json::json!({"stdout":"provider-shape-is-not-forwarded"}),
129                )
130                .expect("result");
131        }
132        let text = String::from_utf8(output).expect("UTF-8");
133        assert_eq!(text.lines().count(), 2);
134        for line in text.lines() {
135            serde_json::from_str::<Value>(line).expect("JSONL record");
136        }
137        assert!(!text.contains("choices"));
138    }
139
140    #[test]
141    fn interruption_event_is_a_normalized_json_record() {
142        let event = ProtocolEvent::TurnInterrupted {
143            reason: "user_cancelled".to_owned(),
144            phase: "provider_stream".to_owned(),
145        };
146        let value = serde_json::to_value(event).expect("event JSON");
147        assert_eq!(value["type"], "turn_interrupted");
148        assert_eq!(value["reason"], "user_cancelled");
149        assert_eq!(value["phase"], "provider_stream");
150        assert!(!serde_json::to_string(&value)
151            .expect("serialized event")
152            .contains("choices"));
153    }
154}