use std::collections::HashMap; use serde::Deserialize; /// Commands sent from the client to the agent (Client -> Agent) #[derive(Debug, Deserialize, PartialEq)] #[serde(tag = "type")] #[serde(rename_all = "snake_case")] pub enum ProtocolCommand { Message { msg_id: String, content: String, #[serde(default)] files: Vec, }, Stop, ToolApprove { call_id: String, #[serde(default)] scope: ApprovalScope, }, ToolDeny { call_id: String, #[serde(default)] reason: String, }, InitHistory { text: String, }, SetMode { mode: SessionMode, }, SetConfig { #[serde(default)] model: Option, #[serde(default)] thinking: Option, #[serde(default)] thinking_budget: Option, #[serde(default)] effort: Option, #[serde(default)] compaction: Option, }, AddMcpServer { name: String, transport: String, #[serde(default)] command: Option, #[serde(default)] args: Option>, #[serde(default)] env: Option>, #[serde(default)] url: Option, #[serde(default)] headers: Option>, }, Ping, } #[derive(Debug, Deserialize, Default, PartialEq, Eq)] #[serde(rename_all = "snake_case")] pub enum ApprovalScope { #[default] Once, Always, } #[derive(Debug, Deserialize, PartialEq, Eq)] #[serde(rename_all = "snake_case")] pub enum SessionMode { Default, AutoEdit, Yolo, } #[cfg(test)] mod tests { use super::*; #[test] fn set_config_debug_format() { let cmd = ProtocolCommand::SetConfig { model: Some("test-model".into()), thinking: None, thinking_budget: None, effort: None, compaction: None, }; let dbg = format!("{cmd:?}"); assert!(dbg.contains("SetConfig")); assert!(dbg.contains("test-model")); } #[test] fn set_config_equality() { let a = ProtocolCommand::SetConfig { model: Some("m".into()), thinking: None, thinking_budget: None, effort: None, compaction: None, }; let b = ProtocolCommand::SetConfig { model: Some("m".into()), thinking: None, thinking_budget: None, effort: None, compaction: None, }; assert_eq!(a, b); let c = ProtocolCommand::SetConfig { model: None, thinking: None, thinking_budget: None, effort: None, compaction: None, }; assert_ne!(a, c); } #[test] fn set_config_with_all_fields_equality() { let a = ProtocolCommand::SetConfig { model: Some("m".into()), thinking: Some("enabled".into()), thinking_budget: Some(8000), effort: Some("high".into()), compaction: None, }; let b = ProtocolCommand::SetConfig { model: Some("m".into()), thinking: Some("enabled".into()), thinking_budget: Some(8000), effort: Some("high".into()), compaction: None, }; assert_eq!(a, b); } #[test] fn set_config_all_none_fields() { let cmd = ProtocolCommand::SetConfig { model: None, thinking: None, thinking_budget: None, effort: None, compaction: None, }; let dbg = format!("{cmd:?}"); assert!(dbg.contains("SetConfig")); } #[test] fn set_config_with_compaction() { let json = r#"{"type":"set_config","compaction":"full"}"#; let cmd: ProtocolCommand = serde_json::from_str(json).unwrap(); match cmd { ProtocolCommand::SetConfig { compaction, .. } => { assert_eq!(compaction.unwrap(), "full"); } _ => panic!("expected SetConfig"), } } #[test] fn set_config_compaction_none_by_default() { let json = r#"{"type":"set_config","model":"test"}"#; let cmd: ProtocolCommand = serde_json::from_str(json).unwrap(); match cmd { ProtocolCommand::SetConfig { compaction, .. } => { assert!(compaction.is_none()); } _ => panic!("expected SetConfig"), } } #[test] fn add_mcp_server_stdio_deserialize() { let json = r#"{ "type": "add_mcp_server", "name": "team-tools", "transport": "stdio", "command": "node", "args": ["bridge.js", "--port", "9000"], "env": {"TOKEN": "abc123"} }"#; let cmd: ProtocolCommand = serde_json::from_str(json).unwrap(); match cmd { ProtocolCommand::AddMcpServer { name, transport, command, args, env, url, headers, } => { assert_eq!(name, "team-tools"); assert_eq!(transport, "stdio"); assert_eq!(command.unwrap(), "node"); assert_eq!(args.unwrap(), vec!["bridge.js", "--port", "9000"]); assert_eq!(env.unwrap().get("TOKEN").unwrap(), "abc123"); assert!(url.is_none()); assert!(headers.is_none()); } _ => panic!("expected AddMcpServer"), } } #[test] fn ping_deserialize() { let json = r#"{"type":"ping"}"#; let cmd: ProtocolCommand = serde_json::from_str(json).unwrap(); assert_eq!(cmd, ProtocolCommand::Ping); } #[test] fn add_mcp_server_sse_deserialize() { let json = r#"{ "type": "add_mcp_server", "name": "remote-tools", "transport": "sse", "url": "http://localhost:8080/sse", "headers": {"Authorization": "Bearer tok"} }"#; let cmd: ProtocolCommand = serde_json::from_str(json).unwrap(); match cmd { ProtocolCommand::AddMcpServer { name, transport, command, url, headers, .. } => { assert_eq!(name, "remote-tools"); assert_eq!(transport, "sse"); assert!(command.is_none()); assert_eq!(url.unwrap(), "http://localhost:8080/sse"); assert_eq!(headers.unwrap().get("Authorization").unwrap(), "Bearer tok"); } _ => panic!("expected AddMcpServer"), } } }