codex/codex-rs/rmcp-client/tests/mcp_events.rs

mod streamable_http_test_support;

use std::convert::Infallible;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::time::Duration;

use anyhow::Context as _;
use axum::Json;
use axum::Router;
use axum::http::StatusCode;
use axum::response::IntoResponse;
use axum::response::sse::Event;
use axum::response::sse::Sse;
use axum::routing::post;
use codex_api::AuthError;
use codex_api::AuthProvider;
use codex_config::types::AuthKeyringBackendKind;
use codex_config::types::OAuthCredentialsStoreMode;
use codex_exec_server::Environment;
use codex_rmcp_client::RmcpClient;
use futures::StreamExt as _;
use futures::stream;
use http::HeaderMap;
use http::HeaderValue;
use pretty_assertions::assert_eq;
use serde_json::Value;
use serde_json::json;
use tokio::net::TcpListener;
use tokio::sync::Notify;
use tokio::sync::mpsc;
use tokio::time::timeout;

use streamable_http_test_support::create_client;
use streamable_http_test_support::initialize_client;

#[derive(Default)]
struct RevocableAuth {
    unavailable: AtomicBool,
}

impl AuthProvider for RevocableAuth {
    fn add_auth_headers(&self, headers: &mut HeaderMap) -> Result<(), AuthError> {
        if self.unavailable.load(Ordering::SeqCst) {
            return Err(AuthError::Transient(
                "authentication unavailable".to_string(),
            ));
        }
        headers.insert(
            http::header::AUTHORIZATION,
            HeaderValue::from_static("Bearer test-bearer"),
        );
        Ok(())
    }
}

struct StreamClosed {
    event_name: String,
    closed: mpsc::UnboundedSender<String>,
}

impl Drop for StreamClosed {
    fn drop(&mut self) {
        let _ = self.closed.send(self.event_name.clone());
    }
}

fn initialize_response(message: &Value) -> axum::response::Response {
    Json(json!({
        "jsonrpc": "2.0",
        "id": message["id"],
        "result": {
            "protocolVersion": "2025-06-18",
            "capabilities": {},
            "serverInfo": {"name": "plugin-runtime", "version": "1.0.0"},
        },
    }))
    .into_response()
}

fn plugin_runtime_notifications(event_name: &str, request_id: Value) -> [Value; 2] {
    let metadata = json!({"io.modelcontextprotocol/subscriptionId": request_id});
    [
        json!({
            "jsonrpc": "2.0",
            "method": "notifications/events/active",
            "params": {"cursor": null, "truncated": false, "_meta": metadata},
        }),
        json!({
            "jsonrpc": "2.0",
            "method": "notifications/events/event",
            "params": {
                "eventId": format!("event-{event_name}"),
                "name": event_name,
                "timestamp": "2026-08-07T12:00:00Z",
                "data": {"issue": 42},
                "cursor": null,
                "_meta": metadata,
            },
        }),
    ]
}

#[tokio::test]
async fn plugin_runtime_event_streams_are_isolated_and_cancel_locally() -> anyhow::Result<()> {
    let (stream_closed_tx, mut stream_closed_rx) = mpsc::unbounded_channel::<String>();
    let router = Router::new().route(
        "/mcp",
        post(move |Json(message): Json<Value>| {
            let stream_closed_tx = stream_closed_tx.clone();
            async move {
                match message["method"].as_str() {
                    Some("initialize") => initialize_response(&message),
                    Some("notifications/initialized") => StatusCode::ACCEPTED.into_response(),
                    Some("events/stream") => {
                        let event_name = message["params"]["name"]
                            .as_str()
                            .expect("event stream request must contain a name")
                            .to_string();
                        assert_eq!(
                            message.pointer("/params/arguments"),
                            Some(&json!({"project": "codex"}))
                        );
                        assert!(message["params"].get("cursor").is_none());

                        let closed = StreamClosed {
                            event_name: event_name.clone(),
                            closed: stream_closed_tx,
                        };
                        let events = stream::iter(
                            plugin_runtime_notifications(&event_name, message["id"].clone())
                                .into_iter()
                                .map(|notification| {
                                    Ok::<_, Infallible>(
                                        Event::default()
                                            .event("message")
                                            .data(notification.to_string()),
                                    )
                                }),
                        )
                        .chain(stream::pending())
                        .map(move |event| {
                            let _ = &closed;
                            event
                        });

                        Sse::new(events).into_response()
                    }
                    Some("notifications/cancelled") => {
                        panic!("event cancellation should close the local stream without a POST")
                    }
                    method => panic!("unexpected Plugin Runtime request: {method:?}"),
                }
            }
        }),
    );

    let listener = TcpListener::bind("127.0.0.1:0").await?;
    let base_url = format!("http://{}", listener.local_addr()?);
    let server = tokio::spawn(async move { axum::serve(listener, router).await });
    let auth = Arc::new(RevocableAuth::default());
    let client = RmcpClient::new_streamable_http_client(
        "test-streamable-http",
        &format!("{base_url}/mcp"),
        /*bearer_token*/ None,
        /*http_headers*/ None,
        /*env_http_headers*/ None,
        OAuthCredentialsStoreMode::File,
        AuthKeyringBackendKind::default(),
        Environment::default_for_tests().get_http_client(),
        Some(auth.clone()),
    )
    .await?;
    initialize_client(&client).await?;

    let mut first = client
        .send_event_stream_request(Some(json!({
            "name": "github.pull_request.opened",
            "arguments": {"project": "codex"},
        })))
        .await?;
    let mut second = client
        .send_event_stream_request(Some(json!({
            "name": "gmail.message.received",
            "arguments": {"project": "codex"},
        })))
        .await?;

    for (event_name, request) in [
        ("github.pull_request.opened", &mut first),
        ("gmail.message.received", &mut second),
    ] {
        let active = timeout(Duration::from_secs(5), request.notifications.recv())
            .await?
            .context("event stream closed before activation")?;
        assert_eq!(active.method, "notifications/events/active");
        assert_eq!(
            active.params,
            Some(json!({"cursor": null, "truncated": false}))
        );

        let event = timeout(Duration::from_secs(5), request.notifications.recv())
            .await?
            .context("event stream closed before delivery")?;
        assert_eq!(event.method, "notifications/events/event");
        assert_eq!(
            event.params.as_ref().and_then(|params| params.get("name")),
            Some(&json!(event_name))
        );
    }

    first
        .handle
        .cancel(Some("event subscription closed".to_string()))
        .await?;
    let closed = timeout(Duration::from_secs(5), stream_closed_rx.recv())
        .await?
        .context("cancelled stream did not close")?;
    assert_eq!(closed, "github.pull_request.opened");

    // Local cleanup must still work if authentication fails after the stream starts.
    auth.unavailable.store(true, Ordering::SeqCst);
    second
        .handle
        .cancel(Some("event subscription closed".to_string()))
        .await?;
    let closed = timeout(Duration::from_secs(5), stream_closed_rx.recv())
        .await?
        .context("cancelled stream did not close")?;
    assert_eq!(closed, "gmail.message.received");

    client.shutdown().await;
    server.abort();
    let _ = server.await;
    Ok(())
}

#[tokio::test]
async fn plugin_runtime_event_stream_times_out_before_stalled_headers() -> anyhow::Result<()> {
    let request_started = std::sync::Arc::new(Notify::new());
    let server_request_started = std::sync::Arc::clone(&request_started);
    let router = Router::new().route(
        "/mcp",
        post(move |Json(message): Json<Value>| {
            let request_started = std::sync::Arc::clone(&server_request_started);
            async move {
                match message["method"].as_str() {
                    Some("initialize") => initialize_response(&message),
                    Some("notifications/initialized") => StatusCode::ACCEPTED.into_response(),
                    Some("events/stream") => {
                        request_started.notify_one();
                        std::future::pending().await
                    }
                    method => panic!("unexpected Plugin Runtime request: {method:?}"),
                }
            }
        }),
    );

    let listener = TcpListener::bind("127.0.0.1:0").await?;
    let base_url = format!("http://{}", listener.local_addr()?);
    let server = tokio::spawn(async move { axum::serve(listener, router).await });
    let client = create_client(&base_url).await?;
    let request = client
        .send_event_stream_request(Some(json!({
            "name": "github.pull_request.opened",
            "arguments": {},
        })))
        .await?;

    timeout(Duration::from_secs(5), request_started.notified())
        .await
        .context("event stream request did not reach the server")?;
    tokio::time::pause();
    tokio::time::advance(Duration::from_secs(31)).await;
    let error = request
        .handle
        .rx
        .await?
        .expect_err("stalled response headers must time out");
    tokio::time::resume();
    assert!(error.to_string().contains("timed out"));

    client.shutdown().await;
    server.abort();
    let _ = server.await;
    Ok(())
}