agent_commander/streaming/
input_stream.rs

1//! JSON Input Stream
2//! Creates NDJSON input for CLI tools
3
4use super::ndjson::stringify_ndjson_line;
5use serde_json::{json, Value};
6
7/// JSON Input Stream struct
8/// Builds NDJSON input for streaming to CLI tools
9#[derive(Debug, Clone)]
10pub struct JsonInputStream {
11    compact: bool,
12    messages: Vec<Value>,
13}
14
15impl Default for JsonInputStream {
16    fn default() -> Self {
17        Self::new(true)
18    }
19}
20
21impl JsonInputStream {
22    /// Create a new JSON input stream
23    ///
24    /// # Arguments
25    /// * `compact` - Use compact JSON (default: true)
26    pub fn new(compact: bool) -> Self {
27        Self {
28            compact,
29            messages: Vec::new(),
30        }
31    }
32
33    /// Add a message to the stream
34    ///
35    /// # Arguments
36    /// * `message` - Message to add
37    ///
38    /// # Returns
39    /// Self for chaining
40    pub fn add(&mut self, message: Value) -> &mut Self {
41        if !message.is_null() {
42            self.messages.push(message);
43        }
44        self
45    }
46
47    /// Add a user prompt message
48    ///
49    /// # Arguments
50    /// * `content` - Prompt content
51    ///
52    /// # Returns
53    /// Self for chaining
54    pub fn add_prompt(&mut self, content: &str) -> &mut Self {
55        self.add(json!({
56            "type": "user_prompt",
57            "content": content
58        }))
59    }
60
61    /// Add a system message
62    ///
63    /// # Arguments
64    /// * `content` - System message content
65    ///
66    /// # Returns
67    /// Self for chaining
68    pub fn add_system_message(&mut self, content: &str) -> &mut Self {
69        self.add(json!({
70            "type": "system",
71            "content": content
72        }))
73    }
74
75    /// Add a configuration message
76    ///
77    /// # Arguments
78    /// * `config` - Configuration object
79    ///
80    /// # Returns
81    /// Self for chaining
82    pub fn add_config(&mut self, config: Value) -> &mut Self {
83        let mut msg = json!({"type": "config"});
84        if let (Some(obj), Some(cfg)) = (msg.as_object_mut(), config.as_object()) {
85            for (k, v) in cfg {
86                obj.insert(k.clone(), v.clone());
87            }
88        }
89        self.add(msg)
90    }
91
92    /// Convert the stream to NDJSON string
93    pub fn to_string(&self) -> String {
94        self.messages
95            .iter()
96            .map(|msg| stringify_ndjson_line(msg, self.compact))
97            .collect()
98    }
99
100    /// Convert the stream to bytes
101    pub fn to_bytes(&self) -> Vec<u8> {
102        self.to_string().into_bytes()
103    }
104
105    /// Get the number of messages in the stream
106    pub fn size(&self) -> usize {
107        self.messages.len()
108    }
109
110    /// Clear all messages
111    pub fn clear(&mut self) -> &mut Self {
112        self.messages.clear();
113        self
114    }
115
116    /// Get all messages
117    pub fn get_messages(&self) -> &[Value] {
118        &self.messages
119    }
120
121    /// Create from a vector of messages
122    ///
123    /// # Arguments
124    /// * `messages` - Vector of messages
125    /// * `compact` - Use compact JSON
126    pub fn from_messages(messages: Vec<Value>, compact: bool) -> Self {
127        Self { compact, messages }
128    }
129}
130
131impl std::fmt::Display for JsonInputStream {
132    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
133        write!(f, "{}", self.to_string())
134    }
135}
136
137#[cfg(test)]
138mod tests {
139    use super::*;
140    use serde_json::json;
141
142    #[test]
143    fn test_add_message() {
144        let mut stream = JsonInputStream::new(true);
145        stream.add(json!({"type": "hello"}));
146        assert_eq!(stream.size(), 1);
147    }
148
149    #[test]
150    fn test_to_string_produces_ndjson() {
151        let mut stream = JsonInputStream::new(true);
152        stream.add(json!({"a": 1}));
153        stream.add(json!({"b": 2}));
154
155        let output = stream.to_string();
156        assert_eq!(output, "{\"a\":1}\n{\"b\":2}\n");
157    }
158
159    #[test]
160    fn test_add_prompt() {
161        let mut stream = JsonInputStream::new(true);
162        stream.add_prompt("Hello");
163
164        let messages = stream.get_messages();
165        assert_eq!(messages.len(), 1);
166        assert_eq!(messages[0]["type"], "user_prompt");
167        assert_eq!(messages[0]["content"], "Hello");
168    }
169
170    #[test]
171    fn test_add_system_message() {
172        let mut stream = JsonInputStream::new(true);
173        stream.add_system_message("You are helpful");
174
175        let messages = stream.get_messages();
176        assert_eq!(messages[0]["type"], "system");
177    }
178
179    #[test]
180    fn test_chaining() {
181        let mut stream = JsonInputStream::new(true);
182        stream
183            .add_system_message("System")
184            .add_prompt("User")
185            .add(json!({"custom": true}));
186
187        assert_eq!(stream.size(), 3);
188    }
189
190    #[test]
191    fn test_clear() {
192        let mut stream = JsonInputStream::new(true);
193        stream.add(json!({"a": 1}));
194        stream.clear();
195        assert_eq!(stream.size(), 0);
196    }
197
198    #[test]
199    fn test_from_messages() {
200        let messages = vec![json!({"a": 1}), json!({"b": 2})];
201        let stream = JsonInputStream::from_messages(messages, true);
202        assert_eq!(stream.size(), 2);
203    }
204
205    #[test]
206    fn test_to_bytes() {
207        let mut stream = JsonInputStream::new(true);
208        stream.add(json!({"test": true}));
209
210        let bytes = stream.to_bytes();
211        let str_result = String::from_utf8(bytes).unwrap();
212        assert_eq!(str_result, "{\"test\":true}\n");
213    }
214}
agent_commander/streaming/input_stream.rs

agent_commander/streaming/
input_stream.rs
