agentmux_srv/backend/blockcontroller/persistent/

lifecycle.rs

// Copyright 2026, AgentMux Corp.
// SPDX-License-Identifier: Apache-2.0

//! Stop/restart and session-identity queries: `stop_process`, the kill request
//! path, `session_id`, `needs_spawn`, and the deferred-restart flag.

use super::*;

impl PersistentSubprocessController {
    pub fn stop_process(&self, force: bool) -> Result<(), String> {
        let request = if force {
            KillRequest::Force
        } else {
            KillRequest::Graceful(std::time::Instant::now() + super::super::SHUTDOWN_GRACE)
        };
        self.request_stop(request)
    }

    /// Another process already running on session `sid`: `(block id, closing)`.
    /// Checks live persistent controllers in the registry, then blocks
    /// mid-close (out of the registry, process not yet exited).
    pub(super) fn session_held_elsewhere(&self, sid: &str) -> Option<(String, bool)> {
        for (block_id, ctrl) in super::super::get_all_controllers() {
            if block_id == self.block_id {
                continue;
            }
            if let Some(other) = ctrl.as_any().downcast_ref::<PersistentSubprocessController>() {
                let g = other.inner.lock().unwrap();
                if g.current_pid.is_some() && g.session_id.as_deref() == Some(sid) {
                    return Some((block_id, false));
                }
            }
        }
        let store = self.mstore.as_ref()?;
        super::super::closing_blocks_still_running()
            .into_iter()
            .filter(|block_id| *block_id != self.block_id)
            .find(|block_id| {
                store
                    .get::<crate::backend::obj::Block>(block_id)
                    .ok()
                    .flatten()
                    .is_some_and(|b| crate::backend::obj::meta_get_string(&b.meta, core::META_SESSION_ID, "") == sid)
            })
            .map(|block_id| (block_id, true))
    }

    /// Another pane of this process holding agent `uid`'s single-live-instance
    /// lease: its block id. Whatever session either pane holds — a pane with
    /// a stale id refused because the agent's live pane holds the lease is
    /// still "open in another pane", not "another instance" (Camper on #3935).
    pub(super) fn agent_held_by_other_pane(&self, uid: &str) -> Option<String> {
        super::super::get_all_controllers().into_iter().find_map(|(block_id, ctrl)| {
            if block_id == self.block_id {
                return None;
            }
            let other = ctrl.as_any().downcast_ref::<PersistentSubprocessController>()?;
            let holds = other.inner.lock().unwrap().agent_lease.as_ref().is_some_and(|l| l.agent_uid() == uid);
            holds.then_some(block_id)
        })
    }

    pub(super) fn request_stop(&self, request: KillRequest) -> Result<(), String> {
        Self::request_stop_on(&self.inner, request)
    }

    /// [`request_stop`] over just the shared state, for the `'static`
    /// future [`Controller::shutdown`] returns.
    pub(super) fn request_stop_on(inner_arc: &Arc<Mutex<PersistentInner>>, request: KillRequest) -> Result<(), String> {
        Self::request_stop_inner(inner_arc, request, false, false);
        Ok(())
    }

    /// Apply a deferred runtime-config restart: stop the current process (as
    /// `stop_process(false)` does) WITHOUT clearing the restart token the turn
    /// boundary set (`PersistentInner::config_restart_generation`), so the exit
    /// handler brings the replacement up. It never sets the token either: a
    /// Stop that landed between the boundary and this call already cleared it
    /// and took the kill channel, and must stay in force (codex P1 on #3990).
    pub(super) fn stop_for_config_restart(&self) {
        let request = KillRequest::Graceful(std::time::Instant::now() + super::super::SHUTDOWN_GRACE);
        Self::request_stop_inner(&self.inner, request, false, true);
    }

    /// A kill for teardown: drains the deferred queue in the SAME `inner`
    /// acquisition that issues the kill, and returns what it drained for the
    /// caller to report. Draining in a second acquisition left a window
    /// after the kill request where a flush could write a deferred message
    /// into the process being killed, and leave the report an empty queue:
    /// lost, with no stranded warning (codex P2 on #3562). Once this lock is
    /// released, `stop_pending` holds anything queued later.
    pub(super) fn request_stop_draining_deferred(&self, request: KillRequest) -> Vec<String> {
        Self::request_stop_inner(&self.inner, request, true, false)
    }

    /// `config_restart`: this stop is a deferred runtime-config restart's own
    /// ([`Self::stop_for_config_restart`]), which leaves the restart token as
    /// the boundary left it. Every other stop clears it in this same
    /// acquisition, so a Stop or pane close landing after the restart was
    /// committed wins, and nothing respawns behind it.
    pub(super) fn request_stop_inner(
        inner_arc: &Arc<Mutex<PersistentInner>>,
        request: KillRequest,
        drain_deferred: bool,
        config_restart: bool,
    ) -> Vec<String> {
        let (kill_tx, drained) = {
            let mut inner = inner_arc.lock().unwrap();
            // Recorded unconditionally, not only when `kill_tx` is already
            // `None` — codex P1 on PR #2360 (round 16, commit ce1642d90):
            // `stop_process` can race a process that already exited (a
            // confirmed stale-`--resume` death) and is about to be
            // silently retried — sending through `kill_tx` is futile in
            // that window regardless of whether it's already `None` (the
            // exit-handler cleared it) or still `Some` (`tokio::select!`
            // already committed to the `child.wait()` exit arm before
            // this call reached the lock, so the `kill_rx` arm will never
            // be polled again even if the send succeeds). Recording this
            // as a `StopRequested` event — resolved by
            // `persistent_resume::update` once `ProcessExited` arrives —
            // is the only way the exit-handler's retry decision can know
            // the user explicitly asked to stop.
            let generation = inner.spawn_generation;
            inner.apply_resume_event(persistent_resume::ResumeEvent::StopRequested { generation });
            inner.stop_pending = true;
            if !config_restart {
                inner.config_restart_generation = None;
            }
            let drained = if drain_deferred {
                inner.deferred_deliveries.drain(..).collect()
            } else {
                Vec::new()
            };
            (inner.kill_tx.take(), drained)
        };
        if let Some(tx) = kill_tx {
            let _ = tx.send(request);
        }
        drained
    }

    pub fn session_id(&self) -> Option<String> {
        self.inner.lock().unwrap().session_id.clone()
    }

    /// True when this controller is registered but has **no live process and no
    /// spawn already in flight** — the state a lazily-registered controller sits
    /// in between `register` ("spawns on first message") and its first
    /// `send_message`.
    ///
    /// In that state `inject_message` cannot deliver: it has no spawn config and
    /// only writes to a live `stdin_tx`, so it returns "persistent process not
    /// running". The caller can, however, start a *turn* — `run_agent_turn`
    /// re-reads the spawn config from block metadata and calls `send_message`,
    /// which does spawn. This predicate is what lets the reactive delivery path
    /// tell those two situations apart. See
    /// `docs/reports/REPORT_JEKT_DELIVERY_DROPS_UNSPAWNED_PERSISTENT_AGENTS_2026_09_03.md`.
    ///
    /// **The invariant this must uphold:** `needs_spawn() == true` implies a
    /// subsequent `send_message` takes `decide_send_action`'s `BecomeSpawner`
    /// branch — i.e. it really spawns and delivers. It must never be true in a
    /// state where `send_message` would instead return `SendAction::Queued`,
    /// because `Queued` returns `Ok(())` having delivered nothing, the reactive
    /// path maps that to `Ok(true)`, and the caller is told the message landed
    /// when it is still sitting in the queue. `cloud_subscriber` only retries
    /// on `!success`, so a false success is a permanently lost message.
    ///
    /// That is why all three exclusions are here, and each mirrors one of
    /// `decide_send_action`'s own guards:
    ///
    /// - `stdin_tx.is_none()` — a live process is steerable; `inject_message`
    ///   handles it and no turn start is wanted.
    /// - `!spawning_in_progress` — a caller has already claimed this spawn
    ///   round (see that field's doc comment on the concurrent-spawn TOCTOU
    ///   race). Reporting "needs spawn" here would invite a second, racing
    ///   spawn and orphan a child.
    /// - `!drain_claim` — a `RetryFlush` drain owns the round. This one is
    ///   subtle and was missed on the first cut (reagent P1 on PR #2960): when
    ///   a drain's target process dies mid-flush, the drain **deliberately
    ///   retains** its claim for the fallback respawn while the exit handler
    ///   clears `stdin_tx`. Without this term, that window reports "needs
    ///   spawn", `decide_send_action` then returns `Queued` on the
    ///   still-held claim, and the message is silently queued while the caller
    ///   is told it was delivered.
    ///
    /// Both excluded windows are retryable rather than lost: the caller gets
    /// the original delivery error back, which is what lets it come again.
    ///
    /// Racy by nature, and safe to be: the process can exit the instant after
    /// this returns `false`. It is a routing hint, not a guarantee — the spawn
    /// claim inside `send_message` is what actually serialises spawners. What
    /// it must not do is report `true` for a state that cannot spawn.
    pub fn needs_spawn(&self) -> bool {
        let inner = self.inner.lock().unwrap();
        inner.stdin_tx.is_none() && !inner.spawning_in_progress && !inner.drain_claim
    }

    /// Ask this controller to restart itself once the current turn ends,
    /// instead of being torn down and replaced right now.
    ///
    /// Returns `true` when the restart was deferred (a turn is in flight, so
    /// the caller must NOT replace the controller), `false` when the pane is
    /// idle and an immediate replace is safe — the caller then proceeds
    /// exactly as before.
    ///
    /// Called from `resync_controller`'s forced-replace path. See
    /// [`PersistentInner::restart_when_idle`] for what the old
    /// kill-immediately behaviour destroyed.
    ///
    /// Deliberately keyed on the health monitor's `is_active_turn` rather
    /// than on `stdin_tx.is_some()`: a live process between turns is idle and
    /// should be replaced immediately (that's the common `/model` case, and
    /// deferring it would leave the change unapplied until the next turn
    /// happened to end). It's specifically an IN-FLIGHT turn that must not be
    /// interrupted.
    pub fn request_restart_when_idle(&self) -> bool {
        // The turn-active read and the flag write happen under ONE acquisition
        // of `inner`, and the turn-end consumer clears `active_turn` under that
        // same lock — so the two serialize (codex P2 on PR #2858). Interleaved,
        // this would otherwise observe an active turn, have the consumer run to
        // completion (seeing `restart_when_idle` still false, so restarting
        // nothing), and then set the flag — leaving it stuck until the NEXT
        // turn ended, so the user's next prompt ran with the old config.
        //
        // Lock order is `inner` -> health monitor here and in the consumer;
        // `TurnActivityTracker` never holds its own lock across an `inner`
        // acquisition, so the order can't invert.
        let mut inner = self.inner.lock().unwrap();
        if !self.health_monitor.is_active_turn() {
            return false;
        }
        inner.restart_when_idle = true;
        drop(inner);
        // Tell the menu a restart is coming, so a selection that has not
        // reached the process yet reads as pending and not as a mismatch.
        self.publish_status();
        tracing::info!(
            block_id = %self.block_id,
            "runtime-config change arrived mid-turn — deferring the restart to the end of this turn \
             instead of killing the in-flight message"
        );
        true
    }
}

/// The MCP self-heal's steps for this pane (`backend::mcp_health` decides
/// them; docs/specs/SPEC_AGENT_MCP_SELF_HEAL_2026_10_10.md).
impl PersistentSubprocessController {
    /// A frame from process `generation`: the `agentmux` MCP server's status
    /// (`system/init`, our `mcp_status`) or a reconnect's outcome. Anything
    /// else, or a frame from a replaced process, is ignored.
    pub(super) fn observe_mcp_frame(&self, frame: &serde_json::Value, generation: u64) {
        use crate::backend::mcp_health as heal;
        if self.inner.lock().unwrap().spawn_generation != generation {
            return;
        }
        let process = self.mcp_process_key(generation);
        let now = agentmux_common::time::now_ms_u64();
        let report = heal::status_from_init(frame)
            .map(|status| (status, None))
            .or_else(|| heal::status_from_control_response(frame));
        let decision = if let Some((status, error)) = report {
            heal::on_status(&self.block_id, &process, status, error, heal::auto_enabled(), now)
        } else if let Some(result) = heal::reconnect_result(frame) {
            heal::on_reconnect(&self.block_id, &process, result, now)
        } else {
            return;
        };
        self.carry_out_mcp_step(decision, generation);
    }

    /// `agentmux-mcp doctor --heal`: start the repair over now. `Err` says why
    /// nothing started (`mcp_health::heal_now`).
    pub(crate) fn mcp_heal_now(&self) -> Result<(), &'static str> {
        let generation = self.inner.lock().unwrap().spawn_generation;
        let now = agentmux_common::time::now_ms_u64();
        let decision = crate::backend::mcp_health::heal_now(&self.block_id, &self.mcp_process_key(generation), now)?;
        self.carry_out_mcp_step(decision, generation);
        Ok(())
    }

    /// One CLI process of this controller: a replacement controller restarts
    /// its generations, so the generation alone could name an older process.
    fn mcp_process_key(&self, generation: u64) -> String {
        format!("{:p}-{generation}", Arc::as_ptr(&self.inner))
    }

    fn carry_out_mcp_step(&self, decision: crate::backend::mcp_health::Decision, generation: u64) {
        use crate::backend::mcp_health::{self as heal, Action};
        if let Some(notice) = &decision.notice {
            self.show_mcp_notice(notice);
        }
        match decision.action {
            Action::None => {}
            Action::Reconnect { after } => self.send_mcp_request(heal::reconnect_request_line(), after, generation),
            Action::CheckStatus { after } => self.send_mcp_request(heal::status_request_line(), after, generation),
            Action::SafeProfileRestart => self.restart_now_or_after_turn(),
            Action::GiveUp => self.tell_agent_mcp_is_down(),
        }
    }

    /// Write a control request to process `generation`, after `after`; not if
    /// it has been replaced by then.
    fn send_mcp_request(&self, line: String, after: std::time::Duration, generation: u64) {
        let inner = Arc::clone(&self.inner);
        let send = move || {
            if inner.lock().unwrap().spawn_generation == generation {
                Self::push_stdin(&inner, line);
            }
        };
        if after.is_zero() {
            send();
        } else if let Ok(rt) = tokio::runtime::Handle::try_current() {
            rt.spawn(async move {
                tokio::time::sleep(after).await;
                send();
            });
        }
    }

    /// Restart the CLI (the same session) with the config now in force, which
    /// for a pane in the safe profile means without the shell prefix: at once
    /// when idle, else when the turn ends, through the deferred runtime-config
    /// restart (`turn_boundary_locked`, `respawn_after_config_restart`). The
    /// turn check and the flag share one acquisition, as in
    /// [`Self::request_restart_when_idle`].
    fn restart_now_or_after_turn(&self) {
        let mut inner = self.inner.lock().unwrap();
        if inner.stdin_tx.is_none() || inner.stop_pending {
            return;
        }
        if self.health_monitor.is_active_turn() {
            inner.restart_when_idle = true;
            drop(inner);
            self.publish_status();
            return;
        }
        inner.restart_pending = true;
        inner.config_restart_generation = Some(inner.spawn_generation);
        drop(inner);
        tracing::info!(block_id = %self.block_id, "MCP self-heal: restarting the agent in the safe profile");
        self.stop_for_config_restart();
    }

    fn show_mcp_notice(&self, notice: &crate::backend::mcp_health::Notice) {
        tracing::info!(block_id = %self.block_id, state = notice.state, step = notice.step, error = ?notice.error, "MCP self-heal");
        let (Some(broker), Some(filestore), Some(mstore)) = (&self.broker, &self.filestore, &self.mstore) else {
            return;
        };
        let (broker, filestore, mstore) = (Arc::clone(broker), Arc::clone(filestore), Arc::clone(mstore));
        let (block_id, frame) = (self.block_id.clone(), crate::backend::mcp_health::notice_frame(notice));
        let append = move || crate::backend::cli_notice::append_to_pane(&broker, &filestore, &mstore, &block_id, &frame);
        match tokio::runtime::Handle::try_current() {
            Ok(rt) => drop(rt.spawn_blocking(append)),
            Err(_) => append(),
        }
    }

    /// The `[AgentMux]` note that AgentMux gave up, so the agent tells the
    /// operator and knows its tools are gone. By UID when the pane has one.
    fn tell_agent_mcp_is_down(&self) {
        let error = crate::backend::mcp_health::snapshot(&self.block_id).and_then(|s| s.last_error);
        let text = crate::backend::mcp_health::give_up_note(error.as_deref());
        let handler = crate::backend::reactive::get_global_handler();
        let Some(target) = handler
            .list_agents()
            .into_iter()
            .find(|r| r.block_id == self.block_id)
            .map(|r| r.uid.unwrap_or(r.agent_id))
        else {
            tracing::info!(block_id = %self.block_id, "MCP self-heal: no registered agent to tell");
            return;
        };
        if let Err(e) = handler.deliver_system_note(&target, &text) {
            tracing::info!(block_id = %self.block_id, error = %e, "MCP self-heal note not delivered");
        }
    }
}