magic_coder_types/protocol/

types.rs

use serde::{Deserialize, Serialize};
use serde_json::Value;
use uuid::Uuid;

#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub enum Role {
    Assistant,
    User,
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub enum RuntimeMode {
    Agent,
    Plan,
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub enum ToolExecutionTarget {
    Unspecified,
    /// Tool must be executed by a connected client (TUI/VS Code).
    ClientLocal,
    /// Tool is executed server-side (agents/MCP/etc). Clients must not provide outputs.
    ServerAgents,
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub enum ToolCallApproval {
    Unspecified,
    Pending,
    Approved,
    AutoApproved,
    Rejected,
}

#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
pub struct ModelConfig {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub temperature: Option<f32>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub top_p: Option<f32>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub presence_penalty: Option<f32>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub frequency_penalty: Option<f32>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub max_tokens: Option<i32>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub reasoning_effort: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ThreadModelOverride {
    pub model_id: Uuid,
    pub model_config: ModelConfig,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCallOutput {
    /// Tool call id
    pub id: String,

    pub is_error: bool,
    pub output: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub duration_seconds: Option<i32>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum SubagentEscalationResolution {
    Rejected {
        #[serde(default, skip_serializing_if = "Option::is_none")]
        reason: Option<String>,
    },
    ResolvedWithOutput {
        #[serde(default)]
        is_error: bool,
        output: String,
        #[serde(default, skip_serializing_if = "Option::is_none")]
        duration_seconds: Option<i32>,
    },
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct LocalSkillMetadata {
    pub name: String,
    pub description: String,
    pub cwd: String,
}

trait BoolExt {
    fn is_false(&self) -> bool;
}

impl BoolExt for bool {
    fn is_false(&self) -> bool {
        !*self
    }
}

/// Normalized operating system family reported by a client during the handshake.
#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq)]
pub enum Os {
    /// The client could not determine a supported OS family.
    #[default]
    #[serde(rename = "other")]
    Other,
    /// Linux distributions.
    #[serde(rename = "linux")]
    Linux,
    /// macOS / Darwin.
    #[serde(rename = "macos")]
    MacOS,
    /// Microsoft Windows.
    #[serde(rename = "windows")]
    Windows,
}

/// Normalized CPU architecture reported by a client during the handshake.
#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq)]
pub enum Arch {
    /// The client could not determine a supported CPU architecture.
    #[default]
    #[serde(rename = "other")]
    Other,
    /// 32-bit x86.
    #[serde(rename = "x86")]
    X86,
    /// 64-bit x86.
    #[serde(rename = "amd64")]
    Amd64,
    /// 64-bit ARM.
    #[serde(rename = "aarch64")]
    Aarch64,
}

/// Snapshot of coarse client machine characteristics captured when the socket connects.
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
#[serde(default)]
pub struct ClientSystemInfo {
    /// Operating system family.
    pub os: Os,
    /// Human-readable OS version string reported by the client.
    pub os_version: String,
    /// CPU architecture family.
    pub arch: Arch,
    /// Logical CPU core count.
    pub cpu_cores: u16,
    /// Total physical memory reported by the client, in whole megabytes.
    pub ram_mb: u32,
}

impl ClientSystemInfo {
    fn is_unknown(&self) -> bool {
        self == &Self::default()
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ClientMessage {
    HelloMath {
        /// Random per-client-instance id (generated on client startup).
        ///
        /// Used for multi-client tool-call claiming/coordination. Not derived from auth token.
        client_instance_id: String,
        /// Client version.
        version: String,
        /// Minimum supported server version (the client can work with any server >= this).
        min_supported_version: String,
        /// Optional thread root message id used by reconnecting clients to restore active thread state.
        #[serde(default, skip_serializing_if = "Option::is_none")]
        resume_thread_id: Option<Uuid>,
        /// Whether the client is currently running in automagic mode.
        #[serde(default, skip_serializing_if = "BoolExt::is_false")]
        automagic: bool,
        /// Normalized snapshot of the client machine.
        ///
        /// Older clients may omit this; missing values deserialize as "unknown".
        #[serde(default, skip_serializing_if = "ClientSystemInfo::is_unknown")]
        system_info: ClientSystemInfo,
    },
    /// Send a user message.
    ///
    /// If `thread_id` is `None`, the server creates a new thread and returns it in `SendMessageAck`.
    /// `request_id` is a client-generated correlation id for 1:1 request↔response mapping.
    SendMessage {
        request_id: Uuid,
        thread_id: Option<Uuid>,
        text: String,
        /// Optional runtime mode update to apply to the thread before generation.
        #[serde(default, skip_serializing_if = "Option::is_none")]
        runtime_mode: Option<RuntimeMode>,
        /// Optional model override update to apply to the thread before generation.
        ///
        /// - `None`: do not change current override.
        /// - `Some(value)`: set override to `value`.
        #[serde(default, skip_serializing_if = "Option::is_none")]
        model_override: Option<ThreadModelOverride>,
    },
    /// Update/refresh auth token without reconnecting the WebSocket.
    UpdateAuthToken {
        token: String,
    },
    /// Update the client's workspace roots / default root (used for prompt context + validations).
    ///
    /// Sent separately from `HelloMath` so roots can be updated without reconnecting.
    UpdateWorkspaceRoots {
        default_root: String,
        workspace_roots: Vec<String>,
    },
    /// Replace the client-local skill snapshot exposed to the server for this connection.
    UpdateLocalSkills {
        skills: Vec<LocalSkillMetadata>,
    },
    RejectToolCall {
        id: String,
        #[serde(default, skip_serializing_if = "Option::is_none")]
        reason: Option<String>,
    },
    AcceptToolCall {
        id: String,
    },
    ResolveSubagentEscalation {
        parent_message_id: Uuid,
        subagent_run_id: Uuid,
        escalation_id: String,
        resolution: SubagentEscalationResolution,
    },
    ToolCallOutputs {
        outputs: Vec<ToolCallOutput>,
    },
    /// Cancel the current in-progress generation for a specific assistant message.
    CancelGeneration {
        message_id: Uuid,
    },
}

#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
pub struct Usage {
    pub input_tokens: i32,
    pub output_tokens: i32,
    #[serde(default)]
    pub cache_read_input_tokens: i32,
    #[serde(default)]
    pub cache_creation_input_tokens: i32,
    #[serde(default)]
    pub cache_creation_input_tokens_5m: i32,
    #[serde(default)]
    pub cache_creation_input_tokens_1h: i32,
}

#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub enum MessageStatus {
    Completed,
    /// The agent is waiting for the client/user to provide more input (e.g. tool outputs / approvals)
    /// before it can continue generation.
    WaitingForUser,
    Failed,
    Cancelled,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ServerMessage {
    HelloMagic {
        version: String,
        min_supported_version: String,
    },
    VersionMismatch {
        server_version: String,
        server_min_supported_version: String,
    },
    Goodbye {
        reconnect: bool,
    },
    SendMessageAck {
        request_id: Uuid,
        thread_id: Uuid,
        user_message_id: Uuid,
    },
    AuthUpdated,
    RuntimeModeUpdated {
        thread_id: Uuid,
        mode: RuntimeMode,
        #[serde(default, skip_serializing_if = "Option::is_none")]
        changed_by_client_instance_id: Option<String>,
    },
    ThreadModelUpdated {
        thread_id: Uuid,
        #[serde(default, skip_serializing_if = "Option::is_none")]
        model_override: Option<ThreadModelOverride>,
        #[serde(default, skip_serializing_if = "Option::is_none")]
        changed_by_client_instance_id: Option<String>,
    },
    MessageHeader {
        message_id: Uuid,
        thread_id: Uuid,
        role: Role,
        #[serde(default, skip_serializing_if = "Option::is_none")]
        request_id: Option<Uuid>,
    },
    ReasoningDelta {
        message_id: Uuid,
        content: String,
    },
    TextDelta {
        message_id: Uuid,
        content: String,
    },
    ToolCallHeader {
        message_id: Uuid,
        tool_call_id: String,
        name: String,
        execution_target: ToolExecutionTarget,
        approval: ToolCallApproval,
    },
    ToolCallArgumentsDelta {
        message_id: Uuid,
        tool_call_id: String,
        delta: String,
    },
    ToolCall {
        message_id: Uuid,
        tool_call_id: String,
        args: Value,
    },
    ToolCallResult {
        message_id: Uuid,
        tool_call_id: String,
        is_error: bool,
        output: String,
        #[serde(default, skip_serializing_if = "Option::is_none")]
        duration_seconds: Option<i32>,
    },
    ToolCallClaimed {
        message_id: Uuid,
        tool_call_id: String,
        claimed_by_client_instance_id: String,
    },
    ToolCallApprovalUpdated {
        message_id: Uuid,
        tool_call_id: String,
        approval: ToolCallApproval,
    },
    MessageDone {
        message_id: Uuid,
        #[serde(default, skip_serializing_if = "Option::is_none")]
        usage: Option<Usage>,
        status: MessageStatus,
    },
    Error {
        #[serde(default, skip_serializing_if = "Option::is_none")]
        request_id: Option<Uuid>,
        #[serde(default, skip_serializing_if = "Option::is_none")]
        message_id: Option<Uuid>,
        code: String,
        message: String,
    },
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    fn hello_math(resume_thread_id: Option<Uuid>) -> ClientMessage {
        ClientMessage::HelloMath {
            client_instance_id: "client-a".to_string(),
            version: "1.2.3".to_string(),
            min_supported_version: "1.0.0".to_string(),
            resume_thread_id,
            automagic: false,
            system_info: ClientSystemInfo::default(),
        }
    }

    #[test]
    fn send_message_omits_optional_updates_when_not_set() {
        let msg = ClientMessage::SendMessage {
            request_id: Uuid::nil(),
            thread_id: None,
            text: "hello".to_string(),
            runtime_mode: None,
            model_override: None,
        };

        let value = serde_json::to_value(msg).expect("serialize");
        let body = value
            .get("SendMessage")
            .and_then(|v| v.as_object())
            .expect("SendMessage body");

        assert!(body.get("runtime_mode").is_none());
        assert!(body.get("model_override").is_none());
    }

    #[test]
    fn hello_math_omits_resume_thread_id_when_not_set() {
        let msg = hello_math(None);

        let value = serde_json::to_value(msg).expect("serialize");
        let body = value
            .get("HelloMath")
            .and_then(|v| v.as_object())
            .expect("HelloMath body");

        assert!(body.get("resume_thread_id").is_none());
        assert!(body.get("automagic").is_none());
        assert!(body.get("system_info").is_none());
    }

    #[test]
    fn hello_math_round_trip_resume_thread_id() {
        let thread_id = Uuid::new_v4();
        let msg = hello_math(Some(thread_id));

        let value = serde_json::to_value(&msg).expect("serialize");
        let body = value
            .get("HelloMath")
            .and_then(|v| v.as_object())
            .expect("HelloMath body");
        assert_eq!(
            body.get("resume_thread_id"),
            Some(&serde_json::Value::String(thread_id.to_string()))
        );

        let back: ClientMessage = serde_json::from_value(value).expect("deserialize");
        match back {
            ClientMessage::HelloMath {
                resume_thread_id, ..
            } => assert_eq!(resume_thread_id, Some(thread_id)),
            _ => panic!("expected HelloMath"),
        }
    }

    #[test]
    fn hello_math_deserializes_defaults_for_new_fields() {
        let value = json!({
            "HelloMath": {
                "client_instance_id": "client-a",
                "version": "1.2.3",
                "min_supported_version": "1.0.0"
            }
        });

        let back: ClientMessage = serde_json::from_value(value).expect("deserialize");
        match back {
            ClientMessage::HelloMath {
                automagic,
                system_info,
                ..
            } => {
                assert!(!automagic);
                assert_eq!(system_info, ClientSystemInfo::default());
            }
            _ => panic!("expected HelloMath"),
        }
    }

    #[test]
    fn hello_math_serializes_non_default_system_info() {
        let msg = ClientMessage::HelloMath {
            client_instance_id: "client-a".to_string(),
            version: "1.2.3".to_string(),
            min_supported_version: "1.0.0".to_string(),
            resume_thread_id: None,
            automagic: true,
            system_info: ClientSystemInfo {
                os: Os::MacOS,
                os_version: "15.5".to_string(),
                arch: Arch::Amd64,
                cpu_cores: 10,
                ram_mb: 32768,
            },
        };

        let value = serde_json::to_value(msg).expect("serialize");
        let body = value
            .get("HelloMath")
            .and_then(|v| v.as_object())
            .expect("HelloMath body");

        assert_eq!(body.get("automagic"), Some(&json!(true)));
        assert_eq!(
            body.get("system_info"),
            Some(&json!({
                "os": "macos",
                "os_version": "15.5",
                "arch": "amd64",
                "cpu_cores": 10,
                "ram_mb": 32768
            }))
        );
    }

    #[test]
    fn hello_math_deserializes_canonical_arch_names() {
        let value = json!({
            "HelloMath": {
                "client_instance_id": "client-a",
                "version": "1.2.3",
                "min_supported_version": "1.0.0",
                "system_info": {
                    "os": "linux",
                    "os_version": "6.8",
                    "arch": "amd64",
                    "cpu_cores": 8,
                    "ram_mb": 16384
                }
            }
        });

        let back: ClientMessage = serde_json::from_value(value).expect("deserialize");
        match back {
            ClientMessage::HelloMath { system_info, .. } => {
                assert_eq!(system_info.arch, Arch::Amd64);
            }
            _ => panic!("expected HelloMath"),
        }
    }

    #[test]
    fn send_message_serializes_model_override_when_set() {
        let msg = ClientMessage::SendMessage {
            request_id: Uuid::nil(),
            thread_id: Some(Uuid::nil()),
            text: "hello".to_string(),
            runtime_mode: Some(RuntimeMode::Plan),
            model_override: Some(ThreadModelOverride {
                model_id: Uuid::nil(),
                model_config: ModelConfig::default(),
            }),
        };

        let value = serde_json::to_value(msg).expect("serialize");
        let body = value
            .get("SendMessage")
            .and_then(|v| v.as_object())
            .expect("SendMessage body");

        assert_eq!(body.get("runtime_mode"), Some(&json!("Plan")));
        assert!(body.get("model_override").is_some());
    }

    #[test]
    fn send_message_deserializes_model_override_states() {
        let set_json = json!({
            "SendMessage": {
                "request_id": Uuid::nil(),
                "thread_id": Uuid::nil(),
                "text": "hello",
                "runtime_mode": "Agent",
                "model_override": {
                    "model_id": Uuid::nil(),
                    "model_config": {}
                }
            }
        });
        let keep_json = json!({
            "SendMessage": {
                "request_id": Uuid::nil(),
                "thread_id": Uuid::nil(),
                "text": "hello"
            }
        });

        let set_msg: ClientMessage = serde_json::from_value(set_json).expect("deserialize set");
        let keep_msg: ClientMessage = serde_json::from_value(keep_json).expect("deserialize keep");

        match set_msg {
            ClientMessage::SendMessage {
                runtime_mode,
                model_override,
                ..
            } => {
                assert_eq!(runtime_mode, Some(RuntimeMode::Agent));
                assert!(model_override.is_some());
            }
            _ => panic!("expected SendMessage"),
        }

        match keep_msg {
            ClientMessage::SendMessage { model_override, .. } => {
                assert_eq!(model_override, None);
            }
            _ => panic!("expected SendMessage"),
        }
    }

    #[test]
    fn update_local_skills_round_trip_full_and_empty() {
        let full = ClientMessage::UpdateLocalSkills {
            skills: vec![LocalSkillMetadata {
                name: "Build skill".to_string(),
                description: "Run and fix build failures".to_string(),
                cwd: "/Users/dev/project".to_string(),
            }],
        };
        let empty = ClientMessage::UpdateLocalSkills { skills: vec![] };

        let full_json = serde_json::to_value(&full).expect("serialize full");
        let empty_json = serde_json::to_value(&empty).expect("serialize empty");

        let full_back: ClientMessage = serde_json::from_value(full_json).expect("deserialize full");
        let empty_back: ClientMessage =
            serde_json::from_value(empty_json).expect("deserialize empty");

        match full_back {
            ClientMessage::UpdateLocalSkills { skills } => {
                assert_eq!(skills.len(), 1);
                assert_eq!(skills[0].name, "Build skill");
                assert_eq!(skills[0].description, "Run and fix build failures");
                assert_eq!(skills[0].cwd, "/Users/dev/project");
            }
            _ => panic!("expected UpdateLocalSkills"),
        }

        match empty_back {
            ClientMessage::UpdateLocalSkills { skills } => {
                assert!(skills.is_empty());
            }
            _ => panic!("expected UpdateLocalSkills"),
        }
    }

    #[test]
    fn resolve_subagent_escalation_rejected_round_trip() {
        let msg = ClientMessage::ResolveSubagentEscalation {
            parent_message_id: Uuid::nil(),
            subagent_run_id: Uuid::nil(),
            escalation_id: "esc-1".to_string(),
            resolution: SubagentEscalationResolution::Rejected {
                reason: Some("not now".to_string()),
            },
        };

        let value = serde_json::to_value(&msg).expect("serialize");
        let body = value
            .get("ResolveSubagentEscalation")
            .and_then(|v| v.as_object())
            .expect("ResolveSubagentEscalation body");
        assert_eq!(body.get("escalation_id"), Some(&json!("esc-1")));

        let back: ClientMessage = serde_json::from_value(value).expect("deserialize");
        match back {
            ClientMessage::ResolveSubagentEscalation {
                resolution:
                    SubagentEscalationResolution::Rejected {
                        reason: Some(reason),
                    },
                ..
            } => assert_eq!(reason, "not now"),
            _ => panic!("expected ResolveSubagentEscalation::Rejected"),
        }
    }

    #[test]
    fn resolve_subagent_escalation_resolved_with_output_round_trip() {
        let msg = ClientMessage::ResolveSubagentEscalation {
            parent_message_id: Uuid::nil(),
            subagent_run_id: Uuid::nil(),
            escalation_id: "esc-2".to_string(),
            resolution: SubagentEscalationResolution::ResolvedWithOutput {
                is_error: false,
                output: "ok".to_string(),
                duration_seconds: Some(3),
            },
        };

        let value = serde_json::to_value(&msg).expect("serialize");
        let back: ClientMessage = serde_json::from_value(value).expect("deserialize");
        match back {
            ClientMessage::ResolveSubagentEscalation {
                escalation_id,
                resolution:
                    SubagentEscalationResolution::ResolvedWithOutput {
                        is_error,
                        output,
                        duration_seconds,
                    },
                ..
            } => {
                assert_eq!(escalation_id, "esc-2");
                assert!(!is_error);
                assert_eq!(output, "ok");
                assert_eq!(duration_seconds, Some(3));
            }
            _ => panic!("expected ResolveSubagentEscalation::ResolvedWithOutput"),
        }
    }
}