path: root/makima/src/daemon/ws/protocol.rs

//! Protocol types for daemon-server communication.
//!
//! These types mirror the server's protocol exactly for compatibility.

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

/// Message from daemon to server.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "camelCase")]
pub enum DaemonMessage {
    /// Authentication request (first message required).
    Authenticate {
        #[serde(rename = "apiKey")]
        api_key: String,
        #[serde(rename = "machineId")]
        machine_id: String,
        hostname: String,
        #[serde(rename = "maxConcurrentTasks")]
        max_concurrent_tasks: i32,
    },

    /// Periodic heartbeat with current status.
    Heartbeat {
        #[serde(rename = "activeTasks")]
        active_tasks: Vec<Uuid>,
    },

    /// Task output streaming (stdout/stderr from Claude Code).
    TaskOutput {
        #[serde(rename = "taskId")]
        task_id: Uuid,
        output: String,
        #[serde(rename = "isPartial")]
        is_partial: bool,
    },

    /// Task status change notification.
    TaskStatusChange {
        #[serde(rename = "taskId")]
        task_id: Uuid,
        #[serde(rename = "oldStatus")]
        old_status: String,
        #[serde(rename = "newStatus")]
        new_status: String,
    },

    /// Task progress update with summary.
    TaskProgress {
        #[serde(rename = "taskId")]
        task_id: Uuid,
        summary: String,
    },

    /// Task completion notification.
    TaskComplete {
        #[serde(rename = "taskId")]
        task_id: Uuid,
        success: bool,
        error: Option<String>,
    },

    /// Register a tool key for orchestrator API access.
    RegisterToolKey {
        #[serde(rename = "taskId")]
        task_id: Uuid,
        /// The API key for this orchestrator to use when calling mesh endpoints.
        key: String,
    },

    /// Revoke a tool key when task completes.
    RevokeToolKey {
        #[serde(rename = "taskId")]
        task_id: Uuid,
    },

    /// Authentication required - OAuth token expired, provides login URL.
    AuthenticationRequired {
        /// Task ID that triggered the auth error (if any).
        #[serde(rename = "taskId")]
        task_id: Option<Uuid>,
        /// OAuth login URL for remote authentication.
        #[serde(rename = "loginUrl")]
        login_url: String,
        /// Hostname of the daemon requiring auth.
        hostname: Option<String>,
    },

    // =========================================================================
    // Merge Response Messages (sent by daemon after processing merge commands)
    // =========================================================================

    /// Response to ListBranches command.
    BranchList {
        #[serde(rename = "taskId")]
        task_id: Uuid,
        branches: Vec<BranchInfo>,
    },

    /// Response to MergeStatus command.
    MergeStatusResponse {
        #[serde(rename = "taskId")]
        task_id: Uuid,
        #[serde(rename = "inProgress")]
        in_progress: bool,
        #[serde(rename = "sourceBranch")]
        source_branch: Option<String>,
        #[serde(rename = "conflictedFiles")]
        conflicted_files: Vec<String>,
    },

    /// Response to merge operations (MergeStart, MergeResolve, MergeCommit, MergeAbort).
    MergeResult {
        #[serde(rename = "taskId")]
        task_id: Uuid,
        success: bool,
        message: String,
        #[serde(rename = "commitSha")]
        commit_sha: Option<String>,
        /// Present only when conflicts occurred.
        conflicts: Option<Vec<String>>,
    },

    /// Response to CheckMergeComplete command.
    MergeCompleteCheck {
        #[serde(rename = "taskId")]
        task_id: Uuid,
        #[serde(rename = "canComplete")]
        can_complete: bool,
        #[serde(rename = "unmergedBranches")]
        unmerged_branches: Vec<String>,
        #[serde(rename = "mergedCount")]
        merged_count: u32,
        #[serde(rename = "skippedCount")]
        skipped_count: u32,
    },

    // =========================================================================
    // Completion Action Response Messages
    // =========================================================================

    /// Response to RetryCompletionAction command.
    CompletionActionResult {
        #[serde(rename = "taskId")]
        task_id: Uuid,
        success: bool,
        message: String,
        /// PR URL if action was "pr" and successful.
        #[serde(rename = "prUrl")]
        pr_url: Option<String>,
    },

    /// Report daemon's available directories for task output.
    DaemonDirectories {
        /// Current working directory of the daemon.
        #[serde(rename = "workingDirectory")]
        working_directory: String,
        /// Path to ~/.makima/home directory (for cloning completed work).
        #[serde(rename = "homeDirectory")]
        home_directory: String,
        /// Path to worktrees directory (~/.makima/worktrees).
        #[serde(rename = "worktreesDirectory")]
        worktrees_directory: String,
    },

    /// Response to CloneWorktree command.
    CloneWorktreeResult {
        #[serde(rename = "taskId")]
        task_id: Uuid,
        success: bool,
        message: String,
        /// The path where the worktree was cloned.
        #[serde(rename = "targetDir")]
        target_dir: Option<String>,
    },

    /// Response to CheckTargetExists command.
    CheckTargetExistsResult {
        #[serde(rename = "taskId")]
        task_id: Uuid,
        /// Whether the target directory exists.
        exists: bool,
        /// The path that was checked.
        #[serde(rename = "targetDir")]
        target_dir: String,
    },

    // =========================================================================
    // Contract File Response Messages
    // =========================================================================

    /// Response to ReadRepoFile command.
    RepoFileContent {
        /// Request ID from the original command.
        #[serde(rename = "requestId")]
        request_id: Uuid,
        /// Path to the file that was read.
        #[serde(rename = "filePath")]
        file_path: String,
        /// File content (None if error occurred).
        content: Option<String>,
        /// Whether the operation succeeded.
        success: bool,
        /// Error message if operation failed.
        error: Option<String>,
    },

    // =========================================================================
    // Supervisor Git Response Messages
    // =========================================================================

    /// Response to CreateBranch command.
    BranchCreated {
        #[serde(rename = "taskId")]
        task_id: Uuid,
        success: bool,
        #[serde(rename = "branchName")]
        branch_name: String,
        message: String,
    },

    /// Response to MergeTaskToTarget command.
    MergeToTargetResult {
        #[serde(rename = "taskId")]
        task_id: Uuid,
        success: bool,
        message: String,
        #[serde(rename = "commitSha")]
        commit_sha: Option<String>,
        conflicts: Option<Vec<String>>,
    },

    /// Response to CreatePR command.
    PRCreated {
        #[serde(rename = "taskId")]
        task_id: Uuid,
        success: bool,
        message: String,
        #[serde(rename = "prUrl")]
        pr_url: Option<String>,
        #[serde(rename = "prNumber")]
        pr_number: Option<i32>,
    },

    /// Response to GetTaskDiff command.
    TaskDiff {
        #[serde(rename = "taskId")]
        task_id: Uuid,
        success: bool,
        diff: Option<String>,
        error: Option<String>,
    },

    /// Response to CleanupWorktree command.
    CleanupWorktreeResult {
        #[serde(rename = "taskId")]
        task_id: Uuid,
        success: bool,
        message: String,
    },
}

/// Information about a branch (used in BranchList message).
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct BranchInfo {
    /// Full branch name.
    pub name: String,
    /// Task ID extracted from branch name (if parseable).
    #[serde(rename = "taskId")]
    pub task_id: Option<Uuid>,
    /// Whether this branch has been merged.
    #[serde(rename = "isMerged")]
    pub is_merged: bool,
    /// Short SHA of the last commit.
    #[serde(rename = "lastCommit")]
    pub last_commit: String,
    /// Subject line of the last commit.
    #[serde(rename = "lastCommitMessage")]
    pub last_commit_message: String,
}

/// Command from server to daemon.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "camelCase")]
pub enum DaemonCommand {
    /// Confirm successful authentication.
    Authenticated {
        #[serde(rename = "daemonId")]
        daemon_id: Uuid,
    },

    /// Spawn a new task in a container.
    SpawnTask {
        #[serde(rename = "taskId")]
        task_id: Uuid,
        /// Human-readable task name (used for commit messages).
        #[serde(rename = "taskName")]
        task_name: String,
        plan: String,
        #[serde(rename = "repoUrl")]
        repo_url: Option<String>,
        #[serde(rename = "baseBranch")]
        base_branch: Option<String>,
        /// Target branch to merge into (used for completion actions).
        #[serde(rename = "targetBranch")]
        target_branch: Option<String>,
        /// Parent task ID if this is a subtask.
        #[serde(rename = "parentTaskId")]
        parent_task_id: Option<Uuid>,
        /// Depth in task hierarchy (0=top-level, 1=subtask, 2=sub-subtask).
        depth: i32,
        /// Whether this task should run as an orchestrator (true if depth==0 and has subtasks).
        #[serde(rename = "isOrchestrator")]
        is_orchestrator: bool,
        /// Path to user's local repository (outside ~/.makima) for completion actions.
        #[serde(rename = "targetRepoPath")]
        target_repo_path: Option<String>,
        /// Action on completion: "none", "branch", "merge", "pr".
        #[serde(rename = "completionAction")]
        completion_action: Option<String>,
        /// Task ID to continue from (copy worktree from this task).
        #[serde(rename = "continueFromTaskId")]
        continue_from_task_id: Option<Uuid>,
        /// Files to copy from parent task's worktree.
        #[serde(rename = "copyFiles")]
        copy_files: Option<Vec<String>>,
        /// Contract ID if this task is associated with a contract.
        #[serde(rename = "contractId")]
        contract_id: Option<Uuid>,
        /// Whether this task is a supervisor (long-running contract orchestrator).
        #[serde(rename = "isSupervisor", default)]
        is_supervisor: bool,
        /// Whether to run in autonomous loop mode.
        /// When enabled, task will automatically restart with --continue if it exits
        /// without a COMPLETION_GATE indicating ready: true.
        #[serde(rename = "autonomousLoop", default)]
        autonomous_loop: bool,
    },

    /// Pause a running task.
    PauseTask {
        #[serde(rename = "taskId")]
        task_id: Uuid,
    },

    /// Resume a paused task.
    ResumeTask {
        #[serde(rename = "taskId")]
        task_id: Uuid,
    },

    /// Interrupt a task (gracefully or forced).
    InterruptTask {
        #[serde(rename = "taskId")]
        task_id: Uuid,
        graceful: bool,
    },

    /// Send a message to a running task.
    SendMessage {
        #[serde(rename = "taskId")]
        task_id: Uuid,
        message: String,
    },

    /// Inject context about sibling task progress.
    InjectSiblingContext {
        #[serde(rename = "taskId")]
        task_id: Uuid,
        #[serde(rename = "siblingTaskId")]
        sibling_task_id: Uuid,
        #[serde(rename = "siblingName")]
        sibling_name: String,
        #[serde(rename = "siblingStatus")]
        sibling_status: String,
        #[serde(rename = "progressSummary")]
        progress_summary: Option<String>,
        #[serde(rename = "changedFiles")]
        changed_files: Vec<String>,
    },

    // =========================================================================
    // Merge Commands (for orchestrators to merge subtask branches)
    // =========================================================================

    /// List all subtask branches for a task.
    ListBranches {
        #[serde(rename = "taskId")]
        task_id: Uuid,
    },

    /// Start merging a subtask branch.
    MergeStart {
        #[serde(rename = "taskId")]
        task_id: Uuid,
        #[serde(rename = "sourceBranch")]
        source_branch: String,
    },

    /// Get current merge status.
    MergeStatus {
        #[serde(rename = "taskId")]
        task_id: Uuid,
    },

    /// Resolve a merge conflict.
    MergeResolve {
        #[serde(rename = "taskId")]
        task_id: Uuid,
        file: String,
        /// "ours" or "theirs"
        strategy: String,
    },

    /// Commit the current merge.
    MergeCommit {
        #[serde(rename = "taskId")]
        task_id: Uuid,
        message: String,
    },

    /// Abort the current merge.
    MergeAbort {
        #[serde(rename = "taskId")]
        task_id: Uuid,
    },

    /// Skip merging a subtask branch (mark as intentionally not merged).
    MergeSkip {
        #[serde(rename = "taskId")]
        task_id: Uuid,
        #[serde(rename = "subtaskId")]
        subtask_id: Uuid,
        reason: String,
    },

    /// Check if all subtask branches have been merged or skipped (completion gate).
    CheckMergeComplete {
        #[serde(rename = "taskId")]
        task_id: Uuid,
    },

    // =========================================================================
    // Completion Action Commands
    // =========================================================================

    /// Retry a completion action for a completed task.
    RetryCompletionAction {
        #[serde(rename = "taskId")]
        task_id: Uuid,
        /// Human-readable task name (used for commit messages).
        #[serde(rename = "taskName")]
        task_name: String,
        /// The action to execute: "branch", "merge", or "pr".
        action: String,
        /// Path to the target repository.
        #[serde(rename = "targetRepoPath")]
        target_repo_path: String,
        /// Target branch to merge into (for merge/pr actions).
        #[serde(rename = "targetBranch")]
        target_branch: Option<String>,
    },

    /// Clone worktree to a target directory.
    CloneWorktree {
        #[serde(rename = "taskId")]
        task_id: Uuid,
        /// Path to the target directory.
        #[serde(rename = "targetDir")]
        target_dir: String,
    },

    /// Check if a target directory exists.
    CheckTargetExists {
        #[serde(rename = "taskId")]
        task_id: Uuid,
        /// Path to check.
        #[serde(rename = "targetDir")]
        target_dir: String,
    },

    // =========================================================================
    // Contract File Commands
    // =========================================================================

    /// Read a file from a repository linked to a contract.
    ReadRepoFile {
        /// Request ID for correlating response.
        #[serde(rename = "requestId")]
        request_id: Uuid,
        /// Contract ID (used for logging/context).
        #[serde(rename = "contractId")]
        contract_id: Uuid,
        /// Path to the file within the repository.
        #[serde(rename = "filePath")]
        file_path: String,
        /// Full repository path on daemon's filesystem.
        #[serde(rename = "repoPath")]
        repo_path: String,
    },

    // =========================================================================
    // Supervisor Git Commands
    // =========================================================================

    /// Create a new branch in the supervisor's worktree.
    CreateBranch {
        #[serde(rename = "taskId")]
        task_id: Uuid,
        #[serde(rename = "branchName")]
        branch_name: String,
        /// Optional reference to create branch from (task_id or SHA).
        #[serde(rename = "fromRef")]
        from_ref: Option<String>,
    },

    /// Merge a task's changes to a target branch.
    MergeTaskToTarget {
        #[serde(rename = "taskId")]
        task_id: Uuid,
        /// Target branch to merge into (default: task's base branch).
        #[serde(rename = "targetBranch")]
        target_branch: Option<String>,
        /// Whether to squash commits.
        squash: bool,
    },

    /// Create a pull request for a task's changes.
    CreatePR {
        #[serde(rename = "taskId")]
        task_id: Uuid,
        title: String,
        body: Option<String>,
        /// Base branch for the PR (default: main).
        #[serde(rename = "baseBranch")]
        base_branch: String,
    },

    /// Get the diff for a task's changes.
    GetTaskDiff {
        #[serde(rename = "taskId")]
        task_id: Uuid,
    },

    /// Clean up a task's worktree (used when contract is completed/deleted).
    CleanupWorktree {
        #[serde(rename = "taskId")]
        task_id: Uuid,
        /// Whether to delete the associated branch.
        #[serde(rename = "deleteBranch")]
        delete_branch: bool,
    },

    /// Error response.
    Error {
        code: String,
        message: String,
    },
}

impl DaemonMessage {
    /// Create an authentication message.
    pub fn authenticate(
        api_key: &str,
        machine_id: &str,
        hostname: &str,
        max_concurrent_tasks: i32,
    ) -> Self {
        Self::Authenticate {
            api_key: api_key.to_string(),
            machine_id: machine_id.to_string(),
            hostname: hostname.to_string(),
            max_concurrent_tasks,
        }
    }

    /// Create a heartbeat message.
    pub fn heartbeat(active_tasks: Vec<Uuid>) -> Self {
        Self::Heartbeat { active_tasks }
    }

    /// Create a task output message.
    pub fn task_output(task_id: Uuid, output: String, is_partial: bool) -> Self {
        Self::TaskOutput {
            task_id,
            output,
            is_partial,
        }
    }

    /// Create a task status change message.
    pub fn task_status_change(task_id: Uuid, old_status: &str, new_status: &str) -> Self {
        Self::TaskStatusChange {
            task_id,
            old_status: old_status.to_string(),
            new_status: new_status.to_string(),
        }
    }

    /// Create a task progress message.
    pub fn task_progress(task_id: Uuid, summary: String) -> Self {
        Self::TaskProgress { task_id, summary }
    }

    /// Create a task complete message.
    pub fn task_complete(task_id: Uuid, success: bool, error: Option<String>) -> Self {
        Self::TaskComplete {
            task_id,
            success,
            error,
        }
    }

    /// Create a register tool key message.
    pub fn register_tool_key(task_id: Uuid, key: String) -> Self {
        Self::RegisterToolKey { task_id, key }
    }

    /// Create a revoke tool key message.
    pub fn revoke_tool_key(task_id: Uuid) -> Self {
        Self::RevokeToolKey { task_id }
    }
}

#[cfg(test)]
mod tests {
    use crate::daemon::*;

    #[test]
    fn test_daemon_message_serialization() {
        let msg = DaemonMessage::authenticate("key123", "machine-abc", "worker-1", 4);
        let json = serde_json::to_string(&msg).unwrap();
        assert!(json.contains("\"type\":\"authenticate\""));
        assert!(json.contains("\"apiKey\":\"key123\""));
        assert!(json.contains("\"machineId\":\"machine-abc\""));
    }

    #[test]
    fn test_daemon_command_deserialization() {
        let json = r#"{"type":"spawnTask","taskId":"550e8400-e29b-41d4-a716-446655440000","plan":"Build the feature","repoUrl":"https://github.com/test/repo","baseBranch":"main","parentTaskId":null,"depth":0,"isOrchestrator":false}"#;
        let cmd: DaemonCommand = serde_json::from_str(json).unwrap();
        match cmd {
            DaemonCommand::SpawnTask {
                plan,
                repo_url,
                base_branch,
                parent_task_id,
                depth,
                is_orchestrator,
                ..
            } => {
                assert_eq!(plan, "Build the feature");
                assert_eq!(repo_url, Some("https://github.com/test/repo".to_string()));
                assert_eq!(base_branch, Some("main".to_string()));
                assert_eq!(parent_task_id, None);
                assert_eq!(depth, 0);
                assert!(!is_orchestrator);
            }
            _ => panic!("Expected SpawnTask"),
        }
    }

    #[test]
    fn test_orchestrator_spawn_deserialization() {
        let json = r#"{"type":"spawnTask","taskId":"550e8400-e29b-41d4-a716-446655440000","plan":"Coordinate subtasks","repoUrl":"https://github.com/test/repo","baseBranch":"main","parentTaskId":null,"depth":0,"isOrchestrator":true}"#;
        let cmd: DaemonCommand = serde_json::from_str(json).unwrap();
        match cmd {
            DaemonCommand::SpawnTask {
                is_orchestrator,
                depth,
                ..
            } => {
                assert!(is_orchestrator);
                assert_eq!(depth, 0);
            }
            _ => panic!("Expected SpawnTask"),
        }
    }
}