client.rs

use crate::daemon::{Daemon, RunOptions};
use crate::error::IpcError;
use crate::ipc::{IpcRequest, IpcResponse, deserialize, fs_name, serialize};
use crate::{Result, supervisor};
use exponential_backoff::Backoff;
use interprocess::local_socket::tokio::{RecvHalf, SendHalf};
use interprocess::local_socket::traits::tokio::Stream;
use miette::Context;
use std::path::PathBuf;
use std::time::Duration;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::sync::Mutex;
use uuid::Uuid;

pub struct IpcClient {
    _id: String,
    recv: Mutex<BufReader<RecvHalf>>,
    send: Mutex<SendHalf>,
}

const CONNECT_ATTEMPTS: u32 = 5;
const CONNECT_MIN_DELAY: Duration = Duration::from_millis(100);
const CONNECT_MAX_DELAY: Duration = Duration::from_secs(1);
const REQUEST_TIMEOUT: Duration = Duration::from_secs(5);

impl IpcClient {
    pub async fn connect(autostart: bool) -> Result<Self> {
        if autostart {
            supervisor::start_if_not_running()?;
        }
        let id = Uuid::new_v4().to_string();
        let client = Self::connect_(&id, "main").await?;
        trace!("Connected to IPC socket");
        let rsp = client.request(IpcRequest::Connect).await?;
        if !rsp.is_ok() {
            return Err(IpcError::UnexpectedResponse {
                expected: "Ok".to_string(),
                actual: format!("{:?}", rsp),
            }
            .into());
        }
        debug!("Connected to IPC main");
        Ok(client)
    }

    async fn connect_(id: &str, name: &str) -> Result<Self> {
        for duration in Backoff::new(CONNECT_ATTEMPTS, CONNECT_MIN_DELAY, CONNECT_MAX_DELAY) {
            match interprocess::local_socket::tokio::Stream::connect(fs_name(name)?).await {
                Ok(conn) => {
                    let (recv, send) = conn.split();
                    let recv = BufReader::new(recv);

                    return Ok(Self {
                        _id: id.to_string(),
                        recv: Mutex::new(recv),
                        send: Mutex::new(send),
                    });
                }
                Err(err) => {
                    if let Some(duration) = duration {
                        debug!(
                            "Failed to connect to IPC socket: {:?}, retrying in {:?}",
                            err, duration
                        );
                        tokio::time::sleep(duration).await;
                        continue;
                    } else {
                        return Err(IpcError::ConnectionFailed {
                            attempts: CONNECT_ATTEMPTS,
                            details: Some(format!(
                                "{err}\nensure the supervisor is running with: pitchfork supervisor start"
                            )),
                        }
                        .into());
                    }
                }
            }
        }
        Err(IpcError::ConnectionFailed {
            attempts: CONNECT_ATTEMPTS,
            details: Some(
                "ensure the supervisor is running with: pitchfork supervisor start".to_string(),
            ),
        }
        .into())
    }

    pub async fn send(&self, msg: IpcRequest) -> Result<()> {
        let mut msg = serialize(&msg)?;
        if msg.contains(&0) {
            return Err(IpcError::InvalidMessage {
                details: Some("message contains null byte".to_string()),
            }
            .into());
        }
        msg.push(0);
        let mut send = self.send.lock().await;
        send.write_all(&msg)
            .await
            .map_err(|e| IpcError::SendFailed {
                details: Some(e.to_string()),
            })?;
        Ok(())
    }

    async fn read(&self, timeout: Duration) -> Result<IpcResponse> {
        let mut recv = self.recv.lock().await;
        let mut bytes = Vec::new();
        match tokio::time::timeout(timeout, recv.read_until(0, &mut bytes)).await {
            Ok(Ok(_)) => {}
            Ok(Err(err)) => {
                return Err(IpcError::ReadFailed {
                    details: Some(err.to_string()),
                }
                .into());
            }
            Err(_) => {
                return Err(IpcError::Timeout {
                    seconds: timeout.as_secs(),
                }
                .into());
            }
        }
        if bytes.is_empty() {
            return Err(IpcError::ConnectionClosed.into());
        }
        deserialize(&bytes).wrap_err("failed to deserialize IPC response")
    }

    async fn request(&self, msg: IpcRequest) -> Result<IpcResponse> {
        self.request_with_timeout(msg, REQUEST_TIMEOUT).await
    }

    fn unexpected_response(expected: &str, actual: &IpcResponse) -> IpcError {
        IpcError::UnexpectedResponse {
            expected: expected.to_string(),
            actual: format!("{:?}", actual),
        }
    }

    async fn request_with_timeout(
        &self,
        msg: IpcRequest,
        timeout: Duration,
    ) -> Result<IpcResponse> {
        self.send(msg).await?;
        self.read(timeout).await
    }

    pub async fn enable(&self, id: String) -> Result<bool> {
        let rsp = self.request(IpcRequest::Enable { id: id.clone() }).await?;
        match rsp {
            IpcResponse::Yes => {
                info!("enabled daemon {}", id);
                Ok(true)
            }
            IpcResponse::No => {
                info!("daemon {} already enabled", id);
                Ok(false)
            }
            rsp => Err(Self::unexpected_response("Yes or No", &rsp).into()),
        }
    }

    pub async fn disable(&self, id: String) -> Result<bool> {
        let rsp = self.request(IpcRequest::Disable { id: id.clone() }).await?;
        match rsp {
            IpcResponse::Yes => {
                info!("disabled daemon {}", id);
                Ok(true)
            }
            IpcResponse::No => {
                info!("daemon {} already disabled", id);
                Ok(false)
            }
            rsp => Err(Self::unexpected_response("Yes or No", &rsp).into()),
        }
    }

    pub async fn run(&self, opts: RunOptions) -> Result<(Vec<String>, Option<i32>)> {
        info!("starting daemon {}", opts.id);
        let start_time = chrono::Local::now();
        // Use longer timeout for daemon start - ready_delay can be up to 60s+
        let timeout = Duration::from_secs(opts.ready_delay.unwrap_or(3) + 60);
        let rsp = self
            .request_with_timeout(IpcRequest::Run(opts.clone()), timeout)
            .await?;
        let mut started_daemons = vec![];
        let mut exit_code = None;
        match rsp {
            IpcResponse::DaemonStart { daemon } => {
                started_daemons.push(daemon.id.clone());
                info!("started {}", daemon.id);
            }
            IpcResponse::DaemonReady { daemon } => {
                started_daemons.push(daemon.id.clone());
                info!("started {}", daemon.id);
            }
            IpcResponse::DaemonFailedWithCode { exit_code: code } => {
                let code = code.unwrap_or(1);
                exit_code = Some(code);
                error!("daemon {} failed with exit code {}", opts.id, code);

                // Print logs from the time we started this specific daemon
                if let Err(e) =
                    crate::cli::logs::print_logs_for_time_range(&opts.id, start_time, None)
                {
                    error!("Failed to print logs: {}", e);
                }
            }
            IpcResponse::DaemonAlreadyRunning => {
                warn!("daemon {} already running", opts.id);
            }
            IpcResponse::DaemonFailed { error } => {
                error!("Failed to start daemon {}: {}", opts.id, error);
                exit_code = Some(1);

                // Print logs from the time we started this specific daemon
                if let Err(e) =
                    crate::cli::logs::print_logs_for_time_range(&opts.id, start_time, None)
                {
                    error!("Failed to print logs: {}", e);
                }
            }
            rsp => {
                return Err(Self::unexpected_response("DaemonStart or DaemonReady", &rsp).into());
            }
        }
        Ok((started_daemons, exit_code))
    }

    pub async fn active_daemons(&self) -> Result<Vec<Daemon>> {
        let rsp = self.request(IpcRequest::GetActiveDaemons).await?;
        match rsp {
            IpcResponse::ActiveDaemons(daemons) => Ok(daemons),
            rsp => Err(Self::unexpected_response("ActiveDaemons", &rsp).into()),
        }
    }

    pub async fn update_shell_dir(&self, shell_pid: u32, dir: PathBuf) -> Result<()> {
        let rsp = self
            .request(IpcRequest::UpdateShellDir {
                shell_pid,
                dir: dir.clone(),
            })
            .await?;
        match rsp {
            IpcResponse::Ok => {
                trace!("updated shell dir for pid {shell_pid} to {}", dir.display());
            }
            rsp => return Err(Self::unexpected_response("Ok", &rsp).into()),
        }
        Ok(())
    }

    pub async fn clean(&self) -> Result<()> {
        let rsp = self.request(IpcRequest::Clean).await?;
        match rsp {
            IpcResponse::Ok => {
                trace!("cleaned");
            }
            rsp => return Err(Self::unexpected_response("Ok", &rsp).into()),
        }
        Ok(())
    }

    pub async fn get_disabled_daemons(&self) -> Result<Vec<String>> {
        let rsp = self.request(IpcRequest::GetDisabledDaemons).await?;
        match rsp {
            IpcResponse::DisabledDaemons(daemons) => Ok(daemons),
            rsp => Err(Self::unexpected_response("DisabledDaemons", &rsp).into()),
        }
    }

    pub async fn get_notifications(&self) -> Result<Vec<(log::LevelFilter, String)>> {
        let rsp = self.request(IpcRequest::GetNotifications).await?;
        match rsp {
            IpcResponse::Notifications(notifications) => Ok(notifications),
            rsp => Err(Self::unexpected_response("Notifications", &rsp).into()),
        }
    }

    pub async fn stop(&self, id: String) -> Result<()> {
        let rsp = self.request(IpcRequest::Stop { id: id.clone() }).await?;
        match rsp {
            IpcResponse::Ok => {
                info!("stopped daemon {}", id);
                Ok(())
            }
            IpcResponse::DaemonAlreadyStopped => {
                warn!("daemon {} is not running", id);
                Ok(())
            }
            rsp => Err(Self::unexpected_response("Ok or DaemonAlreadyStopped", &rsp).into()),
        }
    }
}