hatchery_node/server/

node_wake.rs

//! Letter -> node -> g4a wake -> session turn.
//!
//! m4a only hands a letter over (its `M4A_LETTER_HOOK`); the node knows which of its sessions
//! owns the addressed nick (it assigned that account at spawn) and asks g4a
//! (`gate4agent_wake::hosted`) what to do. A live session gets the framed prompt as its next
//! turn; the node's policy ([`crate::wake_pipeline`]) gates everything else: typing into a raw
//! terminal needs an explicit terminal grant, a nick bound to a routine gets one routine POST via
//! g4a, and anything else is refused (m4a keeps the letter in its durable inbox).
//! Callers present the node's wake token (a 0600 file next to the accounts root), so another
//! local user cannot type into a session through the door.

use std::collections::BTreeMap;
use std::sync::Arc;

use gate4agent_wake::hosted::{plan_hosted_wake, HostedSessions, HostedTarget, HostedWake};
use gate4agent_wake::{ProviderKind, WakeLetter};

use crate::wake_pipeline::{deliver_routine, launch_spawn_profile, route, Route, RoutineLetter, SessionShape};

use super::*;

/// One letter for a session of this node.
pub struct NodeLetter {
    pub nick: String,
    pub from: String,
    pub event_id: String,
    pub body: String,
    /// The node spawned this session itself: a raw terminal gets the short wake line typed by
    /// the node (no terminal grant needed; the line carries no letter text).
    pub own_spawn: bool,
}

/// A letter for a busy live session is retried this often, at most this many times (~30 min).
const BUSY_RETRY: std::time::Duration = std::time::Duration::from_secs(10);
const BUSY_RETRY_MAX: u32 = 180;

#[derive(Clone)]
pub struct NodeWake {
    shared: Arc<NodeShared>,
    token: Arc<str>,
    policy: Arc<crate::wake_pipeline::WakePolicy>,
}

struct Snapshot(BTreeMap<String, HostedTarget>);

impl HostedSessions for Snapshot {
    fn lookup(&self, nick: &str) -> Option<HostedTarget> {
        self.0.get(nick).cloned()
    }
}

fn kind_of(agent: &AgentId) -> Option<ProviderKind> {
    ProviderKind::parse(crate::m4a_accounts::m4a_provider(agent.as_str())?)
}

impl NodeWake {
    pub(super) fn new(shared: Arc<NodeShared>, token: String, policy: Arc<crate::wake_pipeline::WakePolicy>) -> Self {
        Self { shared, token: token.into(), policy }
    }

    /// Constant-time check of the caller's token.
    pub fn authorized(&self, token: &str) -> bool {
        let (a, b) = (self.token.as_bytes(), token.as_bytes());
        a.len() == b.len() && a.iter().zip(b).fold(0u8, |acc, (x, y)| acc | (x ^ y)) == 0
    }

    fn snapshot(&self) -> (Snapshot, BTreeMap<String, SessionAddress>) {
        let live = self.shared.m4a_live.lock().unwrap_or_else(|p| p.into_inner()).clone();
        let bindings = self.shared.session_bindings.lock().unwrap_or_else(|p| p.into_inner());
        let mut targets = BTreeMap::new();
        let mut addresses = BTreeMap::new();
        for (nick, (instance_id, agent)) in live {
            let Some(provider) = kind_of(&agent) else { continue };
            let binding = bindings.get(&instance_id);
            if let Some(b) = binding {
                addresses.insert(
                    nick.clone(),
                    SessionAddress { workspace_id: b.workspace_id.clone(), session: SessionKey { instance_id, generation: b.generation } },
                );
            }
            let spawn_allowed = self.policy.spawn_profile(&nick).is_some();
            targets.insert(nick.clone(), HostedTarget { nick, provider, live: binding.is_some(), spawn_allowed });
        }
        (Snapshot(targets), addresses)
    }

    async fn resume_with_wake(&self, nick: &str, address: &SessionAddress, prompt: String) -> Result<&'static str, String> {
        let record_id = self
            .shared
            .session_bindings
            .lock()
            .unwrap_or_else(|p| p.into_inner())
            .get(&address.session.instance_id)
            .and_then(|b| b.record_id.clone())
            .ok_or("exited session has no record to resume")?;
        let size = gate4agent_types::TerminalSize { rows: 40, columns: 120 };
        let prompt_retry = prompt.clone();
        match self.shared.resume_session_record(&record_id, size, Some(prompt)).await {
            Ok(_) => {
                tracing::info!(%nick, "letter resumed an exited session");
                Ok("resumed")
            }
            Err(e) => {
                tracing::warn!(%nick, code = ?e.code, "exited session not resumed");
                // The record is still held (SessionRecordBusy: the old process is going away or a
                // live session owns it): try the resume again later instead of losing the letter.
                if format!("{:?}", e.code).contains("SessionRecordBusy") {
                    self.queue_resume(nick.to_owned(), record_id, prompt_retry);
                    return Ok("queued");
                }
                Err("exited session could not be resumed".into())
            }
        }
    }

    /// Retries resuming a held record with the wake prompt, every [`BUSY_RETRY`] up to
    /// [`BUSY_RETRY_MAX`] tries.
    fn queue_resume(&self, nick: String, record_id: super::SessionRecordId, prompt: String) {
        tracing::info!(%nick, "letter queued: session record busy, resume retried later");
        let this = self.clone();
        tokio::spawn(async move {
            for attempt in 1..=BUSY_RETRY_MAX {
                tokio::time::sleep(BUSY_RETRY).await;
                let size = gate4agent_types::TerminalSize { rows: 40, columns: 120 };
                match this.shared.resume_session_record(&record_id, size, Some(prompt.clone())).await {
                    Ok(_) => {
                        tracing::info!(%nick, attempt, "queued letter resumed an exited session");
                        return;
                    }
                    Err(e) if format!("{:?}", e.code).contains("SessionRecordBusy") => continue,
                    Err(e) => {
                        tracing::warn!(%nick, attempt, code = ?e.code, "queued resume failed (letter stays in the mailbox)");
                        return;
                    }
                }
            }
            tracing::warn!(%nick, "queued resume not done in time (letter stays in the mailbox)");
        });
    }

    /// Retries a semantic prompt into a live session that is busy with a turn, every
    /// [`BUSY_RETRY`] up to [`BUSY_RETRY_MAX`] tries; stops when the session goes away.
    fn queue_for_busy(&self, nick: String, address: SessionAddress, agent: AgentId, prompt: String) {
        tracing::info!(%nick, "letter queued for a busy session");
        let this = self.clone();
        tokio::spawn(async move {
            for attempt in 1..=BUSY_RETRY_MAX {
                tokio::time::sleep(BUSY_RETRY).await;
                if this.shared.validate_address(&address).is_err() {
                    tracing::warn!(%nick, attempt, "queued letter dropped: session ended (letter stays in the mailbox)");
                    return;
                }
                let action = InputAction::SubmitPrompt(PromptPayload { text: prompt.clone(), framing: prompt_framing(&agent) });
                if this.shared.dispatch_input_bounded(&address, action).await.is_ok() {
                    tracing::info!(%nick, attempt, "queued letter woke a live session");
                    return;
                }
            }
            tracing::warn!(%nick, "queued letter not delivered in time (letter stays in the mailbox)");
        });
    }

    /// Wake signal for a session this node spawned (from its mailbox listener): no letter text.
    pub async fn wake_spawned(&self, nick: String, from: String, event_id: String, body: String) -> Result<&'static str, String> {
        self.wake(NodeLetter { nick, from, event_id, body, own_spawn: true }).await
    }

    /// Wakes the session that owns `letter.nick`. `Ok` names what happened; `Err` is safe to
    /// log (it never carries the letter text).
    pub async fn wake(&self, letter: NodeLetter) -> Result<&'static str, String> {
        let (nick, from) = (letter.nick.clone(), letter.from.clone());
        let r = self.wake_letter(letter).await;
        if let Err(reason) = &r {
            // The letter stays in the mailbox; the operator hears about it instead of nobody.
            let (_, addresses) = self.snapshot();
            if let Some(a) = addresses.get(&nick) {
                let id = a.session.instance_id.0;
                let provider = crate::auto_permit::acp_provider(id).unwrap_or("session");
                crate::session_signals::emit(a.workspace_id.as_str(), id, provider, crate::session_signals::Signal::LetterStuck { hint: letter_stuck_hint(&nick, &from, reason) });
            }
        }
        r
    }

    async fn wake_letter(&self, letter: NodeLetter) -> Result<&'static str, String> {
        let (sessions, addresses) = self.snapshot();
        let wake = WakeLetter { body: &letter.body, from_nick: &letter.from, event_id: &letter.event_id, room: None, reaction: None };
        match plan_hosted_wake(&sessions, &letter.nick, &wake) {
            HostedWake::Submit { nick, prompt } => {
                let prompt = if self.policy.deliver_body { prompt } else { crate::wake_pipeline::wake_line(&letter.from, &letter.event_id) };
                let address = addresses.get(&nick).cloned().ok_or("session vanished")?;
                let wake_line = crate::wake_pipeline::wake_line(&letter.from, &letter.event_id);
                let agent = match self.shared.validate_address(&address) {
                    Ok(agent) => agent,
                    Err(_) if letter.own_spawn => return self.resume_with_wake(&nick, &address, wake_line).await,
                    Err(_) => return Err("session is not live".to_owned()),
                };
                let refused = |e: NodeFailure| {
                    tracing::warn!(nick = %nick, code = ?e.code, "letter wake input refused");
                    "session did not take the turn".to_owned()
                };
                // An ACP session takes prompts only (typing into it fails with "Acp does not
                // support this input action"), even while its runtime is busy or limited and the
                // semantic-prompt policy check does not pass.
                let acp = crate::auto_permit::acp_provider(address.session.instance_id.0).is_some();
                let shape = if acp || self.shared.require_session_runtime_policy(&address, ProviderRuntimeRequirement::SemanticPrompt).is_ok() {
                    SessionShape::Semantic
                } else {
                    SessionShape::RawTerminal
                };
                let routed = match route(&self.policy, &nick, shape) {
                    Route::Refuse(_) if letter.own_spawn && shape == SessionShape::RawTerminal => Route::TerminalTyping,
                    // A session this node spawned in ACP takes the wake as a prompt turn.
                    Route::Refuse(_) if letter.own_spawn && shape == SessionShape::Semantic => Route::Deliver,
                    other => other,
                };
                // A spawned semantic session gets the letter itself (the CLI/ACP analogue of the app's chat
                // widget); a raw terminal keeps the one-line notice (typed text must stay one line).
                let prompt = match (letter.own_spawn, shape) {
                    (true, SessionShape::Semantic) => crate::wake_pipeline::letter_prompt(&letter.from, &letter.event_id, &letter.body),
                    (true, _) => crate::wake_pipeline::wake_line(&letter.from, &letter.event_id),
                    _ => prompt,
                };
                match routed {
                    Route::Deliver => {
                        let queued = prompt.clone();
                        let action = InputAction::SubmitPrompt(PromptPayload { text: prompt, framing: prompt_framing(&agent) });
                        match self.shared.dispatch_input_bounded(&address, action).await {
                            // A session this node spawned whose process exited is resumed with
                            // the wake as its first prompt (same record, same m4a account).
                            Err(_) if letter.own_spawn && self.shared.validate_address(&address).is_err() => {
                                return self.resume_with_wake(&nick, &address, wake_line).await
                            }
                            // Still live but did not take the prompt (mid-turn): keep the letter
                            // and hand it over once the session is free, instead of resuming a
                            // record the live session holds (SessionRecordBusy, letter lost).
                            Err(_) if letter.own_spawn && shape == SessionShape::Semantic => {
                                self.queue_for_busy(nick.clone(), address.clone(), agent.clone(), queued);
                                return Ok("queued");
                            }
                            other => other.map_err(refused)?,
                        };
                    }
                    Route::TerminalTyping => {
                        // Terminal capability, granted explicitly for this nick: type the prompt
                        // as one line, then Enter.
                        let text = prompt.replace(['\n', '\r'], " ");
                        match self.shared.dispatch_input_bounded(&address, InputAction::TerminalText(TerminalText { text })).await {
                            // A session this node spawned whose process exited is resumed with
                            // the wake as its first prompt (same record, same m4a account).
                            Err(_) if letter.own_spawn && self.shared.validate_address(&address).is_err() => {
                                return self.resume_with_wake(&nick, &address, wake_line).await
                            }
                            other => other.map_err(refused)?,
                        };
                        // TUIs (kimi) treat a fast Enter right after the text as part of a paste.
                        tokio::time::sleep(std::time::Duration::from_millis(400)).await;
                        self.shared.dispatch_input_bounded(&address, InputAction::TerminalControl(TerminalControl::ControlM)).await.map_err(refused)?;
                    }
                    Route::Refuse(why) => {
                        tracing::info!(nick = %nick, "letter not typed: no terminal grant");
                        return Err(why);
                    }
                    Route::Routine(_) => return Err("unexpected routine route for a live session".into()),
                    Route::Spawn(_) => return Err("unexpected spawn route for a live session".into()),
                }
                tracing::info!(nick = %nick, "letter woke a live session");
                Ok("submitted")
            }
            HostedWake::Spawn { nick, .. } => match self.policy.spawn_profile(&nick).cloned() {
                Some(profile) => {
                    tokio::task::spawn_blocking(move || launch_spawn_profile(&profile))
                        .await
                        .map_err(|_| "spawn profile task failed".to_owned())?
                        .map_err(|e| e)?;
                    tracing::info!(nick = %nick, "letter woke an offline session via spawn profile");
                    Ok("spawned")
                }
                None => Err("spawn on letter is not enabled on this node".into()),
            },
            HostedWake::Refuse(why) => match route(&self.policy, &letter.nick, SessionShape::None) {
                Route::Spawn(profile) => {
                    let nick = letter.nick.clone();
                    tokio::task::spawn_blocking(move || launch_spawn_profile(&profile))
                        .await
                        .map_err(|_| "spawn profile task failed".to_owned())?
                        .map_err(|e| e)?;
                    tracing::info!(nick = %nick, "letter woke an offline session via spawn profile");
                    Ok("spawned")
                }
                Route::Routine(target) => {
                    let routine_letter = RoutineLetter {
                        to: letter.nick.clone(),
                        from: letter.from.clone(),
                        event_id: letter.event_id.clone(),
                        subject: None,
                        room: None,
                        body: letter.body.clone(),
                        reply: None,
                        note: None,
                    };
                    let nick = letter.nick.clone();
                    match tokio::task::spawn_blocking(move || deliver_routine(&target, &routine_letter)).await {
                        Ok(Ok(_)) => {
                            tracing::info!(nick = %nick, "letter delivered to a routine");
                            Ok("routine")
                        }
                        Ok(Err(e)) => Err(e),
                        Err(_) => Err("routine delivery task failed".into()),
                    }
                }
                _ => Err(why),
            },
        }
    }
}

/// `nick=<n> from=<f>: <reason>` -- short, control characters dropped, no letter text.
pub fn letter_stuck_hint(nick: &str, from: &str, reason: &str) -> String {
    format!("nick={nick} from={from}: {reason}").chars().filter(|c| !c.is_control()).take(160).collect()
}

#[cfg(test)]
mod letter_stuck_tests {
    #[test]
    fn hint_names_the_session_and_sender_only() {
        let h = super::letter_stuck_hint("codex-21", "hatcheryweb", "exited session could not be resumed\n");
        assert_eq!(h, "nick=codex-21 from=hatcheryweb: exited session could not be resumed");
    }
}