rust-sdk/crates/rmcp/tests/test_message_protocol.rs

//cargo test --test test_message_protocol --features "client server"
#![cfg(not(feature = "local"))]

mod common;
use common::handlers::{TestClientHandler, TestServer};
use rmcp::{
    ServiceExt,
    model::*,
    service::{RequestContext, Service},
};

// Tests start here
#[tokio::test]
async fn test_message_roles() {
    let messages = vec![
        SamplingMessage::user_text("user message"),
        SamplingMessage::assistant_text("assistant message"),
    ];

    // Verify all roles can be serialized/deserialized correctly
    let json = serde_json::to_string(&messages).unwrap();
    let deserialized: Vec<SamplingMessage> = serde_json::from_str(&json).unwrap();
    assert_eq!(messages, deserialized);
}

#[tokio::test]
async fn test_context_inclusion_integration() -> anyhow::Result<()> {
    let (server_transport, client_transport) = tokio::io::duplex(4096);

    // Start server
    let server_handle = tokio::spawn(async move {
        let server = TestServer::new().serve(server_transport).await?;
        server.waiting().await?;
        anyhow::Ok(())
    });

    // Start client that honors context requests
    let handler = TestClientHandler::new(true, true);
    let client = handler.clone().serve(client_transport).await?;

    // Test ThisServer context inclusion
    let request = ServerRequest::CreateMessageRequest(CreateMessageRequest::new(
        CreateMessageRequestParams::new(vec![SamplingMessage::user_text("test message")], 100)
            .with_include_context(ContextInclusion::ThisServer),
    ));

    let result = handler
        .handle_request(
            request.clone(),
            RequestContext::new(NumberOrString::Number(1), client.peer().clone()),
        )
        .await?;

    if let ClientResult::CreateMessageResult(result) = result {
        let text = result
            .message
            .content
            .first()
            .unwrap()
            .as_text()
            .unwrap()
            .text
            .as_str();
        assert!(
            text.contains("test context"),
            "Response should include context for ThisServer"
        );
    } else {
        panic!("Expected CreateMessageResult");
    }

    // Test AllServers context inclusion
    let request = ServerRequest::CreateMessageRequest(CreateMessageRequest::new(
        CreateMessageRequestParams::new(vec![SamplingMessage::user_text("test message")], 100)
            .with_include_context(ContextInclusion::AllServers),
    ));

    let result = handler
        .handle_request(
            request.clone(),
            RequestContext::new(NumberOrString::Number(2), client.peer().clone()),
        )
        .await?;

    if let ClientResult::CreateMessageResult(result) = result {
        let text = result
            .message
            .content
            .first()
            .unwrap()
            .as_text()
            .unwrap()
            .text
            .as_str();
        assert!(
            text.contains("test context"),
            "Response should include context for AllServers"
        );
    } else {
        panic!("Expected CreateMessageResult");
    }

    // Test No context inclusion
    let request = ServerRequest::CreateMessageRequest(CreateMessageRequest::new(
        CreateMessageRequestParams::new(vec![SamplingMessage::user_text("test message")], 100)
            .with_include_context(ContextInclusion::None),
    ));

    let result = handler
        .handle_request(
            request.clone(),
            RequestContext::new(NumberOrString::Number(3), client.peer().clone()),
        )
        .await?;

    if let ClientResult::CreateMessageResult(result) = result {
        let text = result
            .message
            .content
            .first()
            .unwrap()
            .as_text()
            .unwrap()
            .text
            .as_str();
        assert!(
            !text.contains("test context"),
            "Response should not include context for None"
        );
    } else {
        panic!("Expected CreateMessageResult");
    }

    client.cancel().await?;
    server_handle.await??;
    Ok(())
}

#[tokio::test]
async fn test_context_inclusion_ignored_integration() -> anyhow::Result<()> {
    let (server_transport, client_transport) = tokio::io::duplex(4096);

    // Start server
    let server_handle = tokio::spawn(async move {
        let server = TestServer::new().serve(server_transport).await?;
        server.waiting().await?;
        anyhow::Ok(())
    });

    // Start client that ignores context requests
    let handler = TestClientHandler::new(false, false);
    let client = handler.clone().serve(client_transport).await?;

    // Test that context requests are ignored
    let request = ServerRequest::CreateMessageRequest(CreateMessageRequest::new(
        CreateMessageRequestParams::new(vec![SamplingMessage::user_text("test message")], 100)
            .with_include_context(ContextInclusion::ThisServer),
    ));

    let result = handler
        .handle_request(
            request.clone(),
            RequestContext::new(NumberOrString::Number(1), client.peer().clone()),
        )
        .await?;

    if let ClientResult::CreateMessageResult(result) = result {
        let text = result
            .message
            .content
            .first()
            .unwrap()
            .as_text()
            .unwrap()
            .text
            .as_str();
        assert!(
            !text.contains("test context"),
            "Context should be ignored when client chooses not to honor requests"
        );
    } else {
        panic!("Expected CreateMessageResult");
    }

    client.cancel().await?;
    server_handle.await??;
    Ok(())
}

#[tokio::test]
async fn test_message_sequence_integration() -> anyhow::Result<()> {
    let (server_transport, client_transport) = tokio::io::duplex(4096);

    // Start server
    let server_handle = tokio::spawn(async move {
        let server = TestServer::new().serve(server_transport).await?;
        server.waiting().await?;
        anyhow::Ok(())
    });

    // Start client
    let handler = TestClientHandler::new(true, true);
    let client = handler.clone().serve(client_transport).await?;

    let request = ServerRequest::CreateMessageRequest(CreateMessageRequest::new(
        CreateMessageRequestParams::new(
            vec![
                SamplingMessage::user_text("first message"),
                SamplingMessage::assistant_text("second message"),
            ],
            100,
        )
        .with_include_context(ContextInclusion::ThisServer),
    ));

    let result = handler
        .handle_request(
            request.clone(),
            RequestContext::new(NumberOrString::Number(1), client.peer().clone()),
        )
        .await?;

    if let ClientResult::CreateMessageResult(result) = result {
        let text = result
            .message
            .content
            .first()
            .unwrap()
            .as_text()
            .unwrap()
            .text
            .as_str();
        assert!(
            text.contains("test context"),
            "Response should include context when ThisServer is specified"
        );
        assert_eq!(result.model, "test-model");
        assert_eq!(
            result.stop_reason,
            Some(CreateMessageResult::STOP_REASON_END_TURN.to_string())
        );
    } else {
        panic!("Expected CreateMessageResult");
    }

    client.cancel().await?;
    server_handle.await??;
    Ok(())
}

#[tokio::test]
async fn test_message_sequence_validation_integration() -> anyhow::Result<()> {
    let (server_transport, client_transport) = tokio::io::duplex(4096);

    let server_handle = tokio::spawn(async move {
        let server = TestServer::new().serve(server_transport).await?;
        server.waiting().await?;
        anyhow::Ok(())
    });

    let handler = TestClientHandler::new(true, true);
    let client = handler.clone().serve(client_transport).await?;

    // Test valid sequence: User -> Assistant -> User
    let request = ServerRequest::CreateMessageRequest(CreateMessageRequest::new(
        CreateMessageRequestParams::new(
            vec![
                SamplingMessage::user_text("first user message"),
                SamplingMessage::assistant_text("first assistant response"),
                SamplingMessage::user_text("second user message"),
            ],
            100,
        ),
    ));

    let result = handler
        .handle_request(
            request.clone(),
            RequestContext::new(NumberOrString::Number(1), client.peer().clone()),
        )
        .await?;

    assert!(matches!(result, ClientResult::CreateMessageResult(_)));

    // Test invalid: No user message
    let request = ServerRequest::CreateMessageRequest(CreateMessageRequest::new(
        CreateMessageRequestParams::new(
            vec![SamplingMessage::assistant_text("assistant message")],
            100,
        ),
    ));

    let result = handler
        .handle_request(
            request.clone(),
            RequestContext::new(NumberOrString::Number(2), client.peer().clone()),
        )
        .await;

    assert!(result.is_err());

    client.cancel().await?;
    server_handle.await??;
    Ok(())
}

#[tokio::test]
async fn test_selective_context_handling_integration() -> anyhow::Result<()> {
    let (server_transport, client_transport) = tokio::io::duplex(4096);

    let server_handle = tokio::spawn(async move {
        let server = TestServer::new().serve(server_transport).await?;
        server.waiting().await?;
        anyhow::Ok(())
    });

    // Client that only honors ThisServer but ignores AllServers
    let handler = TestClientHandler::new(true, false);
    let client = handler.clone().serve(client_transport).await?;

    // Test ThisServer is honored
    let request = ServerRequest::CreateMessageRequest(CreateMessageRequest::new(
        CreateMessageRequestParams::new(vec![SamplingMessage::user_text("test message")], 100)
            .with_include_context(ContextInclusion::ThisServer),
    ));

    let result = handler
        .handle_request(
            request.clone(),
            RequestContext::new(NumberOrString::Number(1), client.peer().clone()),
        )
        .await?;

    if let ClientResult::CreateMessageResult(result) = result {
        let text = result
            .message
            .content
            .first()
            .unwrap()
            .as_text()
            .unwrap()
            .text
            .as_str();
        assert!(
            text.contains("test context"),
            "ThisServer context request should be honored"
        );
    }

    // Test AllServers is ignored
    let request = ServerRequest::CreateMessageRequest(CreateMessageRequest::new(
        CreateMessageRequestParams::new(vec![SamplingMessage::user_text("test message")], 100)
            .with_include_context(ContextInclusion::AllServers),
    ));

    let result = handler
        .handle_request(
            request.clone(),
            RequestContext::new(NumberOrString::Number(2), client.peer().clone()),
        )
        .await?;

    if let ClientResult::CreateMessageResult(result) = result {
        let text = result
            .message
            .content
            .first()
            .unwrap()
            .as_text()
            .unwrap()
            .text
            .as_str();
        assert!(
            !text.contains("test context"),
            "AllServers context request should be ignored"
        );
    }

    client.cancel().await?;
    server_handle.await??;
    Ok(())
}

#[tokio::test]
async fn test_context_inclusion() -> anyhow::Result<()> {
    let (server_transport, client_transport) = tokio::io::duplex(4096);
    let server_handle = tokio::spawn(async move {
        let server = TestServer::new().serve(server_transport).await?;
        server.waiting().await?;
        anyhow::Ok(())
    });

    let handler = TestClientHandler::new(true, true);
    let client = handler.clone().serve(client_transport).await?;

    // Test context handling
    let request = ServerRequest::CreateMessageRequest(CreateMessageRequest::new(
        CreateMessageRequestParams::new(vec![SamplingMessage::user_text("test")], 100)
            .with_include_context(ContextInclusion::ThisServer),
    ));

    let result = handler
        .handle_request(
            request.clone(),
            RequestContext::new(NumberOrString::Number(1), client.peer().clone()),
        )
        .await?;

    if let ClientResult::CreateMessageResult(result) = result {
        let text = result
            .message
            .content
            .first()
            .unwrap()
            .as_text()
            .unwrap()
            .text
            .as_str();
        assert!(text.contains("test context"));
    }

    client.cancel().await?;
    server_handle.await??;
    Ok(())
}