1use crate::error::ProtocolError;
2use serde::{Deserialize, Serialize};
3use std::sync::atomic::{AtomicU32, Ordering};
4
5const PREFIX: &str = "\x1bP@kitty-cmd";
6const SUFFIX: &str = "\x1b\\";
7const MAX_CHUNK_SIZE: usize = 4096;
8
9static STREAM_ID_COUNTER: AtomicU32 = AtomicU32::new(1);
10
11#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct KittyMessage {
13 pub cmd: String,
14 pub version: Vec<u32>,
15 #[serde(skip_serializing_if = "Option::is_none")]
16 pub no_response: Option<bool>,
17 #[serde(skip_serializing_if = "Option::is_none")]
18 pub kitty_window_id: Option<String>,
19 #[serde(skip_serializing_if = "Option::is_none")]
20 pub payload: Option<serde_json::Value>,
21 #[serde(skip_serializing_if = "Option::is_none")]
22 pub async_id: Option<String>,
23 #[serde(skip_serializing_if = "Option::is_none")]
24 pub cancel_async: Option<bool>,
25 #[serde(skip_serializing_if = "Option::is_none")]
26 pub stream_id: Option<String>,
27 #[serde(skip_serializing_if = "Option::is_none")]
28 pub stream: Option<bool>,
29}
30
31impl KittyMessage {
32 pub fn new(cmd: impl Into<String>, version: impl Into<Vec<u32>>) -> Self {
33 Self {
34 cmd: cmd.into(),
35 version: version.into(),
36 no_response: None,
37 kitty_window_id: None,
38 payload: None,
39 async_id: None,
40 cancel_async: None,
41 stream_id: None,
42 stream: None,
43 }
44 }
45
46 pub fn no_response(mut self, value: bool) -> Self {
47 self.no_response = Some(value);
48 self
49 }
50
51 pub fn kitty_window_id(mut self, id: impl Into<String>) -> Self {
52 self.kitty_window_id = Some(id.into());
53 self
54 }
55
56 pub fn payload(mut self, payload: serde_json::Value) -> Self {
57 self.payload = Some(payload);
58 self
59 }
60
61 pub fn async_id(mut self, id: impl Into<String>) -> Self {
62 self.async_id = Some(id.into());
63 self
64 }
65
66 pub fn cancel_async(mut self, value: bool) -> Self {
67 self.cancel_async = Some(value);
68 self
69 }
70
71 pub fn stream_id(mut self, id: impl Into<String>) -> Self {
72 self.stream_id = Some(id.into());
73 self
74 }
75
76 pub fn stream(mut self, value: bool) -> Self {
77 self.stream = Some(value);
78 self
79 }
80
81 pub fn generate_unique_id() -> String {
82 let id = STREAM_ID_COUNTER.fetch_add(1, Ordering::Relaxed);
83 format!("{:x}", id)
84 }
85
86 pub fn needs_streaming(&self) -> bool {
87 if let Some(payload) = &self.payload {
88 if let Some(obj) = payload.as_object() {
89 for (_key, value) in obj {
90 if let Some(s) = value.as_str() {
91 if s.len() > MAX_CHUNK_SIZE {
92 return true;
93 }
94 }
95 }
96 }
97 }
98 false
99 }
100
101 pub fn into_chunks(mut self) -> Vec<KittyMessage> {
102 let mut chunks = Vec::new();
103
104 if !self.needs_streaming() {
105 return vec![self];
106 }
107
108 if let Some(payload) = self.payload.take() {
109 if let Some(obj) = payload.as_object() {
110 let stream_id = Self::generate_unique_id();
111
112 for (_key, value) in obj {
113 if let Some(s) = value.as_str() {
114 if s.len() > MAX_CHUNK_SIZE {
115 for (i, chunk_data) in s.as_bytes().chunks(MAX_CHUNK_SIZE).enumerate() {
116 let mut chunk_msg = self.clone();
117 chunk_msg.stream_id = Some(stream_id.clone());
118 chunk_msg.stream = Some(true);
119
120 let mut chunk_payload = serde_json::Map::new();
121 chunk_payload.insert("data".to_string(), serde_json::Value::String(String::from_utf8_lossy(chunk_data).to_string()));
122 chunk_payload.insert("chunk_num".to_string(), serde_json::json!(i));
123 chunk_msg.payload = Some(serde_json::Value::Object(chunk_payload));
124
125 chunks.push(chunk_msg);
126 }
127
128 let mut end_chunk = self.clone();
129 end_chunk.stream_id = Some(stream_id);
130 end_chunk.stream = Some(true);
131 let mut end_payload = serde_json::Map::new();
132 end_payload.insert("data".to_string(), serde_json::Value::String(String::new()));
133 end_chunk.payload = Some(serde_json::Value::Object(end_payload));
134 chunks.push(end_chunk);
135
136 return chunks;
137 }
138 }
139 }
140 }
141 }
142
143 chunks.push(self);
144 chunks
145 }
146
147 pub fn encode(&self) -> Result<Vec<u8>, ProtocolError> {
148 let json = serde_json::to_string(self)?;
149 let message = format!("{}{}{}", PREFIX, json, SUFFIX);
150 Ok(message.into_bytes())
151 }
152
153 pub fn decode(data: &[u8]) -> Result<Self, ProtocolError> {
154 let s = std::str::from_utf8(data)
155 .map_err(|e| ProtocolError::InvalidMessageFormat(e.to_string()))?;
156
157 if !s.starts_with(PREFIX) {
158 return Err(ProtocolError::InvalidEscapeSequence);
159 }
160
161 if !s.ends_with(SUFFIX) {
162 return Err(ProtocolError::InvalidEscapeSequence);
163 }
164
165 let json_start = PREFIX.len();
166 let json_end = s.len() - SUFFIX.len();
167 let json_str = &s[json_start..json_end];
168
169 serde_json::from_str(json_str).map_err(ProtocolError::JsonError)
170 }
171}
172
173#[derive(Debug, Clone, Serialize, Deserialize)]
174pub struct KittyResponse {
175 pub ok: bool,
176 pub data: Option<serde_json::Value>,
177 pub error: Option<String>,
178}
179
180impl KittyResponse {
181 pub fn decode(data: &[u8]) -> Result<Self, ProtocolError> {
182 let s = std::str::from_utf8(data)
183 .map_err(|e| ProtocolError::EnvelopeParseError(e.to_string()))?;
184
185 if !s.starts_with("\x1bP@kitty-cmd") {
186 return Err(ProtocolError::EnvelopeParseError(
187 "Invalid response prefix".to_string(),
188 ));
189 }
190
191 if !s.ends_with("\x1b\\") {
192 return Err(ProtocolError::EnvelopeParseError(
193 "Invalid response suffix".to_string(),
194 ));
195 }
196
197 let json_start = PREFIX.len();
198 let json_end = s.len() - SUFFIX.len();
199 let json_str = &s[json_start..json_end];
200
201 let msg: serde_json::Value = serde_json::from_str(json_str)
202 .map_err(ProtocolError::JsonError)?;
203
204 if !msg.is_object() {
205 return Err(ProtocolError::EnvelopeParseError(
206 "Response is not a JSON object".to_string(),
207 ));
208 }
209
210 serde_json::from_value(msg).map_err(ProtocolError::JsonError)
211 }
212}
213
214#[cfg(test)]
215mod tests {
216 use super::*;
217
218 #[test]
219 fn test_message_encode() {
220 let msg = KittyMessage::new("ls", vec![0, 14, 2]);
221 let encoded = msg.encode().unwrap();
222 let decoded = KittyMessage::decode(&encoded).unwrap();
223 assert_eq!(decoded.cmd, "ls");
224 assert_eq!(decoded.version, vec![0, 14, 2]);
225 }
226
227 #[test]
228 fn test_message_with_payload() {
229 let msg = KittyMessage::new("send-text", vec![0, 14, 2])
230 .payload(serde_json::json!({"match": "id:1", "data": "text:hello"}));
231 let encoded = msg.encode().unwrap();
232 let decoded = KittyMessage::decode(&encoded).unwrap();
233 assert_eq!(decoded.cmd, "send-text");
234 assert!(decoded.payload.is_some());
235 }
236
237 #[test]
238 fn test_message_no_response() {
239 let msg = KittyMessage::new("close-window", vec![0, 14, 2])
240 .no_response(true);
241 let encoded = msg.encode().unwrap();
242 let decoded = KittyMessage::decode(&encoded).unwrap();
243 assert_eq!(decoded.no_response, Some(true));
244 }
245
246 #[test]
247 fn test_invalid_escape_sequence() {
248 let data = b"invalid message";
249 let result = KittyMessage::decode(data);
250 assert!(result.is_err());
251 }
252
253 #[test]
254 fn test_response_decode() {
255 let raw = b"\x1bP@kitty-cmd{\"ok\":true,\"data\":[{\"id\":1,\"title\":\"test\"}]}\x1b\\";
256 let response = KittyResponse::decode(raw).unwrap();
257 assert!(response.ok);
258 assert!(response.data.is_some());
259 }
260
261 #[test]
262 fn test_async_id() {
263 let msg = KittyMessage::new("select-window", vec![0, 14, 2])
264 .async_id("abc123");
265 let encoded = msg.encode().unwrap();
266 let decoded = KittyMessage::decode(&encoded).unwrap();
267 assert_eq!(decoded.async_id, Some("abc123".to_string()));
268 }
269
270 #[test]
271 fn test_cancel_async() {
272 let msg = KittyMessage::new("select-window", vec![0, 14, 2])
273 .async_id("abc123")
274 .cancel_async(true);
275 let encoded = msg.encode().unwrap();
276 let decoded = KittyMessage::decode(&encoded).unwrap();
277 assert_eq!(decoded.cancel_async, Some(true));
278 }
279
280 #[test]
281 fn test_unique_id_generation() {
282 let id1 = KittyMessage::generate_unique_id();
283 let id2 = KittyMessage::generate_unique_id();
284 assert_ne!(id1, id2);
285 }
286
287 #[test]
288 fn test_needs_streaming_false() {
289 let msg = KittyMessage::new("send-text", vec![0, 14, 2])
290 .payload(serde_json::json!({"data": "hello"}));
291 assert!(!msg.needs_streaming());
292 }
293
294 #[test]
295 fn test_needs_streaming_true() {
296 let large_data = "x".repeat(5000);
297 let msg = KittyMessage::new("send-text", vec![0, 14, 2])
298 .payload(serde_json::json!({"data": large_data}));
299 assert!(msg.needs_streaming());
300 }
301
302 #[test]
303 fn test_into_chunks_no_streaming() {
304 let msg = KittyMessage::new("send-text", vec![0, 14, 2])
305 .payload(serde_json::json!({"data": "hello"}));
306 let chunks = msg.into_chunks();
307 assert_eq!(chunks.len(), 1);
308 }
309
310 #[test]
311 fn test_into_chunks_with_streaming() {
312 let large_data = "x".repeat(5000);
313 let msg = KittyMessage::new("set-background-image", vec![0, 14, 2])
314 .payload(serde_json::json!({"data": large_data}));
315 let chunks = msg.into_chunks();
316 assert!(chunks.len() > 1);
317 assert!(chunks.iter().all(|c| c.stream_id.is_some()));
318 assert!(chunks.iter().all(|c| c.stream == Some(true)));
319 }
320}