klieo-workflow-api 3.2.0

//! Per-run live-progress channels.
//!
//! A lightweight, self-contained alternative to `klieo-ops-viz`'s multi-run
//! `LiveProgress`: the run-service only tails ONE run at a time, so each run
//! gets its own `tokio::sync::broadcast` channel keyed by run id. This keeps
//! the run-service free of the ops-viz SPA crate; it depends only on `tokio`
//! and `klieo_core::AgentEvent` (already a dependency).
//!
//! Memory is bounded: a channel is created at submit and removed when the run
//! finalizes, so the map holds at most the set of in-flight runs (itself
//! capped by the concurrency limit). Each channel buffers up to
//! [`PROGRESS_BUFFER`] events; a subscriber that lags past it sees a `Lagged`
//! gap, never an error.

use klieo_core::agent::AgentEvent;
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use tokio::sync::broadcast;

/// Per-run broadcast buffer depth. A subscriber lagging past this many
/// unread events skips the gap (a `Lagged`) rather than erroring.
const PROGRESS_BUFFER: usize = 256;

/// Registry of per-run live-progress channels, shared across the router.
#[derive(Clone, Default)]
pub(crate) struct ProgressHub {
    channels: Arc<Mutex<HashMap<String, broadcast::Sender<AgentEvent>>>>,
}

impl ProgressHub {
    /// Create a run's channel and return the sender to install as
    /// `ctx.progress`. Called at submit so a client can attach before the
    /// executor starts emitting.
    pub fn register(&self, run_id: &str) -> broadcast::Sender<AgentEvent> {
        let (sender, _receiver) = broadcast::channel(PROGRESS_BUFFER);
        self.lock().insert(run_id.to_string(), sender.clone());
        sender
    }

    /// Remove a run's channel (call on terminal). Dropping the last sender
    /// closes any live subscriber's stream cleanly.
    pub fn deregister(&self, run_id: &str) {
        self.lock().remove(run_id);
    }

    /// Subscribe to a run's live events, or `None` when no run is currently
    /// live under `run_id` (unknown or already finalized).
    pub fn subscribe(&self, run_id: &str) -> Option<broadcast::Receiver<AgentEvent>> {
        self.lock().get(run_id).map(broadcast::Sender::subscribe)
    }

    fn lock(&self) -> std::sync::MutexGuard<'_, HashMap<String, broadcast::Sender<AgentEvent>>> {
        self.channels.lock().expect("progress hub mutex poisoned")
    }
}

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

    #[test]
    fn subscribe_reflects_register_and_deregister() {
        let hub = ProgressHub::default();
        assert!(hub.subscribe("run-1").is_none());
        let _sender = hub.register("run-1");
        assert!(hub.subscribe("run-1").is_some());
        hub.deregister("run-1");
        assert!(hub.subscribe("run-1").is_none());
    }

    #[test]
    fn registered_sender_delivers_to_subscriber() {
        let hub = ProgressHub::default();
        let sender = hub.register("run-1");
        let mut receiver = hub.subscribe("run-1").unwrap();
        sender.send(AgentEvent::LlmCallStarted).unwrap();
        assert!(matches!(
            receiver.try_recv(),
            Ok(AgentEvent::LlmCallStarted)
        ));
    }
}