huskies/server/src/http/mcp/progress.rs

//! Task-local plumbing for MCP `notifications/progress` events.
//!
//! Long-running tool handlers (notably `tool_run_tests`) emit progress events
//! during execution so the MCP client's transport timer resets and the call
//! never surfaces as a tool-call error to the agent. The HTTP MCP handler
//! installs an emitter via `tokio::task_local!` before dispatching the tool
//! and converts emitted events into SSE messages on the response stream.
//!
//! Tool handlers call [`emit_progress`] unconditionally — when no emitter is
//! installed (plain JSON response path, or test code), the call is a no-op.
//! That keeps handler code free of `if let Some(...) =` ceremony for the
//! progress-aware branch.
//!
//! Wire format follows MCP 2025-03-26: the emitter wraps the data in a
//! standard JSON-RPC notification with method `notifications/progress` and a
//! `progressToken` echoed back to the client.

use serde_json::{Value, json};
use tokio::sync::mpsc::UnboundedSender;

/// Per-request channel for emitting progress notifications back to an MCP
/// client. Installed in a `tokio::task_local!` scope by the SSE response
/// path and consumed by the SSE stream producer.
#[derive(Clone)]
pub struct ProgressEmitter {
    /// The client-supplied opaque token from `params._meta.progressToken`.
    /// Echoed back unchanged in every progress notification so the client
    /// can correlate progress with the originating request.
    pub token: Value,
    /// Channel into the SSE response stream. Each value sent is a
    /// fully-formed `notifications/progress` JSON-RPC message ready to be
    /// serialised as an SSE `data:` field.
    pub tx: UnboundedSender<Value>,
}

tokio::task_local! {
    /// Set by the SSE response path before dispatching a tool call. Unset
    /// in the plain JSON path and in tests, where [`emit_progress`] no-ops.
    pub static EMITTER: ProgressEmitter;
}

/// Emit a progress notification to the current request's SSE stream, if one
/// is attached. No-op when no emitter is in scope (plain JSON path).
///
/// `progress` is a monotonically increasing value (typically seconds elapsed
/// for long-running tools); `total` is optional and omitted when unknown;
/// `message` is a short human-readable description of the current state.
pub fn emit_progress(progress: f64, total: Option<f64>, message: Option<&str>) {
    let _ = EMITTER.try_with(|e| {
        let mut params = json!({
            "progressToken": e.token,
            "progress": progress,
        });
        if let Some(t) = total {
            params["total"] = json!(t);
        }
        if let Some(m) = message {
            params["message"] = json!(m);
        }
        let notification = json!({
            "jsonrpc": "2.0",
            "method": "notifications/progress",
            "params": params,
        });
        // Send is fire-and-forget. If the receiver dropped (client
        // disconnected mid-stream), we don't care — the tool dispatch
        // task keeps running and writes its final state to the CRDT.
        let _ = e.tx.send(notification);
    });
}

#[cfg(test)]
mod tests {
    use super::*;
    use tokio::sync::mpsc::unbounded_channel;

    #[tokio::test]
    async fn emit_progress_no_op_without_emitter() {
        // Calling outside a task_local scope must not panic.
        emit_progress(1.0, None, Some("hello"));
    }

    #[tokio::test]
    async fn emit_progress_sends_notification_when_emitter_installed() {
        let (tx, mut rx) = unbounded_channel();
        let emitter = ProgressEmitter {
            token: json!("test-token"),
            tx,
        };
        EMITTER
            .scope(emitter, async {
                emit_progress(5.0, Some(10.0), Some("halfway"));
            })
            .await;

        let notif = rx.recv().await.expect("notification must be delivered");
        assert_eq!(notif["method"], "notifications/progress");
        assert_eq!(notif["params"]["progressToken"], "test-token");
        assert_eq!(notif["params"]["progress"], 5.0);
        assert_eq!(notif["params"]["total"], 10.0);
        assert_eq!(notif["params"]["message"], "halfway");
    }

    #[tokio::test]
    async fn emit_progress_omits_optional_fields() {
        let (tx, mut rx) = unbounded_channel();
        let emitter = ProgressEmitter {
            token: json!(42),
            tx,
        };
        EMITTER
            .scope(emitter, async {
                emit_progress(1.0, None, None);
            })
            .await;

        let notif = rx.recv().await.unwrap();
        assert_eq!(notif["params"]["progressToken"], 42);
        assert_eq!(notif["params"]["progress"], 1.0);
        assert!(notif["params"].get("total").is_none());
        assert!(notif["params"].get("message").is_none());
    }
}