batty_cli/shim/
kiro_types.rs

1//! JSON-RPC 2.0 message types for the Kiro CLI Agent Client Protocol (ACP).
2//!
3//! These types model the messages exchanged over stdin/stdout when Kiro CLI
4//! runs in `kiro-cli acp --trust-all-tools` mode. Each line is a complete
5//! JSON-RPC 2.0 message (NDJSON transport).
6//!
7//! Protocol reference: <https://agentclientprotocol.com/protocol/overview>
8
9use serde::{Deserialize, Serialize};
10use serde_json::Value;
11
12// ---------------------------------------------------------------------------
13// JSON-RPC 2.0 request/response wrappers
14// ---------------------------------------------------------------------------
15
16/// A JSON-RPC 2.0 request (sent TO Kiro on stdin).
17#[derive(Debug, Serialize)]
18pub struct JsonRpcRequest {
19    pub jsonrpc: &'static str,
20    pub id: u64,
21    pub method: String,
22    #[serde(skip_serializing_if = "Option::is_none")]
23    pub params: Option<Value>,
24}
25
26impl JsonRpcRequest {
27    pub fn new(id: u64, method: &str, params: Option<Value>) -> Self {
28        Self {
29            jsonrpc: "2.0",
30            id,
31            method: method.to_string(),
32            params,
33        }
34    }
35
36    /// Serialize to a single NDJSON line (no trailing newline).
37    pub fn to_ndjson(&self) -> String {
38        serde_json::to_string(self).expect("JsonRpcRequest is always serializable")
39    }
40}
41
42/// A JSON-RPC 2.0 response (sent TO Kiro on stdin, e.g. permission reply).
43#[derive(Debug, Serialize)]
44pub struct JsonRpcResponse {
45    pub jsonrpc: &'static str,
46    pub id: u64,
47    pub result: Value,
48}
49
50impl JsonRpcResponse {
51    pub fn new(id: u64, result: Value) -> Self {
52        Self {
53            jsonrpc: "2.0",
54            id,
55            result,
56        }
57    }
58
59    pub fn to_ndjson(&self) -> String {
60        serde_json::to_string(self).expect("JsonRpcResponse is always serializable")
61    }
62}
63
64/// A single NDJSON message received from Kiro's stdout.
65///
66/// Can be a response (has `id` + `result`/`error`), a notification (has
67/// `method` + `params` but no `id`), or a request from the agent (has
68/// `id` + `method` + `params`, e.g. `session/request_permission`).
69#[derive(Debug, Deserialize)]
70pub struct AcpMessage {
71    #[serde(default)]
72    pub jsonrpc: Option<String>,
73
74    /// Present on responses and agent-initiated requests.
75    #[serde(default)]
76    pub id: Option<u64>,
77
78    /// Present on notifications and agent-initiated requests.
79    #[serde(default)]
80    pub method: Option<String>,
81
82    /// Present on notifications and agent-initiated requests.
83    #[serde(default)]
84    pub params: Option<Value>,
85
86    /// Present on successful responses.
87    #[serde(default)]
88    pub result: Option<Value>,
89
90    /// Present on error responses.
91    #[serde(default)]
92    pub error: Option<Value>,
93}
94
95impl AcpMessage {
96    /// Is this a response to a request we sent (has id + result/error, no method)?
97    pub fn is_response(&self) -> bool {
98        self.id.is_some() && self.method.is_none()
99    }
100
101    /// Is this a notification from the agent (has method, no id)?
102    pub fn is_notification(&self) -> bool {
103        self.method.is_some() && self.id.is_none()
104    }
105
106    /// Is this a request from the agent to us (has id + method)?
107    pub fn is_agent_request(&self) -> bool {
108        self.id.is_some() && self.method.is_some()
109    }
110}
111
112// ---------------------------------------------------------------------------
113// ACP-specific request builders
114// ---------------------------------------------------------------------------
115
116/// Build the `initialize` request.
117pub fn initialize_request(id: u64) -> JsonRpcRequest {
118    let params = serde_json::json!({
119        "protocolVersion": 1,
120        "clientCapabilities": {
121            "fs": { "readTextFile": false, "writeTextFile": false },
122            "terminal": false
123        },
124        "clientInfo": {
125            "name": "batty",
126            "version": env!("CARGO_PKG_VERSION")
127        }
128    });
129    JsonRpcRequest::new(id, "initialize", Some(params))
130}
131
132/// Build the `session/new` request.
133pub fn session_new_request(id: u64, cwd: &str) -> JsonRpcRequest {
134    let params = serde_json::json!({
135        "cwd": cwd,
136        "mcpServers": []
137    });
138    JsonRpcRequest::new(id, "session/new", Some(params))
139}
140
141/// Build the `session/load` request (resume existing session).
142pub fn session_load_request(id: u64, session_id: &str) -> JsonRpcRequest {
143    let params = serde_json::json!({
144        "sessionId": session_id
145    });
146    JsonRpcRequest::new(id, "session/load", Some(params))
147}
148
149/// Build a `session/prompt` request.
150pub fn session_prompt_request(id: u64, session_id: &str, text: &str) -> JsonRpcRequest {
151    let params = serde_json::json!({
152        "sessionId": session_id,
153        "prompt": [{ "type": "text", "text": text }]
154    });
155    JsonRpcRequest::new(id, "session/prompt", Some(params))
156}
157
158/// Build a `session/cancel` request.
159pub fn session_cancel_request(id: u64, session_id: &str) -> JsonRpcRequest {
160    let params = serde_json::json!({
161        "sessionId": session_id
162    });
163    JsonRpcRequest::new(id, "session/cancel", Some(params))
164}
165
166/// Build a permission approval response for `session/request_permission`.
167pub fn permission_approve_response(request_id: u64) -> JsonRpcResponse {
168    JsonRpcResponse::new(
169        request_id,
170        serde_json::json!({
171            "outcome": {
172                "outcome": "selected",
173                "optionId": "allow_once"
174            }
175        }),
176    )
177}
178
179// ---------------------------------------------------------------------------
180// Session update extraction helpers
181// ---------------------------------------------------------------------------
182
183/// Extract the `sessionUpdate` discriminator from a `session/update` params.
184pub fn extract_update_type(params: &Value) -> Option<&str> {
185    params
186        .get("update")
187        .and_then(|u| u.get("sessionUpdate"))
188        .and_then(|v| v.as_str())
189}
190
191/// Extract text content from an `agent_message_chunk` update.
192pub fn extract_message_chunk_text(params: &Value) -> Option<&str> {
193    params
194        .get("update")
195        .and_then(|u| u.get("content"))
196        .and_then(|c| c.get("text"))
197        .and_then(|t| t.as_str())
198}
199
200/// Extract the session ID from a `session/new` or `session/load` response result.
201pub fn extract_session_id(result: &Value) -> Option<&str> {
202    result.get("sessionId").and_then(|v| v.as_str())
203}
204
205/// Extract the stop reason from a prompt result.
206pub fn extract_stop_reason(result: &Value) -> Option<&str> {
207    result.get("stopReason").and_then(|v| v.as_str())
208}
209
210/// Extract context usage percentage from `_kiro.dev/metadata` params.
211pub fn extract_context_usage(params: &Value) -> Option<f64> {
212    params
213        .get("contextUsagePercentage")
214        .and_then(|v| v.as_f64())
215}
216
217/// Build the kiro-cli ACP launch command.
218///
219/// `system_prompt`: optional system prompt passed via --agent config.
220/// Returns the shell command to launch kiro-cli in ACP mode.
221pub fn kiro_acp_command(program: &str) -> String {
222    format!("exec {program} acp --trust-all-tools")
223}
224
225// ---------------------------------------------------------------------------
226// Tests
227// ---------------------------------------------------------------------------
228
229#[cfg(test)]
230mod tests {
231    use super::*;
232
233    // --- Request serialization ---
234
235    #[test]
236    fn initialize_request_serializes() {
237        let req = initialize_request(0);
238        let json: Value = serde_json::from_str(&req.to_ndjson()).unwrap();
239        assert_eq!(json["jsonrpc"], "2.0");
240        assert_eq!(json["id"], 0);
241        assert_eq!(json["method"], "initialize");
242        assert_eq!(json["params"]["protocolVersion"], 1);
243        assert!(json["params"]["clientInfo"]["name"].as_str().is_some());
244    }
245
246    #[test]
247    fn session_new_request_serializes() {
248        let req = session_new_request(1, "/home/user/project");
249        let json: Value = serde_json::from_str(&req.to_ndjson()).unwrap();
250        assert_eq!(json["method"], "session/new");
251        assert_eq!(json["params"]["cwd"], "/home/user/project");
252    }
253
254    #[test]
255    fn session_load_request_serializes() {
256        let req = session_load_request(1, "sess-abc");
257        let json: Value = serde_json::from_str(&req.to_ndjson()).unwrap();
258        assert_eq!(json["method"], "session/load");
259        assert_eq!(json["params"]["sessionId"], "sess-abc");
260    }
261
262    #[test]
263    fn session_prompt_request_serializes() {
264        let req = session_prompt_request(2, "sess-abc", "Fix the bug");
265        let json: Value = serde_json::from_str(&req.to_ndjson()).unwrap();
266        assert_eq!(json["method"], "session/prompt");
267        assert_eq!(json["params"]["sessionId"], "sess-abc");
268        assert_eq!(json["params"]["prompt"][0]["type"], "text");
269        assert_eq!(json["params"]["prompt"][0]["text"], "Fix the bug");
270    }
271
272    #[test]
273    fn session_cancel_request_serializes() {
274        let req = session_cancel_request(3, "sess-abc");
275        let json: Value = serde_json::from_str(&req.to_ndjson()).unwrap();
276        assert_eq!(json["method"], "session/cancel");
277        assert_eq!(json["params"]["sessionId"], "sess-abc");
278    }
279
280    #[test]
281    fn permission_approve_response_serializes() {
282        let resp = permission_approve_response(5);
283        let json: Value = serde_json::from_str(&resp.to_ndjson()).unwrap();
284        assert_eq!(json["jsonrpc"], "2.0");
285        assert_eq!(json["id"], 5);
286        assert_eq!(json["result"]["outcome"]["outcome"], "selected");
287        assert_eq!(json["result"]["outcome"]["optionId"], "allow_once");
288    }
289
290    // --- AcpMessage deserialization ---
291
292    #[test]
293    fn parse_response_message() {
294        let line =
295            r#"{"jsonrpc":"2.0","id":0,"result":{"protocolVersion":1,"agentCapabilities":{}}}"#;
296        let msg: AcpMessage = serde_json::from_str(line).unwrap();
297        assert!(msg.is_response());
298        assert!(!msg.is_notification());
299        assert!(!msg.is_agent_request());
300        assert_eq!(msg.id, Some(0));
301        assert!(msg.result.is_some());
302    }
303
304    #[test]
305    fn parse_notification_message() {
306        let line = r#"{"jsonrpc":"2.0","method":"session/update","params":{"sessionId":"s1","update":{"sessionUpdate":"agent_message_chunk","content":{"type":"text","text":"hello"}}}}"#;
307        let msg: AcpMessage = serde_json::from_str(line).unwrap();
308        assert!(msg.is_notification());
309        assert!(!msg.is_response());
310        assert_eq!(msg.method.as_deref(), Some("session/update"));
311    }
312
313    #[test]
314    fn parse_agent_request_message() {
315        let line = r#"{"jsonrpc":"2.0","id":5,"method":"session/request_permission","params":{"sessionId":"s1","toolCall":{"toolCallId":"c1","title":"Running: ls","kind":"execute"},"options":[{"optionId":"allow_once","name":"Yes"},{"optionId":"deny","name":"No"}]}}"#;
316        let msg: AcpMessage = serde_json::from_str(line).unwrap();
317        assert!(msg.is_agent_request());
318        assert_eq!(msg.method.as_deref(), Some("session/request_permission"));
319        assert_eq!(msg.id, Some(5));
320    }
321
322    #[test]
323    fn parse_error_response() {
324        let line =
325            r#"{"jsonrpc":"2.0","id":1,"error":{"code":-32600,"message":"invalid request"}}"#;
326        let msg: AcpMessage = serde_json::from_str(line).unwrap();
327        assert!(msg.is_response());
328        assert!(msg.error.is_some());
329        assert!(msg.result.is_none());
330    }
331
332    #[test]
333    fn unknown_fields_tolerated() {
334        let line = r#"{"jsonrpc":"2.0","method":"_kiro.dev/metadata","params":{"credits":8.5,"contextUsagePercentage":62.3,"future_field":true}}"#;
335        let msg: AcpMessage = serde_json::from_str(line).unwrap();
336        assert!(msg.is_notification());
337    }
338
339    // --- Extraction helpers ---
340
341    #[test]
342    fn extract_update_type_agent_message_chunk() {
343        let params = serde_json::json!({
344            "sessionId": "s1",
345            "update": {
346                "sessionUpdate": "agent_message_chunk",
347                "content": { "type": "text", "text": "hello" }
348            }
349        });
350        assert_eq!(extract_update_type(&params), Some("agent_message_chunk"));
351    }
352
353    #[test]
354    fn extract_update_type_tool_call() {
355        let params = serde_json::json!({
356            "sessionId": "s1",
357            "update": {
358                "sessionUpdate": "tool_call",
359                "toolCallId": "c1",
360                "title": "Reading file",
361                "kind": "read",
362                "status": "pending"
363            }
364        });
365        assert_eq!(extract_update_type(&params), Some("tool_call"));
366    }
367
368    #[test]
369    fn extract_update_type_turn_end() {
370        let params = serde_json::json!({
371            "sessionId": "s1",
372            "update": { "sessionUpdate": "TurnEnd" }
373        });
374        assert_eq!(extract_update_type(&params), Some("TurnEnd"));
375    }
376
377    #[test]
378    fn extract_message_chunk_text_present() {
379        let params = serde_json::json!({
380            "update": {
381                "sessionUpdate": "agent_message_chunk",
382                "content": { "type": "text", "text": "Here is the fix" }
383            }
384        });
385        assert_eq!(extract_message_chunk_text(&params), Some("Here is the fix"));
386    }
387
388    #[test]
389    fn extract_message_chunk_text_missing() {
390        let params = serde_json::json!({
391            "update": {
392                "sessionUpdate": "tool_call",
393                "toolCallId": "c1"
394            }
395        });
396        assert_eq!(extract_message_chunk_text(&params), None);
397    }
398
399    #[test]
400    fn extract_session_id_from_result() {
401        let result = serde_json::json!({"sessionId": "sess-xyz-123"});
402        assert_eq!(extract_session_id(&result), Some("sess-xyz-123"));
403    }
404
405    #[test]
406    fn extract_session_id_missing() {
407        let result = serde_json::json!({"other": "field"});
408        assert_eq!(extract_session_id(&result), None);
409    }
410
411    #[test]
412    fn extract_context_usage_present() {
413        let params = serde_json::json!({
414            "credits": 8.5,
415            "contextUsagePercentage": 62.3
416        });
417        assert_eq!(extract_context_usage(&params), Some(62.3));
418    }
419
420    #[test]
421    fn extract_context_usage_missing() {
422        let params = serde_json::json!({"credits": 8.5});
423        assert_eq!(extract_context_usage(&params), None);
424    }
425
426    #[test]
427    fn kiro_acp_command_format() {
428        let cmd = kiro_acp_command("kiro-cli");
429        assert_eq!(cmd, "exec kiro-cli acp --trust-all-tools");
430    }
431
432    #[test]
433    fn kiro_acp_command_custom_binary() {
434        let cmd = kiro_acp_command("/opt/kiro-cli");
435        assert_eq!(cmd, "exec /opt/kiro-cli acp --trust-all-tools");
436    }
437
438    // --- JsonRpcRequest ---
439
440    #[test]
441    fn request_with_no_params() {
442        let req = JsonRpcRequest::new(99, "ping", None);
443        let json: Value = serde_json::from_str(&req.to_ndjson()).unwrap();
444        assert_eq!(json["method"], "ping");
445        assert!(json.get("params").is_none());
446    }
447
448    // --- JsonRpcResponse ---
449
450    #[test]
451    fn response_serializes() {
452        let resp = JsonRpcResponse::new(42, serde_json::json!({"ok": true}));
453        let json: Value = serde_json::from_str(&resp.to_ndjson()).unwrap();
454        assert_eq!(json["id"], 42);
455        assert_eq!(json["result"]["ok"], true);
456    }
457}
batty_cli/shim/kiro_types.rs

batty_cli/shim/
kiro_types.rs
