swarm / office/src/activity.rs

//! Live LLM/GPU activity: every decision and write job the pool runs,
//! in flight and finished, with the prompt and the model's reply. The TUI's
//! "llm" tab reads the JSON snapshot file this module flushes to disk.
//!
//! One lock, one flush — the snapshot is serialized INSIDE the lock and
//! written to disk AFTER the lock is released (flush() must never re-lock).

use serde::Serialize;
use std::collections::BTreeMap;
use std::path::PathBuf;
use std::sync::{Mutex, OnceLock};
use std::time::{SystemTime, UNIX_EPOCH};

static ROOT: OnceLock<PathBuf> = OnceLock::new();
static FLIGHTS: Mutex<Option<BTreeMap<String, Flight>>> = Mutex::new(None);

#[derive(Serialize, Clone)]
pub struct Flight {
    pub person: String,
    pub kind: String,
    /// endpoint label, e.g. "http://localhost:8085 [Qwen3.8-27B…]"
    pub model: String,
    /// queued | computing | done | failed
    pub status: String,
    /// wall-clock seconds since epoch (TUI sorts newest first)
    pub started: f64,
    /// the user prompt (the living context) truncated
    pub prompt: String,
    /// model output (or verbatim dialogue turns)
    pub reply: String,
}

const MAX_PROMPT: usize = 1600;
const MAX_REPLY: usize = 2400;
const MAX_ENTRIES: usize = 60;

pub fn set_root(path: PathBuf) {
    if let Some(dir) = path.parent() {
        let _ = std::fs::create_dir_all(dir);
    }
    let _ = ROOT.set(path);
}

pub fn key(person: &str, write: bool) -> String {
    if write {
        format!("{person} · write")
    } else {
        person.to_string()
    }
}

fn now() -> f64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_secs_f64())
        .unwrap_or(0.0)
}

fn clip(s: &str, n: usize) -> String {
    s.chars().take(n).collect()
}

/// Mutate under the lock, snapshot inside the same critical section,
/// then write AFTER the mutex is dropped. One lock acquisition per event.
fn with_map(edit: impl FnOnce(&mut BTreeMap<String, Flight>) -> Vec<Flight>) {
    let payload = {
        let mut guard = FLIGHTS.lock().unwrap_or_else(|e| e.into_inner());
        if guard.is_none() {
            *guard = Some(BTreeMap::new());
        }
        let map = guard.as_mut().expect("activity map");
        let out = edit(map);
        out
    };
    write_snapshot(&payload);
}

fn write_snapshot(mut v: &[Flight]) {
    let Some(root) = ROOT.get() else { return };
    // Arena copy so we can sort without holding the map borrow.
    let mut copy: Vec<Flight> = v.to_vec();
    copy.sort_by(|a, b| {
        b.started
            .partial_cmp(&a.started)
            .unwrap_or(std::cmp::Ordering::Equal)
    });
    if let Ok(text) = serde_json::to_string_pretty(&copy) {
        let tmp = root.with_extension("tmp.json");
        if std::fs::write(&tmp, text).is_ok() {
            let _ = std::fs::rename(&tmp, root);
        }
    }
    let _ = &mut v.clone();
}

/// A decision/write job queued in the pool.
pub fn dispatch(person: &str, write: bool, kind: &str, system: &str, user: &str) {
    let k = key(person, write);
    let flight = Flight {
        person: person.to_string(),
        kind: kind.to_string(),
        model: String::new(),
        status: "queued".into(),
        started: now(),
        prompt: clip(user, MAX_PROMPT),
        reply: format!("(persona tail: {}…)", clip(system, 200)),
    };
    let snapshot = with_map(|m| {
        m.insert(k, flight);
        prune(m);
        m.values().cloned().collect()
    });
    let _ = snapshot; // snapshot already flushed by with_map
}

/// Convenience: mutate + return nothing extra (single snapshot serialize
/// inside `with_map`); keeps one critical section per transition.
fn prune(map: &mut BTreeMap<String, Flight>) {
    if map.len() > MAX_ENTRIES {
        let mut done_keys: Vec<String> = map
            .iter()
            .filter(|(_, f)| f.status.starts_with("done") || f.status.starts_with("failed"))
            .map(|(k, _)| k.clone())
            .collect();
        done_keys.sort();
        let excess = map.len() - MAX_ENTRIES;
        for k in done_keys.into_iter().take(excess) {
            map.remove(&k);
        }
    }
}

/// A worker picked it up and the GPU is generating right now.
pub fn started(person: &str, write: bool, endpoint_label: &str) {
    let k = key(person, write);
    with_map(|m| {
        if let Some(f) = m.get_mut(&k) {
            f.model = endpoint_label.to_string();
            f.status = "computing".into();
        }
        m.values().cloned().collect()
    });
}

/// A dialogue turn lands (Discuss sessions run several calls).
pub fn turn(person: &str, who: &str, line: &str) {
    let k = key(person, false);
    with_map(|m| {
        if let Some(f) = m.get_mut(&k) {
            if !f.reply.is_empty() {
                f.reply.push('\n');
            }
            f.reply.push_str(&format!("{who}: {}", clip(line, 400)));
        }
        m.values().cloned().collect()
    });
}

pub fn finish(person: &str, write: bool, ok: bool, reply: &str) {
    let k = key(person, write);
    with_map(|m| {
        if let Some(f) = m.get_mut(&k) {
            f.status = if ok { "done".into() } else { "failed".into() };
            f.reply = clip(reply, MAX_REPLY);
        }
        m.values().cloned().collect()
    });
}

/// In-memory snapshot for anyone in this process (tests/debug).
pub fn snapshot() -> Vec<Flight> {
    let guard = FLIGHTS.lock().unwrap_or_else(|e| e.into_inner());
    guard
        .as_ref()
        .map(|m| m.values().cloned().collect())
        .unwrap_or_default()
}