path: root/makima/src/server/state.rs

//! Application state holding shared ML models and database pool.

use std::sync::Arc;
use dashmap::DashMap;
use sqlx::PgPool;
use tokio::sync::{broadcast, mpsc, Mutex};
use uuid::Uuid;

use crate::listen::{DiarizationConfig, ParakeetEOU, ParakeetTDT, Sortformer};
use crate::server::auth::{AuthConfig, JwtVerifier};

/// Notification payload for file updates (broadcast to WebSocket subscribers).
#[derive(Debug, Clone)]
pub struct FileUpdateNotification {
    /// ID of the updated file
    pub file_id: Uuid,
    /// New version number after update
    pub version: i32,
    /// List of fields that were updated
    pub updated_fields: Vec<String>,
    /// Source of the update: "user", "llm", or "system"
    pub updated_by: String,
}

// =============================================================================
// Task/Mesh Notifications
// =============================================================================

/// Notification payload for task updates (broadcast to WebSocket subscribers).
#[derive(Debug, Clone)]
pub struct TaskUpdateNotification {
    /// ID of the updated task
    pub task_id: Uuid,
    /// Owner ID for data isolation (notifications are scoped to owner)
    pub owner_id: Option<Uuid>,
    /// New version number after update
    pub version: i32,
    /// Current task status
    pub status: String,
    /// List of fields that were updated
    pub updated_fields: Vec<String>,
    /// Source of the update: "user", "daemon", or "system"
    pub updated_by: String,
}

/// Notification for streaming task output from Claude Code containers.
#[derive(Debug, Clone, serde::Serialize)]
#[serde(rename_all = "camelCase")]
pub struct TaskOutputNotification {
    /// ID of the task producing output
    pub task_id: Uuid,
    /// Owner ID for data isolation (notifications are scoped to owner)
    #[serde(skip)]
    pub owner_id: Option<Uuid>,
    /// Type of message: "assistant", "tool_use", "tool_result", "result", "system", "error", "raw"
    pub message_type: String,
    /// Main text content of the message
    pub content: String,
    /// Tool name if this is a tool_use message
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_name: Option<String>,
    /// Tool input (JSON) if this is a tool_use message
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_input: Option<serde_json::Value>,
    /// Whether tool result was an error
    #[serde(skip_serializing_if = "Option::is_none")]
    pub is_error: Option<bool>,
    /// Cost in USD if this is a result message
    #[serde(skip_serializing_if = "Option::is_none")]
    pub cost_usd: Option<f64>,
    /// Duration in milliseconds if this is a result message
    #[serde(skip_serializing_if = "Option::is_none")]
    pub duration_ms: Option<u64>,
    /// Whether this is a partial line (more coming) or complete
    pub is_partial: bool,
}

/// Command sent from server to daemon.
#[derive(Debug, Clone, serde::Serialize)]
#[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>>,
    },
    /// 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,
    },

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

/// Active daemon connection info stored in state.
#[derive(Debug)]
pub struct DaemonConnectionInfo {
    /// Database ID of the daemon
    pub id: Uuid,
    /// Owner ID for data isolation (from API key authentication)
    pub owner_id: Uuid,
    /// WebSocket connection identifier
    pub connection_id: String,
    /// Daemon hostname
    pub hostname: Option<String>,
    /// Machine identifier
    pub machine_id: Option<String>,
    /// Channel to send commands to this daemon
    pub command_sender: mpsc::Sender<DaemonCommand>,
    /// Current working directory of the daemon
    pub working_directory: Option<String>,
    /// Path to ~/.makima/home directory on daemon (for cloning completed work)
    pub home_directory: Option<String>,
    /// Path to worktrees directory (~/.makima/worktrees) on daemon
    pub worktrees_directory: Option<String>,
}

/// Shared application state containing ML models and database pool.
///
/// Models are wrapped in `Mutex` for thread-safe mutable access during inference.
pub struct AppState {
    /// Speech-to-text model (Parakeet TDT)
    pub parakeet: Mutex<ParakeetTDT>,
    /// End-of-Utterance detection model for streaming
    pub parakeet_eou: Mutex<ParakeetEOU>,
    /// Speaker diarization model (Sortformer)
    pub sortformer: Mutex<Sortformer>,
    /// Optional database connection pool
    pub db_pool: Option<PgPool>,
    /// Broadcast channel for file update notifications
    pub file_updates: broadcast::Sender<FileUpdateNotification>,
    /// Broadcast channel for task update notifications
    pub task_updates: broadcast::Sender<TaskUpdateNotification>,
    /// Broadcast channel for task output streaming
    pub task_output: broadcast::Sender<TaskOutputNotification>,
    /// Active daemon connections (keyed by connection_id)
    pub daemon_connections: DashMap<String, DaemonConnectionInfo>,
    /// Tool keys for orchestrator API access (key -> task_id)
    pub tool_keys: DashMap<String, Uuid>,
    /// JWT verifier for Supabase authentication (None if not configured)
    pub jwt_verifier: Option<JwtVerifier>,
}

impl AppState {
    /// Load all ML models from the specified directories.
    ///
    /// # Arguments
    /// * `parakeet_model_dir` - Path to the Parakeet TDT model directory
    /// * `parakeet_eou_dir` - Path to the Parakeet EOU model directory
    /// * `sortformer_model_path` - Path to the Sortformer diarization model file
    pub fn new(
        parakeet_model_dir: &str,
        parakeet_eou_dir: &str,
        sortformer_model_path: &str,
    ) -> Result<Self, Box<dyn std::error::Error + Send + Sync>> {
        let parakeet = ParakeetTDT::from_pretrained(parakeet_model_dir, None)?;
        let parakeet_eou = ParakeetEOU::from_pretrained(parakeet_eou_dir, None)?;
        let sortformer = Sortformer::with_config(
            sortformer_model_path,
            None,
            DiarizationConfig::callhome(),
        )?;

        // Create broadcast channels with buffer for 256 messages
        let (file_updates, _) = broadcast::channel(256);
        let (task_updates, _) = broadcast::channel(256);
        let (task_output, _) = broadcast::channel(1024); // Larger buffer for output streaming

        // Initialize JWT verifier from environment (optional)
        // Requires SUPABASE_URL and either SUPABASE_JWT_PUBLIC_KEY (RS256) or SUPABASE_JWT_SECRET (HS256)
        let jwt_verifier = match AuthConfig::from_env() {
            Some(config) => match JwtVerifier::new(config) {
                Ok(verifier) => {
                    tracing::info!("JWT authentication configured");
                    Some(verifier)
                }
                Err(e) => {
                    tracing::error!("Failed to initialize JWT verifier: {}", e);
                    None
                }
            },
            None => {
                // Log which env vars are missing
                let has_url = std::env::var("SUPABASE_URL").is_ok();
                let has_public_key = std::env::var("SUPABASE_JWT_PUBLIC_KEY").is_ok();
                let has_secret = std::env::var("SUPABASE_JWT_SECRET").is_ok();

                if !has_url {
                    tracing::info!("JWT authentication not configured (SUPABASE_URL not set)");
                } else if !has_public_key && !has_secret {
                    tracing::info!("JWT authentication not configured (set SUPABASE_JWT_PUBLIC_KEY for RS256 or SUPABASE_JWT_SECRET for HS256)");
                }
                None
            }
        };

        Ok(Self {
            parakeet: Mutex::new(parakeet),
            parakeet_eou: Mutex::new(parakeet_eou),
            sortformer: Mutex::new(sortformer),
            db_pool: None,
            file_updates,
            task_updates,
            task_output,
            daemon_connections: DashMap::new(),
            tool_keys: DashMap::new(),
            jwt_verifier,
        })
    }

    /// Set the database pool.
    pub fn with_db_pool(mut self, pool: PgPool) -> Self {
        self.db_pool = Some(pool);
        self
    }

    /// Broadcast a file update notification to all subscribers.
    ///
    /// This is a no-op if there are no subscribers (ignores send errors).
    pub fn broadcast_file_update(&self, notification: FileUpdateNotification) {
        // Ignore send errors - they just mean no one is listening
        let _ = self.file_updates.send(notification);
    }

    /// Broadcast a task update notification to all subscribers.
    ///
    /// This is a no-op if there are no subscribers (ignores send errors).
    pub fn broadcast_task_update(&self, notification: TaskUpdateNotification) {
        let _ = self.task_updates.send(notification);
    }

    /// Broadcast task output to all subscribers.
    ///
    /// Used for streaming Claude Code container output to frontend clients.
    pub fn broadcast_task_output(&self, notification: TaskOutputNotification) {
        let _ = self.task_output.send(notification);
    }

    /// Register a new daemon connection.
    ///
    /// Returns the connection_id for later reference.
    pub fn register_daemon(
        &self,
        connection_id: String,
        daemon_id: Uuid,
        owner_id: Uuid,
        hostname: Option<String>,
        machine_id: Option<String>,
        command_sender: mpsc::Sender<DaemonCommand>,
    ) {
        self.daemon_connections.insert(
            connection_id.clone(),
            DaemonConnectionInfo {
                id: daemon_id,
                owner_id,
                connection_id,
                hostname,
                machine_id,
                command_sender,
                working_directory: None,
                home_directory: None,
                worktrees_directory: None,
            },
        );
    }

    /// Update daemon directory information.
    pub fn update_daemon_directories(
        &self,
        connection_id: &str,
        working_directory: String,
        home_directory: String,
        worktrees_directory: String,
    ) {
        if let Some(mut entry) = self.daemon_connections.get_mut(connection_id) {
            entry.working_directory = Some(working_directory);
            entry.home_directory = Some(home_directory);
            entry.worktrees_directory = Some(worktrees_directory);
        }
    }

    /// Unregister a daemon connection.
    pub fn unregister_daemon(&self, connection_id: &str) {
        self.daemon_connections.remove(connection_id);
    }

    /// Get a daemon connection by connection_id.
    pub fn get_daemon(&self, connection_id: &str) -> Option<dashmap::mapref::one::Ref<'_, String, DaemonConnectionInfo>> {
        self.daemon_connections.get(connection_id)
    }

    /// Get a daemon by its database ID.
    pub fn get_daemon_by_id(&self, daemon_id: Uuid) -> Option<dashmap::mapref::one::Ref<'_, String, DaemonConnectionInfo>> {
        self.daemon_connections
            .iter()
            .find(|entry| entry.value().id == daemon_id)
            .map(|entry| {
                // Return a reference to the found entry
                self.daemon_connections.get(entry.key()).unwrap()
            })
    }

    /// Send a command to a specific daemon by its database ID.
    pub async fn send_daemon_command(&self, daemon_id: Uuid, command: DaemonCommand) -> Result<(), String> {
        if let Some(daemon) = self.daemon_connections
            .iter()
            .find(|entry| entry.value().id == daemon_id)
        {
            daemon.value().command_sender.send(command).await
                .map_err(|e| format!("Failed to send command to daemon: {}", e))
        } else {
            Err(format!("Daemon {} not connected", daemon_id))
        }
    }

    /// Broadcast sibling progress to all running sibling tasks.
    ///
    /// This is used for sibling awareness - when a task makes progress,
    /// its siblings are notified so they can adjust their work if needed.
    pub async fn broadcast_sibling_progress(
        &self,
        source_task_id: Uuid,
        source_task_name: &str,
        source_task_status: &str,
        progress_summary: Option<String>,
        changed_files: Vec<String>,
        running_sibling_daemon_ids: Vec<(Uuid, Uuid)>, // (task_id, daemon_id)
    ) {
        for (sibling_task_id, daemon_id) in running_sibling_daemon_ids {
            let command = DaemonCommand::InjectSiblingContext {
                task_id: sibling_task_id,
                sibling_task_id: source_task_id,
                sibling_name: source_task_name.to_string(),
                sibling_status: source_task_status.to_string(),
                progress_summary: progress_summary.clone(),
                changed_files: changed_files.clone(),
            };

            // Fire and forget - don't block on sending to all daemons
            if let Err(e) = self.send_daemon_command(daemon_id, command).await {
                tracing::warn!(
                    "Failed to inject sibling context to task {}: {}",
                    sibling_task_id,
                    e
                );
            }
        }
    }

    /// Get list of connected daemon IDs.
    pub fn list_connected_daemon_ids(&self) -> Vec<Uuid> {
        self.daemon_connections
            .iter()
            .map(|entry| entry.value().id)
            .collect()
    }

    // =========================================================================
    // Tool Key Management
    // =========================================================================

    /// Register a tool key for a task.
    ///
    /// This allows orchestrators to authenticate with the API using
    /// the `X-Makima-Tool-Key` header.
    pub fn register_tool_key(&self, key: String, task_id: Uuid) {
        tracing::info!(task_id = %task_id, "Registering tool key");
        self.tool_keys.insert(key, task_id);
    }

    /// Validate a tool key and return the associated task ID.
    pub fn validate_tool_key(&self, key: &str) -> Option<Uuid> {
        self.tool_keys.get(key).map(|entry| *entry.value())
    }

    /// Revoke a tool key for a task.
    ///
    /// This should be called when a task completes or is terminated.
    pub fn revoke_tool_key(&self, task_id: Uuid) {
        // Find and remove the key for this task
        self.tool_keys.retain(|_, v| *v != task_id);
        tracing::info!(task_id = %task_id, "Revoked tool key");
    }
}

/// Type alias for the shared application state.
pub type SharedState = Arc<AppState>;