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

#![allow(clippy::exhaustive_structs)]
//cargo test --test test_structured_output --features "client server macros"
use rmcp::{
    Json, ServerHandler,
    handler::server::{router::tool::ToolRouter, tool::IntoCallToolResult, wrapper::Parameters},
    model::{CallToolResponse, CallToolResult, ContentBlock, ServerResult, Tool},
    tool, tool_handler, tool_router,
};
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};

#[derive(Serialize, Deserialize, JsonSchema)]
pub struct CalculationRequest {
    pub a: i32,
    pub b: i32,
}

#[derive(Serialize, Deserialize, JsonSchema)]
pub struct CalculationResult {
    pub sum: i32,
    pub product: i32,
}

#[derive(Serialize, Deserialize, JsonSchema)]
pub struct UserInfo {
    pub name: String,
    pub age: u32,
}

#[derive(Serialize, Deserialize, JsonSchema)]
pub struct GreetingRequest {
    pub name: String,
}

#[derive(Serialize, Deserialize, JsonSchema)]
pub struct GetUserRequest {
    pub user_id: String,
}

#[tool_handler(router = self.tool_router)]
impl ServerHandler for TestServer {}

#[derive(Debug, Clone)]
pub struct TestServer {
    tool_router: ToolRouter<Self>,
}

impl Default for TestServer {
    fn default() -> Self {
        Self::new()
    }
}

#[tool_router(router = tool_router)]
impl TestServer {
    pub fn new() -> Self {
        Self {
            tool_router: Self::tool_router(),
        }
    }

    /// Tool that returns structured output
    #[tool(name = "calculate", description = "Perform calculations")]
    pub async fn calculate(
        &self,
        params: Parameters<CalculationRequest>,
    ) -> Result<Json<CalculationResult>, String> {
        Ok(Json(CalculationResult {
            sum: params.0.a + params.0.b,
            product: params.0.a * params.0.b,
        }))
    }

    /// Tool that returns regular string output
    #[tool(name = "get-greeting", description = "Get a greeting")]
    pub async fn get_greeting(&self, params: Parameters<GreetingRequest>) -> String {
        format!("Hello, {}!", params.0.name)
    }

    /// Tool that returns structured user info
    #[tool(name = "get-user", description = "Get user info")]
    pub async fn get_user(
        &self,
        params: Parameters<GetUserRequest>,
    ) -> Result<Json<UserInfo>, String> {
        if params.0.user_id == "123" {
            Ok(Json(UserInfo {
                name: "Alice".to_string(),
                age: 30,
            }))
        } else {
            Err("User not found".to_string())
        }
    }

    /// Tool that returns a list of calculation results
    #[tool(
        name = "calculate-list",
        description = "Return a list of calculation results"
    )]
    pub async fn calculate_list(
        &self,
        params: Parameters<CalculationRequest>,
    ) -> Result<Json<Vec<CalculationResult>>, String> {
        Ok(Json(vec![CalculationResult {
            sum: params.0.a + params.0.b,
            product: params.0.a * params.0.b,
        }]))
    }

    /// Tool that returns a count
    #[tool(name = "get-count", description = "Return a count")]
    pub async fn get_count(&self) -> Result<Json<i32>, String> {
        Ok(Json(42))
    }
}

#[tokio::test]
async fn test_tool_with_output_schema() {
    let server = TestServer::new();
    let tools = server.tool_router.list_all();

    // Find the calculate tool
    let calculate_tool = tools.iter().find(|t| t.name == "calculate").unwrap();

    // Verify it has an output schema
    assert!(calculate_tool.output_schema.is_some());

    let schema = calculate_tool.output_schema.as_ref().unwrap();

    // Check that the schema contains expected fields
    let schema_str = serde_json::to_string(schema).unwrap();
    assert!(schema_str.contains("sum"));
    assert!(schema_str.contains("product"));
}

#[tokio::test]
async fn test_tool_without_output_schema() {
    let server = TestServer::new();
    let tools = server.tool_router.list_all();

    // Find the get-greeting tool
    let greeting_tool = tools.iter().find(|t| t.name == "get-greeting").unwrap();

    // Verify it doesn't have an output schema (returns String)
    assert!(greeting_tool.output_schema.is_none());
}

#[tokio::test]
async fn test_structured_content_in_call_result() {
    // Test creating a CallToolResult with structured content
    let structured_data = json!({
        "sum": 7,
        "product": 12
    });

    let result = CallToolResult::structured(structured_data.clone());

    assert!(!result.content.is_empty());
    assert!(result.structured_content.is_some());

    let contents = result.content;

    assert_eq!(contents.len(), 1);

    let content_text = contents.first().unwrap().as_text();

    assert!(content_text.is_some());

    let content_value: Value = serde_json::from_str(&content_text.unwrap().text).unwrap();

    assert_eq!(content_value, structured_data);
    assert_eq!(result.structured_content.unwrap(), structured_data);
    assert_eq!(result.is_error, Some(false));
}

#[tokio::test]
async fn test_structured_error_in_call_result() {
    // Test creating a CallToolResult with structured error
    let error_data = json!({
        "error_code": "NOT_FOUND",
        "message": "User not found"
    });

    let result = CallToolResult::structured_error(error_data.clone());

    assert!(!result.content.is_empty());
    assert!(result.structured_content.is_some());

    let contents = result.content;

    assert_eq!(contents.len(), 1);

    let content_text = contents.first().unwrap().as_text();

    assert!(content_text.is_some());

    let content_value: Value = serde_json::from_str(&content_text.unwrap().text).unwrap();

    assert_eq!(content_value, error_data);
    assert_eq!(result.structured_content.unwrap(), error_data);
    assert_eq!(result.is_error, Some(true));
}

#[tokio::test]
async fn test_mutual_exclusivity_validation() {
    #[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)]
    pub struct Response {
        message: String,
    }
    let response = Response {
        message: "Hello".into(),
    };
    // Test that content and structured_content can both be passed separately
    let content_result =
        CallToolResult::success(vec![ContentBlock::json(response.clone()).unwrap()]);
    let structured_result = CallToolResult::structured(json!({"message": "Hello"}));

    // Verify the validation
    content_result
        .into_typed::<Response>()
        .expect("Failed to extract content");
    structured_result
        .into_typed::<Response>()
        .expect("Failed to extract content");

    // Try to create a result with both fields
    let json_with_both = json!({
        "content": [{"type": "text", "text": "Hello"}],
        "structuredContent": {"message": "Hello"}
    });

    // The deserialization itself should not fail
    let deserialized: Result<CallToolResult, _> = serde_json::from_value(json_with_both);
    assert!(deserialized.is_ok());
}

#[tokio::test]
async fn test_structured_return_conversion() {
    // Test that Json<T> converts to CallToolResult with structured_content
    let calc_result = CalculationResult {
        sum: 7,
        product: 12,
    };

    let structured = Json(calc_result);
    let result: Result<CallToolResponse, rmcp::ErrorData> =
        rmcp::handler::server::tool::IntoCallToolResult::into_call_tool_result(structured);

    assert!(result.is_ok());
    let CallToolResponse::Complete(call_result) = result.unwrap() else {
        panic!("expected complete CallToolResult");
    };

    // Tools which return structured content should also return a serialized version as
    // Content::text for backwards compatibility.
    assert!(!call_result.content.is_empty());
    assert!(call_result.structured_content.is_some());

    let contents = call_result.content;

    assert_eq!(contents.len(), 1);

    let content_text = contents.first().unwrap().as_text();

    assert!(content_text.is_some());

    let content_value: Value = serde_json::from_str(&content_text.unwrap().text).unwrap();
    let structured_value = call_result.structured_content.unwrap();

    assert_eq!(content_value, structured_value);

    assert_eq!(structured_value["sum"], 7);
    assert_eq!(structured_value["product"], 12);
}

#[tokio::test]
async fn test_tool_serialization_with_output_schema() {
    let server = TestServer::new();
    let tools = server.tool_router.list_all();

    let calculate_tool = tools.iter().find(|t| t.name == "calculate").unwrap();

    // Serialize the tool
    let serialized = serde_json::to_value(calculate_tool).unwrap();

    // Check that outputSchema is included
    assert!(serialized["outputSchema"].is_object());

    // Deserialize back
    let deserialized: Tool = serde_json::from_value(serialized).unwrap();
    assert!(deserialized.output_schema.is_some());
}

#[tokio::test]
async fn test_output_schema_requires_structured_content() {
    // Test that tools with output_schema must use structured_content
    let server = TestServer::new();
    let tools = server.tool_router.list_all();

    // The calculate tool should have output_schema
    let calculate_tool = tools.iter().find(|t| t.name == "calculate").unwrap();
    assert!(calculate_tool.output_schema.is_some());

    // Directly call the tool and verify its result structure
    let params = Parameters(CalculationRequest { a: 5, b: 3 });
    let result = server.calculate(params).await.unwrap();

    // Convert the Json<CalculationResult> to CallToolResult
    let call_result: Result<CallToolResponse, rmcp::ErrorData> =
        IntoCallToolResult::into_call_tool_result(result);

    assert!(call_result.is_ok());
    let CallToolResponse::Complete(call_result) = call_result.unwrap() else {
        panic!("expected complete CallToolResult");
    };

    // Verify it has structured_content and content
    assert!(call_result.structured_content.is_some());
    assert!(!call_result.content.is_empty());
}

#[tokio::test]
async fn test_empty_content_array_deserializes() {
    let raw = json!({ "content": [] });
    let result: CallToolResult = serde_json::from_value(raw).unwrap();
    assert!(result.content.is_empty());
    assert!(result.structured_content.is_none());
    assert!(result.is_error.is_none());
}

#[tokio::test]
async fn test_empty_content_array_with_is_error() {
    let raw = json!({ "content": [], "isError": false });
    let result: CallToolResult = serde_json::from_value(raw).unwrap();
    assert!(result.content.is_empty());
    assert_eq!(result.is_error, Some(false));
}

#[test]
fn test_missing_content_defaults_to_empty() {
    let raw = json!({ "isError": false });
    let result: CallToolResult = serde_json::from_value(raw).unwrap();
    assert!(result.content.is_empty());
    assert_eq!(result.is_error, Some(false));
}

#[test]
fn test_missing_content_with_structured_content_deserializes() {
    let raw = json!({ "structuredContent": {"key": "value"}, "isError": false });
    let result: CallToolResult = serde_json::from_value(raw).unwrap();
    assert!(result.content.is_empty());
    assert_eq!(result.structured_content.unwrap()["key"], "value");
}

#[tokio::test]
async fn test_empty_content_deserializes_as_call_tool_result_variant() {
    let raw = json!({ "content": [] });
    let result: ServerResult = serde_json::from_value(raw).unwrap();
    match result {
        ServerResult::CallToolResult(call_result) => {
            assert!(call_result.content.is_empty());
            assert!(call_result.structured_content.is_none());
        }
        other => panic!("Expected CallToolResult, got {:?}", other),
    }
}

#[tokio::test]
async fn test_empty_content_roundtrip() {
    let result = CallToolResult::success(vec![]);
    let v = serde_json::to_value(&result).unwrap();
    assert_eq!(v["content"], json!([]));
    let deserialized: CallToolResult = serde_json::from_value(v).unwrap();
    assert_eq!(deserialized, result);
}

#[test]
fn test_call_tool_result_deserialize_without_content() {
    let json = json!({
        "structuredContent": {"message": "Hello"}
    });
    let result: CallToolResult = serde_json::from_value(json).unwrap();
    assert!(result.content.is_empty());
    assert!(result.structured_content.is_some());
}

#[tokio::test]
async fn test_tool_with_array_output_schema() {
    let server = TestServer::new();
    let tools = server.tool_router.list_all();

    // Find the calculate-list tool
    let calculate_list_tool = tools.iter().find(|t| t.name == "calculate-list").unwrap();

    // Verify it has an output schema
    assert!(calculate_list_tool.output_schema.is_some());

    let schema = calculate_list_tool.output_schema.as_ref().unwrap();

    // Check that the schema contains array type
    let schema_str = serde_json::to_string(schema).unwrap();
    assert!(schema_str.contains("array"));
}

#[tokio::test]
async fn test_tool_with_primitive_output_schema() {
    let server = TestServer::new();
    let tools = server.tool_router.list_all();

    // Find the get-count tool
    let get_count_tool = tools.iter().find(|t| t.name == "get-count").unwrap();

    // Verify it has an output schema
    assert!(get_count_tool.output_schema.is_some());

    let schema = get_count_tool.output_schema.as_ref().unwrap();

    // Check that the schema contains integer type
    assert_eq!(schema.get("type"), Some(&serde_json::json!("integer")));
}