path: root/crates/dinggerd/src/hooks.rs

//! Hook receiver: agents report their own state (herdr's integration model).
//!
//! A unix socket accepts JSON lines from agent extensions:
//!
//! ```json
//! {"method":"pane.report_agent","params":{"pane_id":"%3","agent":"pi","state":"working"}}
//! ```
//!
//! A report may carry the agent's session reference (`agent_session_path`
//! or `agent_session_id`, herdr's names), which restore uses to resume it.
//! A report without a state only records the session (Claude's
//! SessionStart hook, `dingger hook claude`); state stays with the screen.
//! A report may also carry the agent's `importance` (`low`, `normal`,
//! `high`): pi's extension sends it with every report, `dingger importance`
//! alone.
//!
//! Reports are state-change *events*: the watcher applies every one in
//! order, so a turn that starts and finishes between two polls is still a
//! completion, and each report wakes the watcher immediately.
//!
//! `{"method":"pane.mark_unread","params":{"pane_id":"%3"}}` marks an idle
//! agent unread (`dingger unread`, pi's /unread): it shows as done again.
//!
//! A pi report also carries the extension's `extension` version and its
//! `control` socket, where the daemon asks it to reload (`{"method":
//! "reload"}`; it answers whether it does): `dingger reload-agents` sends
//! `{"method":"agents.reload_outdated"}`, and the watcher reloads each agent
//! on an old extension once it is idle and off screen.
//!
//! Each line is answered with `{"ok":true}`.
//!
//! `{"method":"poke"}` only wakes the watcher; `dingger focus` and tmux's
//! focus-change hooks send it so attendance updates without waiting for the
//! next poll. The last reported state stays authoritative for its pane
//! until the next report or until the pane stops running the agent; screen
//! detection is the fallback for panes that never report.

use std::collections::{HashMap, HashSet};
use std::io::{BufRead, BufReader};
use std::os::unix::net::{UnixListener, UnixStream};
use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Condvar, Mutex, MutexGuard, OnceLock};
use std::time::{Duration, Instant};

use detect_core::AgentState;

use crate::importance::Importance;
use serde::Deserialize;
use tracing::{debug, info, warn};

pub fn socket_path() -> PathBuf {
    if let Some(path) = std::env::var_os("DINGGER_AGENT_SOCKET") {
        return PathBuf::from(path);
    }
    crate::snapshot::state_dir().join("agent.sock")
}

#[derive(Debug, Default)]
struct PaneReports {
    /// Last reported state: authoritative until the next report.
    last: Option<AgentState>,
    /// Reports not yet applied by the watcher, oldest first.
    pending: Vec<AgentState>,
    /// The agent's latest session reference.
    session: Option<String>,
    /// The agent's importance, if reported.
    importance: Option<Importance>,
    /// The version of the extension that reported, and its control socket.
    extension: Option<String>,
    control: Option<String>,
    /// When the last report arrived.
    at: Option<Instant>,
}

fn reports() -> MutexGuard<'static, HashMap<String, PaneReports>> {
    static REPORTS: OnceLock<Mutex<HashMap<String, PaneReports>>> = OnceLock::new();
    REPORTS
        .get_or_init(|| Mutex::new(HashMap::new()))
        .lock()
        .unwrap_or_else(|poisoned| poisoned.into_inner())
}

/// Panes marked unread, not yet applied by the watcher.
fn unread_marks() -> MutexGuard<'static, HashSet<String>> {
    static MARKS: OnceLock<Mutex<HashSet<String>>> = OnceLock::new();
    MARKS
        .get_or_init(|| Mutex::new(HashSet::new()))
        .lock()
        .unwrap_or_else(|poisoned| poisoned.into_inner())
}

/// Whether the pane was marked unread since the last call.
pub fn take_unread(pane_id: &str) -> bool {
    unread_marks().remove(pane_id)
}

/// Whether agents on an old extension should be reloaded (`dingger
/// reload-agents`); the watcher clears it when none is left.
static RELOAD_WANTED: AtomicBool = AtomicBool::new(false);

pub fn reload_wanted() -> bool {
    RELOAD_WANTED.load(Ordering::Relaxed)
}

pub fn reload_done() {
    RELOAD_WANTED.store(false, Ordering::Relaxed);
}

fn wake_signal() -> &'static (Mutex<bool>, Condvar) {
    static WAKE: OnceLock<(Mutex<bool>, Condvar)> = OnceLock::new();
    WAKE.get_or_init(|| (Mutex::new(false), Condvar::new()))
}

fn wake() {
    let (flag, condvar) = wake_signal();
    *flag.lock().unwrap_or_else(|poisoned| poisoned.into_inner()) = true;
    condvar.notify_all();
}

/// Sleep until woken by a report or poke, or until `timeout` passes.
pub fn wait_for_wake(timeout: Duration) {
    let (flag, condvar) = wake_signal();
    let guard = flag.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
    let (mut guard, _) = condvar
        .wait_timeout_while(guard, timeout, |woken| !*woken)
        .unwrap_or_else(|poisoned| poisoned.into_inner());
    *guard = false;
}

/// Wake a running watcher (used by `dingger focus` and tmux focus hooks).
pub fn poke() -> Result<(), String> {
    send("{\"method\":\"poke\"}")
}

/// Send one JSON message to the running watcher.
pub fn send(message: &str) -> Result<(), String> {
    use std::io::Write;
    let mut stream = UnixStream::connect(socket_path())
        .map_err(|err| format!("dingger watch is not running: {err}"))?;
    stream
        .write_all(format!("{message}\n").as_bytes())
        .map_err(|err| format!("send failed: {err}"))
}

#[derive(Deserialize)]
struct HookMessage {
    method: String,
    #[serde(default)]
    params: Option<HookParams>,
}

#[derive(Deserialize)]
struct HookParams {
    pane_id: String,
    #[serde(default)]
    state: Option<String>,
    #[serde(default)]
    agent_session_path: Option<String>,
    #[serde(default)]
    agent_session_id: Option<String>,
    #[serde(default)]
    importance: Option<String>,
    #[serde(default)]
    extension: Option<String>,
    #[serde(default)]
    control: Option<String>,
}

fn handle_stream(stream: UnixStream) {
    // Each message gets a reply line: an extension waits for one and sends
    // the message again without it (herdr's protocol), so a report would
    // arrive twice and take seconds.
    let mut replies = stream.try_clone().ok();
    for line in BufReader::new(stream).lines() {
        let Ok(line) = line else { return };
        if let Some(out) = replies.as_mut() {
            use std::io::Write;
            if out.write_all(b"{\"ok\":true}\n").is_err() {
                replies = None; // the sender does not read replies
            }
        }
        let Ok(message) = serde_json::from_str::<HookMessage>(&line) else {
            warn!(%line, "invalid hook message");
            continue;
        };
        if message.method == "poke" {
            wake();
            continue;
        }
        if message.method == "agents.reload_outdated" {
            info!("reload agents on an old extension: requested");
            RELOAD_WANTED.store(true, Ordering::Relaxed);
            wake();
            continue;
        }
        if message.method == "pane.mark_unread" {
            if let Some(params) = message.params {
                info!(pane = %params.pane_id, "marked unread");
                unread_marks().insert(params.pane_id);
                wake();
            }
            continue;
        }
        let (true, Some(params)) = (message.method == "pane.report_agent", message.params) else {
            continue;
        };
        let state = match params.state.as_deref() {
            None => None,
            Some("working") => Some(AgentState::Working),
            Some("blocked") => Some(AgentState::Blocked),
            Some("idle") => Some(AgentState::Idle),
            Some(other) => {
                debug!(state = other, "unknown hook state");
                continue;
            }
        };
        let importance = params.importance.as_deref().and_then(|value| {
            let level = Importance::parse(value);
            if level.is_none() {
                debug!(importance = value, "unknown importance");
            }
            level
        });
        match state {
            Some(state) => info!(pane = %params.pane_id, ?state, ?importance, "hook report"),
            None => info!(pane = %params.pane_id, ?importance, "hook report without state"),
        }
        let reported = Reported {
            session: params.agent_session_path.or(params.agent_session_id),
            importance,
            extension: params.extension,
            control: params.control,
        };
        record(params.pane_id, state, reported);
        wake();
    }
}

/// What a report carries besides its state; each kept until reported again.
#[derive(Debug, Default)]
struct Reported {
    session: Option<String>,
    importance: Option<Importance>,
    extension: Option<String>,
    control: Option<String>,
}

/// A value fit for a tmux option: not empty, no control characters.
fn option_value(value: Option<String>) -> Option<String> {
    value.filter(|v| !v.is_empty() && !v.chars().any(char::is_control))
}

fn record(pane_id: String, state: Option<AgentState>, reported: Reported) {
    let mut reports = reports();
    let pane = reports.entry(pane_id).or_default();
    pane.at = Some(Instant::now());
    if let Some(state) = state {
        pane.last = Some(state);
        pane.pending.push(state);
    }
    if let Some(session) = option_value(reported.session) {
        pane.session = Some(session);
    }
    if reported.importance.is_some() {
        pane.importance = reported.importance;
    }
    if let Some(extension) = option_value(reported.extension) {
        pane.extension = Some(extension);
    }
    if let Some(control) = option_value(reported.control) {
        pane.control = Some(control);
    }
}

/// Start listening; reports accumulate for the watcher to consume.
pub fn spawn_listener() -> Result<(), String> {
    let path = socket_path();
    if let Some(parent) = path.parent() {
        std::fs::create_dir_all(parent).map_err(|err| format!("mkdir failed: {err}"))?;
    }
    let _ = std::fs::remove_file(&path);
    let listener =
        UnixListener::bind(&path).map_err(|err| format!("bind {}: {err}", path.display()))?;
    info!(path = %path.display(), "hook receiver listening");
    std::thread::spawn(move || {
        for stream in listener.incoming() {
            match stream {
                Ok(stream) => {
                    std::thread::spawn(move || handle_stream(stream));
                }
                Err(err) => warn!(%err, "hook listener accept failed"),
            }
        }
    });
    Ok(())
}

/// The pane's hook observations for this tick: every report since the last
/// tick, in order; otherwise the last reported state; `None` if the pane
/// never reported to this process.
pub fn take_observations(pane_id: &str) -> Option<Vec<AgentState>> {
    let mut reports = reports();
    let pane = reports.get_mut(pane_id)?;
    if pane.pending.is_empty() {
        return pane.last.map(|state| vec![state]);
    }
    Some(std::mem::take(&mut pane.pending))
}

/// The last state the pane's extension reported, if any.
pub fn last_state(pane_id: &str) -> Option<AgentState> {
    reports().get(pane_id).and_then(|pane| pane.last)
}

/// The pane's agent session reference, if its extension reported one.
pub fn session(pane_id: &str) -> Option<String> {
    reports().get(pane_id).and_then(|pane| pane.session.clone())
}

/// The pane's agent importance, if reported.
pub fn importance(pane_id: &str) -> Option<Importance> {
    reports().get(pane_id).and_then(|pane| pane.importance)
}

/// The version of the pane's extension and its control socket, if reported.
pub fn extension(pane_id: &str) -> (Option<String>, Option<String>) {
    reports().get(pane_id).map(|pane| (pane.extension.clone(), pane.control.clone())).unwrap_or_default()
}

/// Ask the extension behind `control` to reload; its answer: `Ok(())` when
/// it does, else why not (busy, a draft, or no answer).
pub fn request_reload(control: &str) -> Result<(), String> {
    use std::io::Write;
    let mut stream = UnixStream::connect(control).map_err(|err| format!("connect: {err}"))?;
    stream.set_read_timeout(Some(Duration::from_secs(2))).map_err(|err| err.to_string())?;
    stream.write_all(b"{\"method\":\"reload\"}\n").map_err(|err| format!("send: {err}"))?;
    let mut answer = String::new();
    BufReader::new(stream).read_line(&mut answer).map_err(|err| format!("no answer: {err}"))?;
    #[derive(Deserialize)]
    struct Answer {
        ok: bool,
        #[serde(default)]
        reason: Option<String>,
    }
    match serde_json::from_str::<Answer>(&answer) {
        Ok(Answer { ok: true, .. }) => Ok(()),
        Ok(Answer { reason, .. }) => Err(reason.unwrap_or_else(|| "refused".into())),
        Err(_) => Err(format!("unexpected answer: {}", answer.trim())),
    }
}

/// Drop a pane's reports: the agent was not running when the panes were
/// listed at `listed`. A report since then is kept: the agent started
/// after the listing (Claude reports its session once, as it starts).
pub fn drop_reports(pane_id: &str, listed: Instant) {
    let mut reports = reports();
    if reports.get(pane_id).is_some_and(|pane| pane.at.is_none_or(|at| at < listed)) {
        reports.remove(pane_id);
    }
}

/// Drop reports for panes that no longer exist.
pub fn forget(pane_ids_to_keep: &[String]) {
    reports().retain(|pane_id, _| pane_ids_to_keep.contains(pane_id));
}

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

    #[test]
    fn a_session_report_records_only_the_session() {
        record("%t1".into(), None, Reported { session: Some("/p/s.jsonl".into()), ..Default::default() });
        assert_eq!(session("%t1").as_deref(), Some("/p/s.jsonl"));
        assert_eq!(take_observations("%t1"), None, "state stays with the screen");
        assert_eq!(last_state("%t1"), None);
    }

    #[test]
    fn a_report_after_the_pane_listing_is_not_dropped() {
        let listed = Instant::now();
        record("%t2".into(), None, Reported { session: Some("/p/s.jsonl".into()), ..Default::default() });
        drop_reports("%t2", listed);
        assert!(session("%t2").is_some(), "the agent started after the listing");
        drop_reports("%t2", Instant::now());
        assert_eq!(session("%t2"), None, "listed since, without the agent: dropped");
    }

    #[test]
    fn importance_is_kept_until_reported_again() {
        record("%t3".into(), Some(AgentState::Working), Reported { importance: Some(Importance::High), ..Default::default() });
        record("%t3".into(), Some(AgentState::Idle), Reported::default());
        assert_eq!(importance("%t3"), Some(Importance::High), "a report without it keeps it");
        record("%t3".into(), None, Reported { importance: Some(Importance::Normal), ..Default::default() });
        assert_eq!(importance("%t3"), Some(Importance::Normal));
        assert_eq!(last_state("%t3"), Some(AgentState::Idle), "a report without state keeps the state");
    }

    #[test]
    fn the_extension_and_its_control_socket_are_kept_until_reported_again() {
        let reported = || Reported { extension: Some("v1".into()), control: Some("/s/pi-1.sock".into()), ..Default::default() };
        record("%t4".into(), Some(AgentState::Idle), reported());
        record("%t4".into(), None, Reported { importance: Some(Importance::Low), ..Default::default() });
        assert_eq!(extension("%t4"), (Some("v1".into()), Some("/s/pi-1.sock".into())));
        record("%t4".into(), None, Reported { extension: Some("v2".into()), ..Default::default() });
        assert_eq!(extension("%t4").0.as_deref(), Some("v2"));
    }

    #[test]
    fn a_reload_answer_says_whether_it_reloads() {
        let dir = std::env::temp_dir().join(format!("dingger-reload-{}", std::process::id()));
        std::fs::create_dir_all(&dir).unwrap();
        for (answer, want) in [
            ("{\"ok\":true}", Ok(())),
            ("{\"ok\":false,\"reason\":\"draft\"}", Err("draft".to_string())),
        ] {
            let path = dir.join("pi.sock");
            let _ = std::fs::remove_file(&path);
            let listener = UnixListener::bind(&path).unwrap();
            let server = std::thread::spawn(move || {
                use std::io::Write;
                let (mut stream, _) = listener.accept().unwrap();
                let mut line = String::new();
                BufReader::new(stream.try_clone().unwrap()).read_line(&mut line).unwrap();
                assert_eq!(line, "{\"method\":\"reload\"}\n");
                stream.write_all(format!("{answer}\n").as_bytes()).unwrap();
            });
            assert_eq!(request_reload(path.to_str().unwrap()), want);
            server.join().unwrap();
        }
        assert!(request_reload(dir.join("none.sock").to_str().unwrap()).is_err(), "no extension there");
        let _ = std::fs::remove_dir_all(&dir);
    }
}