agent_procs/daemon/

server.rs

use crate::daemon::wait_engine;
use crate::protocol::{Request, Response};
use std::path::Path;
use std::sync::Arc;
use std::time::Duration;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::net::UnixListener;
use tokio::sync::broadcast;
use tokio::sync::Mutex;

use super::process_manager::ProcessManager;

pub struct DaemonState {
    pub process_manager: ProcessManager,
}

pub async fn run(session: &str, socket_path: &Path) {
    let state = Arc::new(Mutex::new(DaemonState {
        process_manager: ProcessManager::new(session),
    }));

    let listener = UnixListener::bind(socket_path).expect("failed to bind socket");

    loop {
        let (stream, _) = match listener.accept().await {
            Ok(conn) => conn,
            Err(_) => break,
        };

        let state = Arc::clone(&state);
        tokio::spawn(async move {
            let (reader, writer) = stream.into_split();
            let writer = Arc::new(Mutex::new(writer));
            let mut lines = BufReader::new(reader).lines();

            while let Ok(Some(line)) = lines.next_line().await {
                let request: Request = match serde_json::from_str(&line) {
                    Ok(r) => r,
                    Err(e) => {
                        let resp = Response::Error { code: 1, message: format!("invalid request: {}", e) };
                        let _ = send_response(&writer, &resp).await;
                        continue;
                    }
                };

                // Handle follow requests with streaming (before handle_request)
                if let Request::Logs { follow: true, ref target, all, timeout_secs, lines, .. } = request {
                    let output_rx = state.lock().await.process_manager.output_tx.subscribe();
                    let max_lines = lines;
                    let target_filter = target.clone();
                    let show_all = all;

                    handle_follow_stream(
                        &writer, output_rx, target_filter, show_all, timeout_secs, max_lines
                    ).await;
                    continue; // Don't call handle_request
                }

                let is_shutdown = matches!(request, Request::Shutdown);

                let response = handle_request(&state, request).await;
                let _ = send_response(&writer, &response).await;

                if is_shutdown {
                    tokio::time::sleep(Duration::from_millis(100)).await;
                    std::process::exit(0);
                }
            }
        });
    }
}

async fn handle_follow_stream(
    writer: &Arc<Mutex<tokio::net::unix::OwnedWriteHalf>>,
    mut output_rx: broadcast::Receiver<super::log_writer::OutputLine>,
    target: Option<String>,
    all: bool,
    timeout_secs: Option<u64>,
    max_lines: Option<usize>,
) {
    let mut line_count: usize = 0;

    let stream_loop = async {
        loop {
            match output_rx.recv().await {
                Ok(output_line) => {
                    if !all {
                        if let Some(ref t) = target {
                            if output_line.process != *t { continue; }
                        }
                    }

                    let resp = Response::LogLine {
                        process: output_line.process,
                        stream: output_line.stream,
                        line: output_line.line,
                    };
                    if send_response(writer, &resp).await.is_err() {
                        return;
                    }

                    line_count += 1;
                    if let Some(max) = max_lines {
                        if line_count >= max { return; }
                    }
                }
                Err(broadcast::error::RecvError::Lagged(_)) => continue,
                Err(broadcast::error::RecvError::Closed) => return,
            }
        }
    };

    // Apply timeout only if specified; otherwise stream indefinitely
    match timeout_secs {
        Some(secs) => { let _ = tokio::time::timeout(Duration::from_secs(secs), stream_loop).await; }
        None => { stream_loop.await; }
    }

    let _ = send_response(writer, &Response::LogEnd).await;
}

async fn handle_request(state: &Arc<Mutex<DaemonState>>, request: Request) -> Response {
    match request {
        Request::Run { command, name, cwd, env } => {
            state.lock().await.process_manager.spawn_process(&command, name, cwd.as_deref(), env.as_ref()).await
        }
        Request::Stop { target } => {
            state.lock().await.process_manager.stop_process(&target).await
        }
        Request::StopAll => {
            state.lock().await.process_manager.stop_all().await
        }
        Request::Restart { target } => {
            state.lock().await.process_manager.restart_process(&target).await
        }
        Request::Status => {
            state.lock().await.process_manager.status()
        }
        Request::Wait { target, until, regex, exit, timeout_secs } => {
            // Check process exists
            {
                let s = state.lock().await;
                if !s.process_manager.has_process(&target) {
                    return Response::Error { code: 2, message: format!("process not found: {}", target) };
                }
            }
            // Subscribe to output and delegate to wait engine
            let output_rx = state.lock().await.process_manager.output_tx.subscribe();
            let timeout = Duration::from_secs(timeout_secs.unwrap_or(30));
            let state_clone = Arc::clone(state);
            let target_clone = target.clone();
            wait_engine::wait_for(
                output_rx,
                &target,
                until.as_deref(),
                regex,
                exit,
                timeout,
                move || {
                    // This is called synchronously from the wait loop
                    // We can't hold the lock across the whole wait, so we check briefly
                    let state = state_clone.clone();
                    let target = target_clone.clone();
                    // Use try_lock to avoid deadlock
                    let result = match state.try_lock() {
                        Ok(mut s) => s.process_manager.is_process_exited(&target),
                        Err(_) => None,
                    };
                    result
                },
            ).await
        }
        Request::Logs { follow: false, .. } => {
            // Non-follow logs are read directly from files by the CLI — no daemon involvement needed
            Response::Error { code: 1, message: "non-follow logs are read directly from disk by CLI".into() }
        }
        Request::Logs { follow: true, .. } => {
            // Handled separately in connection loop (needs streaming)
            Response::Error { code: 1, message: "follow requests handled in connection loop".into() }
        }
        Request::Shutdown => {
            state.lock().await.process_manager.stop_all().await;
            Response::Ok { message: "daemon shutting down".into() }
        }
    }
}

async fn send_response(writer: &Arc<Mutex<tokio::net::unix::OwnedWriteHalf>>, response: &Response) -> std::io::Result<()> {
    let mut w = writer.lock().await;
    let mut json = serde_json::to_string(response).unwrap();
    json.push('\n');
    w.write_all(json.as_bytes()).await?;
    w.flush().await
}