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
33pub 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}