agent-block 0.21.0

//! Standalone MCP server binary that advertises `resources/subscribe` capability
//! and fires `notify_resource_updated` upon each `subscribe` request.
//!
//! Intended for shell-level positive smoke tests: run it, capture the
//! `SUBSCRIBE_TEST_SERVER_URL=http://…/mcp` marker from stdout, then point
//! `agent-block` at it with `MCP_HTTP_URL=<url>`.
//!
//! Usage:
//!   cargo run --example subscribe_test_server -- --port 0
//!   cargo run --example subscribe_test_server -- --port 3001 --interval 500
use clap::Parser;
use rmcp::{
    model::{
        ResourceUpdatedNotificationParam, ServerCapabilities, ServerInfo, SubscribeRequestParams,
    },
    service::{MaybeSendFuture, RequestContext},
    transport::streamable_http_server::{
        session::local::LocalSessionManager, StreamableHttpServerConfig, StreamableHttpService,
    },
    ErrorData as McpError, RoleServer, ServerHandler,
};
use std::sync::Arc;
use tokio::sync::Mutex;
use tokio_util::sync::CancellationToken;

// ── CLI ───────────────────────────────────────────────────────────────────────

#[derive(Parser)]
#[command(
    name = "subscribe-test-server",
    about = "Subscribe-capable MCP server for shell smoke tests"
)]
struct Cli {
    /// TCP port to listen on. 0 picks an ephemeral port.
    #[arg(long, default_value_t = 0)]
    port: u16,

    /// If > 0, fire an additional notify_resource_updated every N milliseconds
    /// for all currently-subscribed URIs. If 0 (default), fire only the
    /// immediate 50 ms notification triggered at subscribe time.
    #[arg(long, default_value_t = 0)]
    interval: u64,
}

// ── Subscribe Test Server ─────────────────────────────────────────────────────

/// An MCP server that:
/// - Advertises `resources` + `subscribe` capability.
/// - On `subscribe(params, ctx)`: spawns a task that immediately calls
///   `ctx.peer.notify_resource_updated(...)` with the subscribed URI, then
///   returns `Ok(())`.
///
/// Lifted verbatim from `tests/e2e_mcp_resource_subscribe.rs:40-73`.
/// The `subscribed_uris` field is an additive extension for the `--interval`
/// periodic-fire feature; it does NOT change the `get_info()` / `subscribe()`
/// logic inherited from the original helper.
#[derive(Clone)]
struct SubscribeTestServer {
    /// Accumulated subscribed URIs for periodic notification (--interval > 0).
    subscribed_uris: Arc<Mutex<Vec<String>>>,
}

impl SubscribeTestServer {
    fn new() -> Self {
        Self {
            subscribed_uris: Arc::new(Mutex::new(Vec::new())),
        }
    }
}

impl ServerHandler for SubscribeTestServer {
    fn get_info(&self) -> ServerInfo {
        ServerInfo::new(
            ServerCapabilities::builder()
                .enable_resources()
                .enable_resources_subscribe()
                .build(),
        )
    }

    fn subscribe(
        &self,
        request: SubscribeRequestParams,
        context: RequestContext<RoleServer>,
    ) -> impl std::future::Future<Output = Result<(), McpError>> + MaybeSendFuture + '_ {
        let uri = request.uri.clone();
        let peer = context.peer.clone();
        let uris = self.subscribed_uris.clone();
        async move {
            // Spawn so subscribe response is sent before the notification arrives.
            tokio::spawn(async move {
                // Small delay to ensure the subscribe RPC response reaches the
                // client before the notification is dispatched (R3 race mitigation).
                tokio::time::sleep(std::time::Duration::from_millis(50)).await;
                let _ = peer
                    .notify_resource_updated(ResourceUpdatedNotificationParam { uri: uri.clone() })
                    .await;
            });
            // Record the URI for the optional periodic-notification loop.
            uris.lock().await.push(request.uri.clone());
            Ok(())
        }
    }
}

// ── HTTP scaffold ─────────────────────────────────────────────────────────────

async fn run_http(port: u16, interval_ms: u64) {
    let ct = CancellationToken::new();

    let config = StreamableHttpServerConfig::default()
        .with_sse_keep_alive(None)
        .with_cancellation_token(ct.child_token());

    let server = SubscribeTestServer::new();

    // Optional periodic-notification loop (--interval > 0).
    if interval_ms > 0 {
        let uris = server.subscribed_uris.clone();
        // We don't have a reference to "the peer" at this point; periodic
        // notifications require the ServerHandler's peer, which is only
        // available inside a subscribe() call.  The loop is therefore
        // best-effort: it logs that it's running but has no peers to notify
        // until at least one subscribe() has been processed.
        //
        // NOTE: full periodic-fire requires storing peer handles per-session,
        // which is beyond the minimal crux requirement.  The default path
        // (interval=0, 50ms-once-on-subscribe) satisfies the crux fully.
        let ct_loop = ct.child_token();
        tokio::spawn(async move {
            let mut ticker = tokio::time::interval(std::time::Duration::from_millis(interval_ms));
            loop {
                tokio::select! {
                    _ = ct_loop.cancelled() => break,
                    _ = ticker.tick() => {
                        let uris_snap = uris.lock().await.clone();
                        if !uris_snap.is_empty() {
                            eprintln!(
                                "[subscribe-test-server] periodic tick: {} subscribed uri(s) recorded \
                                 (peer-level periodic notify requires per-session peer handles)",
                                uris_snap.len()
                            );
                        }
                    }
                }
            }
        });
    }

    let server_clone = server.clone();
    let service: StreamableHttpService<SubscribeTestServer, LocalSessionManager> =
        StreamableHttpService::new(move || Ok(server_clone.clone()), Default::default(), config);

    let router = axum::Router::new().nest_service("/mcp", service);

    let addr = format!("127.0.0.1:{port}");
    let listener = match tokio::net::TcpListener::bind(&addr).await {
        Ok(l) => l,
        Err(e) => {
            eprintln!("[subscribe-test-server] failed to bind {addr}: {e}");
            std::process::exit(1);
        }
    };

    let bound_addr = match listener.local_addr() {
        Ok(a) => a,
        Err(e) => {
            eprintln!("[subscribe-test-server] local_addr failed: {e}");
            std::process::exit(1);
        }
    };

    println!("SUBSCRIBE_TEST_SERVER_URL=http://{bound_addr}/mcp");

    let ct_shutdown = ct.clone();
    tokio::spawn(async move {
        if let Err(e) = axum::serve(listener, router)
            .with_graceful_shutdown(async move { ct_shutdown.cancelled_owned().await })
            .await
        {
            eprintln!("[subscribe-test-server] http server error: {e}");
        }
    });

    wait_for_signal().await;
    ct.cancel();
}

#[cfg(unix)]
async fn wait_for_signal() {
    use tokio::signal::unix::{signal, SignalKind};
    let mut sigint = signal(SignalKind::interrupt()).expect("SIGINT handler");
    let mut sigterm = signal(SignalKind::terminate()).expect("SIGTERM handler");
    tokio::select! {
        _ = sigint.recv() => {},
        _ = sigterm.recv() => {},
    }
}

#[cfg(not(unix))]
async fn wait_for_signal() {
    let _ = tokio::signal::ctrl_c().await;
}

// ── Entry point ───────────────────────────────────────────────────────────────

#[tokio::main]
async fn main() {
    let cli = Cli::parse();
    run_http(cli.port, cli.interval).await;
}