magi_code/service/persistent/

entrypoint.rs

use super::{Coordinator, ServiceRuntime};
use anyhow::{Result, anyhow};
use crossbeam_channel::{Sender, bounded};
use std::{
    sync::Arc,
    thread,
    time::{Duration, Instant},
};

#[cfg(any(unix, test))]
use std::path::PathBuf;

/// Transport-independent persistent service. This does not listen on a socket or select v2
/// for stdio. A local adapter owns framing, write deadlines, and explicit disconnects.
/// Dropping this owner shuts down the daemon; disconnecting a connection does not.
/// A failed/timeout submission is uncertain: never automatically replay a mutation.
pub struct PersistentService {
    commands: Sender<Command>,
    worker: Option<thread::JoinHandle<()>>,
}

enum Action {
    #[cfg(any(unix, test))]
    StopIfIdle,
    Connect,
    Submit {
        connection: String,
        record: Vec<u8>,
    },
    Disconnect(String),
    Next(String),
    Written {
        connection: String,
        request: String,
    },
}
struct Command {
    action: Action,
    deadline: Instant,
    response: Sender<Result<Option<String>>>,
}

impl PersistentService {
    /// Load the normal runtime once and start its coordinator. No provider work runs until
    /// an initialized connection explicitly claims a durable session and starts a turn.
    pub fn start() -> Result<Self> {
        Self::start_with_loader(ServiceRuntime::load_persistent)
    }

    /// Start the Unix adapter with explicit absolute workspace and state paths.
    #[cfg(any(unix, test))]
    pub(crate) fn start_unix(workspace: PathBuf, state_root: PathBuf) -> Result<Self> {
        let start = || {
            let runtime = ServiceRuntime::load_unix(workspace, state_root)?;
            let mut coordinator = Coordinator::new(Arc::new(runtime))?;
            coordinator.unix_transport = true;
            Self::spawn(coordinator)
        };
        start().map_err(|_: anyhow::Error| anyhow!("application service startup failed"))
    }

    /// Stop only if all earlier admissions and worker cleanup have completed.
    /// Commands behind a successful stop are rejected, never accepted then cancelled.
    #[cfg(any(unix, test))]
    pub(crate) fn stop_if_idle(&self) -> Result<bool> {
        self.call(Action::StopIfIdle).map(|reply| reply.is_some())
    }

    #[cfg(any(unix, test))]
    pub(crate) fn is_finished(&self) -> bool {
        self.worker
            .as_ref()
            .is_none_or(|worker| worker.is_finished())
    }

    fn start_with_loader(load: impl FnOnce() -> Result<ServiceRuntime>) -> Result<Self> {
        let start = || Self::spawn(Coordinator::new(Arc::new(load()?))?);
        start().map_err(|_| anyhow!("application service startup failed"))
    }

    pub(super) fn spawn(coordinator: Coordinator) -> Result<Self> {
        Self::spawn_with_idle_clock(coordinator, Instant::now)
    }

    fn spawn_with_idle_clock(
        mut coordinator: Coordinator,
        mut idle_clock: impl FnMut() -> Instant + Send + 'static,
    ) -> Result<Self> {
        let (commands, receiver) = bounded::<Command>(32);
        let worker = thread::Builder::new()
            .name("magi-persistent-service".into())
            .spawn(move || {
                let started = Instant::now();
                let mut idle_since = started;
                loop {
                    // Only idle accounting uses this clock; request and worker deadlines
                    // remain on the real monotonic clock.
                    let now = idle_clock();
                    coordinator.tick(Instant::now());
                    let idle = coordinator.is_idle();
                    if !idle {
                        idle_since = now;
                    }
                    match receiver.recv_timeout(Duration::from_millis(5)) {
                        Ok(command) => {
                            #[cfg(any(unix, test))]
                            let stopping = matches!(command.action, Action::StopIfIdle);
                            #[cfg(not(any(unix, test)))]
                            let stopping = false;
                            let result = if Instant::now() >= command.deadline {
                                Err(anyhow!(super::Code::RequestTimeout.message()))
                            } else {
                                execute(&mut coordinator, command.action, command.deadline)
                            };
                            let stopped = stopping && matches!(&result, Ok(Some(_)));
                            let _ = command.response.try_send(result);
                            if stopped {
                                // No accepted work exists; do not enter owner-drop cancellation.
                                return;
                            }
                        }
                        Err(crossbeam_channel::RecvTimeoutError::Timeout) => {
                            // Check expiry only after giving queued admissions priority. A send
                            // racing this exit gets a closed response channel, not acceptance.
                            if idle
                                && now.duration_since(started) >= Duration::from_secs(10)
                                && now.duration_since(idle_since) >= Duration::from_secs(60)
                            {
                                return;
                            }
                        }
                        Err(crossbeam_channel::RecvTimeoutError::Disconnected) => break,
                    }
                }
                // Explicit daemon owner shutdown retains stdio's proven cancellation/draining
                // machinery. Ordinary connection loss never reaches this path.
                // Preparation has no effects until commit. Discard it, but keep draining
                // accepted turns before joining a possibly blocked preparation worker.
                let preparation = coordinator.turn_preparation.take();
                let auth_work = coordinator.auth_work.take();
                coordinator.execution.cancel_all();
                while !coordinator.cleaning.is_empty() {
                    coordinator.tick(Instant::now());
                    thread::sleep(Duration::from_millis(1));
                }
                while let Ok(Some(output)) = coordinator.execution.shutdown_next() {
                    if let Some(id) = output.terminal_turn_id {
                        coordinator.execution.finish_worker_output(&id);
                    }
                }
                if let Some(pending) = preparation {
                    let _ = pending.worker.join();
                }
                if let Some(work) = auth_work {
                    let _ = work.worker.join();
                }
            })?;
        Ok(Self {
            commands,
            worker: Some(worker),
        })
    }

    fn call(&self, action: Action) -> Result<Option<String>> {
        let (response, receiver) = bounded(1);
        self.commands
            .try_send(Command {
                action,
                deadline: Instant::now() + Duration::from_secs(30),
                response,
            })
            .map_err(|_| anyhow!("service admission queue unavailable"))?;
        receiver
            .recv_timeout(Duration::from_secs(30))
            .map_err(|_| anyhow!("service response unavailable; admission may have occurred"))?
    }

    /// Allocate a pending connection. Initialize it within five seconds.
    pub fn connect(&self) -> Result<String> {
        self.call(Action::Connect)?
            .ok_or_else(|| anyhow!("connection unavailable"))
    }

    /// Submit one complete, bounded JSON record. Responses are read with `next_record`.
    pub fn submit(&self, connection: &str, record: &[u8]) -> Result<()> {
        if record.len() > crate::service::protocol::MAX_RECORD_BYTES {
            return Err(anyhow!("record too large"));
        }
        self.call(Action::Submit {
            connection: checked_id(connection)?,
            record: record.to_vec(),
        })
        .map(|_| ())
    }

    /// Revoke this connection's control and cancel only its login, not accepted turns/settings.
    pub fn disconnect(&self, connection: &str) -> Result<()> {
        self.call(Action::Disconnect(checked_id(connection)?))
            .map(|_| ())
    }

    /// Remove one encoded record from the outbound queue. A response's request ID remains
    /// reserved until `response_written` or disconnect, even after dequeue.
    pub fn next_record(&self, connection: &str) -> Result<Option<String>> {
        self.call(Action::Next(checked_id(connection)?))
    }

    /// Release response correlation only after the adapter's write attempt. On write failure,
    /// disconnect instead. This acknowledgement never controls execution or lease cleanup.
    pub fn response_written(&self, connection: &str, request: &str) -> Result<()> {
        self.call(Action::Written {
            connection: checked_id(connection)?,
            request: checked_id(request)?,
        })
        .map(|_| ())
    }
}

fn checked_id(id: &str) -> Result<String> {
    anyhow::ensure!(super::wire::valid_id(id), "invalid service identity");
    Ok(id.to_owned())
}

fn execute(
    coordinator: &mut Coordinator,
    action: Action,
    deadline: Instant,
) -> Result<Option<String>> {
    let protocol_error = |code: super::Code| anyhow!(code.message());
    match action {
        #[cfg(any(unix, test))]
        Action::StopIfIdle => Ok(coordinator.is_idle().then(|| "stopped".to_owned())),
        Action::Connect => coordinator
            .connect(Instant::now())
            .map(Some)
            .map_err(protocol_error),
        Action::Submit { connection, record } => coordinator
            .submit_until(&connection, &record, Instant::now(), deadline)
            .map(|_| None)
            .map_err(protocol_error),
        Action::Disconnect(connection) => {
            coordinator.disconnect(&connection);
            Ok(None)
        }
        Action::Next(connection) => {
            let client = coordinator
                .connections
                .get_mut(&connection)
                .ok_or_else(|| protocol_error(super::Code::StaleConnection))?;
            let record = client.queue.pop_front();
            if let Some(record) = &record {
                client.queued_bytes -= record.len();
            }
            Ok(record)
        }
        Action::Written {
            connection,
            request,
        } => {
            let client = coordinator
                .connections
                .get_mut(&connection)
                .ok_or_else(|| protocol_error(super::Code::StaleConnection))?;
            if let Some(count) = client.requests.get_mut(&request) {
                *count -= 1;
                if *count == 0 {
                    client.requests.remove(&request);
                }
            }
            Ok(None)
        }
    }
}

impl Drop for PersistentService {
    fn drop(&mut self) {
        let (replacement, _receiver) = bounded(1);
        drop(std::mem::replace(&mut self.commands, replacement));
        if let Some(worker) = self.worker.take() {
            let _ = worker.join();
        }
    }
}

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

    // The clock gates each coordinator iteration, so queue order and expiry do not
    // depend on scheduling or on waiting sixty real seconds. Drop clock inputs first
    // to unblock the coordinator before PersistentService joins it.
    struct ControlledIdleService {
        clock: Sender<Instant>,
        ready: crossbeam_channel::Receiver<()>,
        service: PersistentService,
    }

    impl ControlledIdleService {
        fn new(coordinator: Coordinator) -> Self {
            let (clock, times) = bounded(1);
            let (ready_sender, ready) = bounded(1);
            let service = PersistentService::spawn_with_idle_clock(coordinator, move || {
                let _ = ready_sender.try_send(());
                times.recv().unwrap_or_else(|_| Instant::now())
            })
            .unwrap();
            ready.recv_timeout(Duration::from_secs(5)).unwrap();
            Self {
                clock,
                ready,
                service,
            }
        }

        fn queue(&self, action: Action) -> crossbeam_channel::Receiver<Result<Option<String>>> {
            let (response, reply) = bounded(1);
            self.service
                .commands
                .try_send(Command {
                    action,
                    deadline: Instant::now() + Duration::from_secs(30),
                    response,
                })
                .unwrap();
            reply
        }

        fn advance(&self, now: Instant) {
            self.clock.send(now).unwrap();
            self.ready.recv_timeout(Duration::from_secs(5)).unwrap();
            assert!(!self.service.is_finished());
        }
    }

    fn idle_coordinator(temp: &tempfile::TempDir) -> Coordinator {
        let runtime =
            ServiceRuntime::load_unix(temp.path().to_owned(), temp.path().join("state")).unwrap();
        Coordinator::new(Arc::new(runtime)).unwrap()
    }

    #[test]
    fn automatic_idle_expiry_prioritizes_queued_admission() {
        let temp = tempfile::tempdir().unwrap();
        let controlled = ControlledIdleService::new(idle_coordinator(&temp));
        let expired = Instant::now() + Duration::from_secs(61);
        let admitted = controlled.queue(Action::Connect);
        controlled.advance(expired);
        let connection = admitted
            .recv_timeout(Duration::from_secs(5))
            .unwrap()
            .unwrap()
            .unwrap();
        // A full timeout iteration after acceptance must not stop the service.
        controlled.advance(expired + Duration::from_secs(61));
        let disconnected = controlled.queue(Action::Disconnect(connection));
        controlled.advance(expired + Duration::from_secs(61));
        disconnected
            .recv_timeout(Duration::from_secs(5))
            .unwrap()
            .unwrap();
        controlled
            .clock
            .send(expired + Duration::from_secs(122))
            .unwrap();
        wait_for_completion(&controlled.service);
        assert!(controlled.service.connect().is_err());
    }

    #[test]
    fn automatic_idle_expiry_rejects_admission_after_exit() {
        let temp = tempfile::tempdir().unwrap();
        let controlled = ControlledIdleService::new(idle_coordinator(&temp));
        controlled
            .clock
            .send(Instant::now() + Duration::from_secs(61))
            .unwrap();
        wait_for_completion(&controlled.service);
        assert!(controlled.service.connect().is_err());
    }

    #[test]
    fn automatic_idle_expiry_waits_for_disconnected_accepted_settings_write() {
        let temp = tempfile::tempdir().unwrap();
        let mut coordinator = idle_coordinator(&temp);
        let connection = coordinator.connect(Instant::now()).unwrap();
        coordinator.submit(&connection, br#"{"protocol_version":2,"kind":"request","request_id":"init","instance_id":null,"connection_id":null,"session_id":null,"operation_id":null,"control":null,"method":"initialize","payload":{"supported_protocol_versions":[2]}}"#, Instant::now()).unwrap();
        let initialized: serde_json::Value =
            serde_json::from_str(coordinator.connections[&connection].queue.front().unwrap())
                .unwrap();
        assert!(initialized["error"].is_null(), "{initialized}");
        let paths = coordinator.runtime.config.paths.clone();
        // Declare the service before the lock so unwinding releases the worker first.
        let controlled;
        let lock = crate::persistence::CrossProcessFileLock::acquire(&paths.settings_file).unwrap();
        let request = serde_json::json!({
            "protocol_version":2,"kind":"request","request_id":"write",
            "instance_id":coordinator.instance,"connection_id":connection,
            "session_id":null,"control":null,
            "operation_id":"settings-write","method":"config.set",
            "payload":{"scope":"global","fast":true}
        });
        coordinator
            .submit(
                &connection,
                &serde_json::to_vec(&request).unwrap(),
                Instant::now(),
            )
            .unwrap();
        assert_eq!(
            coordinator.operations.lookup(
                &coordinator.instance,
                &coordinator.instance,
                "settings-write"
            )["state"],
            "accepted"
        );
        coordinator.disconnect(&connection);
        controlled = ControlledIdleService::new(coordinator);
        let mut now = Instant::now() + Duration::from_secs(61);
        controlled.advance(now);
        now += Duration::from_secs(61);
        controlled.advance(now);
        assert!(!crate::config::read_settings(&paths).unwrap().fast.enabled);
        drop(lock);
        // Let the real worker finish. Virtual idle periods alone cannot bypass it.
        let deadline = Instant::now() + Duration::from_secs(5);
        loop {
            assert!(Instant::now() < deadline, "settings worker did not settle");
            now += Duration::from_secs(61);
            controlled.clock.send(now).unwrap();
            match controlled.ready.recv_timeout(Duration::from_secs(5)) {
                Ok(()) => {}
                Err(crossbeam_channel::RecvTimeoutError::Disconnected) => break,
                Err(error) => panic!("coordinator did not progress: {error}"),
            }
        }
        wait_for_completion(&controlled.service);
        assert!(crate::config::read_settings(&paths).unwrap().fast.enabled);
        assert!(controlled.service.connect().is_err());
    }

    fn wait_for_completion(service: &PersistentService) {
        let deadline = Instant::now() + Duration::from_secs(5);
        while !service.is_finished() {
            assert!(Instant::now() < deadline, "coordinator did not exit");
            thread::sleep(Duration::from_millis(1));
        }
    }

    #[test]
    fn unix_service_refuses_busy_stop_then_exits_without_accepting_more_connections() {
        let temp = tempfile::tempdir().unwrap();
        let service =
            PersistentService::start_unix(temp.path().to_owned(), temp.path().join("state"))
                .unwrap();
        let connection = service.connect().unwrap();
        assert!(!service.stop_if_idle().unwrap());
        assert!(!service.is_finished());
        service.submit(&connection, br#"{"protocol_version":2,"kind":"request","request_id":"init","instance_id":null,"connection_id":null,"session_id":null,"operation_id":null,"control":null,"method":"initialize","payload":{"supported_protocol_versions":[2]}}"#).unwrap();
        let record: serde_json::Value =
            serde_json::from_str(&service.next_record(&connection).unwrap().unwrap()).unwrap();
        assert!(record["error"].is_null(), "{record}");
        assert_eq!(
            record["payload"]["capabilities"]["transports"],
            serde_json::json!(["unix"])
        );
        service.disconnect(&connection).unwrap();
        assert!(service.stop_if_idle().unwrap());
        wait_for_completion(&service);
        assert!(service.connect().is_err());
    }

    #[test]
    fn preparation_blocks_stop_until_worker_completion() {
        let temp = tempfile::tempdir().unwrap();
        let runtime = Arc::new(
            ServiceRuntime::load_unix(temp.path().to_owned(), temp.path().join("state")).unwrap(),
        );
        let mut coordinator = Coordinator::new(Arc::clone(&runtime)).unwrap();
        let (release, released) = bounded(1);
        coordinator.turn_preparation = Some(super::super::PendingTurnPreparation {
            connection: "disconnected".into(),
            request: serde_json::from_value(serde_json::json!({
                "protocol_version":2,"kind":"request","request_id":"prepare",
                "method":"turn.start","payload":{}
            }))
            .unwrap(),
            deadline: Instant::now() + Duration::from_secs(30),
            worker: thread::spawn(move || {
                released.recv().unwrap();
                Ok(runtime)
            }),
        });
        let service = PersistentService::spawn(coordinator).unwrap();
        assert!(!service.stop_if_idle().unwrap());
        release.send(()).unwrap();
        let deadline = Instant::now() + Duration::from_secs(5);
        while !service.stop_if_idle().unwrap() {
            assert!(Instant::now() < deadline, "preparation did not settle");
            thread::sleep(Duration::from_millis(1));
        }
        wait_for_completion(&service);
    }

    #[test]
    fn queued_stop_rejects_later_admissions() {
        let temp = tempfile::tempdir().unwrap();
        let runtime =
            ServiceRuntime::load_unix(temp.path().to_owned(), temp.path().join("state")).unwrap();
        let coordinator = Coordinator::new(Arc::new(runtime)).unwrap();
        assert_eq!(
            super::super::capabilities(coordinator.unix_transport)["transports"],
            serde_json::json!([])
        );
        let service = PersistentService::spawn(coordinator).unwrap();
        let (response, stopped) = bounded(1);
        service
            .commands
            .try_send(Command {
                action: Action::StopIfIdle,
                deadline: Instant::now() + Duration::from_secs(30),
                response,
            })
            .unwrap();
        // This admission is behind the stop even if the worker has not processed it yet.
        assert!(service.connect().is_err());
        assert!(
            stopped
                .recv_timeout(Duration::from_secs(5))
                .unwrap()
                .unwrap()
                .is_some()
        );
        wait_for_completion(&service);
    }

    #[test]
    fn startup_errors_do_not_expose_settings_or_auth_secrets() {
        for marker in ["settings-secret-marker", "auth-secret-marker"] {
            let error = PersistentService::start_with_loader(|| Err(anyhow!(marker)))
                .err()
                .expect("startup must fail");
            assert_eq!(error.to_string(), "application service startup failed");
            assert!(!format!("{error:?}").contains(marker));
        }
    }
}