diff --git a/.changesets/feat_output_schema.md b/.changesets/feat_output_schema.md new file mode 100644 index 00000000..953c86d7 --- /dev/null +++ b/.changesets/feat_output_schema.md @@ -0,0 +1,10 @@ +### Add outputSchema support - @DaleSeo PR #509 + +This PR implements support for the MCP specification's [outputSchema](https://modelcontextprotocol.io/specification/2025-11-25/server/tools#output-schema) field on tools, which allows tools to declare the expected structure of their output. This helps LLMs better understand and reason about GraphQL response data. + +This feature is opt-in to avoid additional token overhead. To enable it, add the following to your config: + +```yaml +overrides: + enable_output_schema: true +``` diff --git a/crates/apollo-mcp-server/src/apps.rs b/crates/apollo-mcp-server/src/apps.rs index c596e7c1..b6ad0ab1 100644 --- a/crates/apollo-mcp-server/src/apps.rs +++ b/crates/apollo-mcp-server/src/apps.rs @@ -86,6 +86,7 @@ pub(crate) fn load_from_path( mutation_mode: MutationMode, disable_type_description: bool, disable_schema_description: bool, + enable_output_schema: bool, ) -> Result, String> { let Ok(apps_dir) = path.read_dir() else { return Ok(Vec::new()); @@ -146,6 +147,7 @@ pub(crate) fn load_from_path( mutation_mode, disable_type_description, disable_schema_description, + enable_output_schema, ) { Err(err) => { return Err(format!( @@ -393,6 +395,7 @@ mod test_load_from_path { MutationMode::All, false, false, + true, ) .expect("Failed to load apps"); assert_eq!(apps.len(), 1); @@ -428,6 +431,7 @@ mod test_load_from_path { MutationMode::All, false, false, + true, ) .expect("Failed to load apps"); assert_eq!(apps.len(), 1); @@ -477,6 +481,7 @@ mod test_load_from_path { MutationMode::All, false, false, + true, ) .expect("Failed to load apps"); assert_eq!(apps.len(), 1); @@ -536,6 +541,7 @@ mod test_load_from_path { MutationMode::All, false, false, + true, ) .expect("Failed to load apps"); assert_eq!(apps.len(), 1); @@ -592,6 +598,7 @@ mod test_load_from_path { MutationMode::All, false, false, + true, ) .expect("Failed to load apps"); @@ -655,6 +662,7 @@ mod test_load_from_path { MutationMode::All, false, false, + true, ); assert!(apps.is_err()); @@ -701,6 +709,7 @@ mod test_load_from_path { MutationMode::All, false, false, + true, ); assert!(apps.is_err()); diff --git a/crates/apollo-mcp-server/src/apps/execution.rs b/crates/apollo-mcp-server/src/apps/execution.rs index 02b8ddc4..d0c1fa53 100644 --- a/crates/apollo-mcp-server/src/apps/execution.rs +++ b/crates/apollo-mcp-server/src/apps/execution.rs @@ -163,7 +163,7 @@ mod tests { "query Primary($apples: Int) { apples(first: $apples) }".to_string(), None, )) - .into_operation(&schema, None, MutationMode::All, true, true) + .into_operation(&schema, None, MutationMode::All, true, true, true) .unwrap() .unwrap(), ); @@ -174,7 +174,7 @@ mod tests { "query FirstPrefetch($bananas: Int) { bananas(first: $bananas) }".to_string(), None, )) - .into_operation(&schema, None, MutationMode::All, true, true) + .into_operation(&schema, None, MutationMode::All, true, true, true) .unwrap() .unwrap(), ); @@ -183,7 +183,7 @@ mod tests { "query SecondPrefetch($oranges: Int) { oranges(first: $oranges) }".to_string(), None, )) - .into_operation(&schema, None, MutationMode::All, true, true) + .into_operation(&schema, None, MutationMode::All, true, true, true) .unwrap() .unwrap(), ); @@ -314,7 +314,7 @@ mod tests { tools: vec![AppTool { operation: Arc::new( RawOperation::from(("query GetId { id }".to_string(), None)) - .into_operation(&schema, None, MutationMode::All, false, false) + .into_operation(&schema, None, MutationMode::All, false, false, true) .unwrap() .unwrap(), ), @@ -361,7 +361,7 @@ mod tests { tools: vec![AppTool { operation: Arc::new( RawOperation::from(("query GetId { id }".to_string(), None)) - .into_operation(&schema, None, MutationMode::All, false, false) + .into_operation(&schema, None, MutationMode::All, false, false, true) .unwrap() .unwrap(), ), @@ -400,7 +400,7 @@ mod tests { tools: vec![AppTool { operation: Arc::new( RawOperation::from(("query GetId { id }".to_string(), None)) - .into_operation(&schema, None, MutationMode::All, false, false) + .into_operation(&schema, None, MutationMode::All, false, false, true) .unwrap() .unwrap(), ), diff --git a/crates/apollo-mcp-server/src/main.rs b/crates/apollo-mcp-server/src/main.rs index 14bb4b59..e339d9a8 100644 --- a/crates/apollo-mcp-server/src/main.rs +++ b/crates/apollo-mcp-server/src/main.rs @@ -128,6 +128,7 @@ async fn main() -> anyhow::Result<()> { .mutation_mode(config.overrides.mutation_mode) .disable_type_description(config.overrides.disable_type_description) .disable_schema_description(config.overrides.disable_schema_description) + .enable_output_schema(config.overrides.enable_output_schema) .disable_auth_token_passthrough(match transport { apollo_mcp_server::server::Transport::Stdio => false, apollo_mcp_server::server::Transport::SSE { auth, .. } => auth diff --git a/crates/apollo-mcp-server/src/operations/execution.rs b/crates/apollo-mcp-server/src/operations/execution.rs index 8a9a113c..bd837450 100644 --- a/crates/apollo-mcp-server/src/operations/execution.rs +++ b/crates/apollo-mcp-server/src/operations/execution.rs @@ -54,7 +54,7 @@ mod tests { .validate() .unwrap(); let operation = RawOperation::from(("query GetHello { hello }".to_string(), None)) - .into_operation(&schema, None, MutationMode::All, true, true) + .into_operation(&schema, None, MutationMode::All, true, true, true) .unwrap() .unwrap(); @@ -78,11 +78,11 @@ mod tests { .unwrap(); let operations = [ RawOperation::from(("query GetHello { hello }".to_string(), None)) - .into_operation(&schema, None, MutationMode::All, true, true) + .into_operation(&schema, None, MutationMode::All, true, true, true) .unwrap() .unwrap(), RawOperation::from(("query GetWorld { hello }".to_string(), None)) - .into_operation(&schema, None, MutationMode::All, true, true) + .into_operation(&schema, None, MutationMode::All, true, true, true) .unwrap() .unwrap(), ]; diff --git a/crates/apollo-mcp-server/src/operations/operation.rs b/crates/apollo-mcp-server/src/operations/operation.rs index 9eed07e7..1e65a029 100644 --- a/crates/apollo-mcp-server/src/operations/operation.rs +++ b/crates/apollo-mcp-server/src/operations/operation.rs @@ -56,6 +56,7 @@ impl Operation { mutation_mode: MutationMode, disable_type_description: bool, disable_schema_description: bool, + enable_output_schema: bool, ) -> Result, OperationError> { if let Some((document, operation, comments)) = operation_defs( &raw_operation.source_text, @@ -111,11 +112,54 @@ impl Operation { )); }; - let tool: Tool = Tool::new(operation_name.clone(), description, schema).annotate( + // Generate output schema from selection set (only if enabled) + let output_schema = if enable_output_schema { + if let Some(root_type_name) = + graphql_schema.root_operation(operation.operation_type) + { + if let Some(root_type) = graphql_schema.types.get(root_type_name) { + let named_fragments: HashMap< + String, + Node, + > = document + .definitions + .iter() + .filter_map(|def| match def { + Definition::FragmentDefinition(fragment_def) => { + Some((fragment_def.name.to_string(), fragment_def.clone())) + } + _ => None, + }) + .collect(); + + serde_json::to_value(schema_walker::selection_set_to_schema( + &operation.selection_set, + root_type, + graphql_schema, + custom_scalar_map, + &named_fragments, + )) + .ok() + .and_then(|v| match v { + Value::Object(obj) => Some(obj), + _ => None, + }) + } else { + None + } + } else { + None + } + } else { + None + }; + + let mut tool: Tool = Tool::new(operation_name.clone(), description, schema).annotate( ToolAnnotations::new() .read_only(operation.operation_type != OperationType::Mutation) .destructive(operation.operation_type == OperationType::Mutation), ); + tool.output_schema = output_schema.map(std::sync::Arc::new); let character_count = tool_character_length(&tool); match character_count { Ok(length) => info!( @@ -512,10 +556,16 @@ fn ensure_properties_exists(json_object: &mut Value) { } fn tool_character_length(tool: &Tool) -> Result { - let tool_schema_string = serde_json::to_string_pretty(&serde_json::json!(tool.input_schema))?; + let input_schema_len = + serde_json::to_string_pretty(&serde_json::json!(tool.input_schema))?.len(); + let output_schema_len = match &tool.output_schema { + Some(schema) => serde_json::to_string_pretty(schema.as_ref())?.len(), + None => 0, + }; Ok(tool.name.len() + tool.description.as_ref().map(|d| d.len()).unwrap_or(0) - + tool_schema_string.len()) + + input_schema_len + + output_schema_len) } #[tracing::instrument(skip_all)] @@ -599,7 +649,7 @@ mod tests { use crate::{ custom_scalar_map::CustomScalarMap, graphql::Executable as _, - operations::{MutationMode, Operation, RawOperation}, + operations::{MutationMode, Operation, RawOperation, operation::tool_character_length}, }; // Example schema for tests @@ -688,6 +738,7 @@ mod tests { MutationMode::None, false, false, + true, ) .unwrap() .unwrap(); @@ -708,7 +759,75 @@ mod tests { }, }, }, - output_schema: None, + output_schema: Some( + { + "type": String("object"), + "properties": Object { + "data": Object { + "type": String("object"), + "properties": Object { + "id": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("null"), + }, + ], + }, + }, + }, + "errors": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "message": Object { + "type": String("string"), + }, + "locations": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "line": Object { + "type": String("integer"), + }, + "column": Object { + "type": String("integer"), + }, + }, + }, + }, + "path": Object { + "type": String("array"), + "items": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("integer"), + }, + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + "required": Array [ + String("message"), + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + }, + ), annotations: Some( ToolAnnotations { title: None, @@ -755,6 +874,7 @@ mod tests { MutationMode::None, false, false, + true, ) .unwrap() .unwrap(); @@ -778,7 +898,75 @@ mod tests { String("id"), ], }, - output_schema: None, + output_schema: Some( + { + "type": String("object"), + "properties": Object { + "data": Object { + "type": String("object"), + "properties": Object { + "id": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("null"), + }, + ], + }, + }, + }, + "errors": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "message": Object { + "type": String("string"), + }, + "locations": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "line": Object { + "type": String("integer"), + }, + "column": Object { + "type": String("integer"), + }, + }, + }, + }, + "path": Object { + "type": String("array"), + "items": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("integer"), + }, + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + "required": Array [ + String("message"), + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + }, + ), annotations: Some( ToolAnnotations { title: None, @@ -826,6 +1014,7 @@ mod tests { MutationMode::None, false, false, + true, ) .unwrap() .unwrap(); @@ -859,7 +1048,75 @@ mod tests { String("id"), ], }, - output_schema: None, + output_schema: Some( + { + "type": String("object"), + "properties": Object { + "data": Object { + "type": String("object"), + "properties": Object { + "id": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("null"), + }, + ], + }, + }, + }, + "errors": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "message": Object { + "type": String("string"), + }, + "locations": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "line": Object { + "type": String("integer"), + }, + "column": Object { + "type": String("integer"), + }, + }, + }, + }, + "path": Object { + "type": String("array"), + "items": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("integer"), + }, + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + "required": Array [ + String("message"), + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + }, + ), annotations: Some( ToolAnnotations { title: None, @@ -917,6 +1174,7 @@ mod tests { MutationMode::None, false, false, + true, ) .unwrap() .unwrap(); @@ -943,7 +1201,75 @@ mod tests { String("id"), ], }, - output_schema: None, + output_schema: Some( + { + "type": String("object"), + "properties": Object { + "data": Object { + "type": String("object"), + "properties": Object { + "id": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("null"), + }, + ], + }, + }, + }, + "errors": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "message": Object { + "type": String("string"), + }, + "locations": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "line": Object { + "type": String("integer"), + }, + "column": Object { + "type": String("integer"), + }, + }, + }, + }, + "path": Object { + "type": String("array"), + "items": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("integer"), + }, + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + "required": Array [ + String("message"), + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + }, + ), annotations: Some( ToolAnnotations { title: None, @@ -994,6 +1320,7 @@ mod tests { MutationMode::None, false, false, + true, ) .unwrap() .unwrap(); @@ -1024,7 +1351,75 @@ mod tests { }, }, }, - output_schema: None, + output_schema: Some( + { + "type": String("object"), + "properties": Object { + "data": Object { + "type": String("object"), + "properties": Object { + "id": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("null"), + }, + ], + }, + }, + }, + "errors": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "message": Object { + "type": String("string"), + }, + "locations": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "line": Object { + "type": String("integer"), + }, + "column": Object { + "type": String("integer"), + }, + }, + }, + }, + "path": Object { + "type": String("array"), + "items": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("integer"), + }, + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + "required": Array [ + String("message"), + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + }, + ), annotations: Some( ToolAnnotations { title: None, @@ -1079,6 +1474,7 @@ mod tests { MutationMode::None, false, false, + true, ) .unwrap() .unwrap(); @@ -1102,7 +1498,75 @@ mod tests { }, }, }, - output_schema: None, + output_schema: Some( + { + "type": String("object"), + "properties": Object { + "data": Object { + "type": String("object"), + "properties": Object { + "id": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("null"), + }, + ], + }, + }, + }, + "errors": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "message": Object { + "type": String("string"), + }, + "locations": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "line": Object { + "type": String("integer"), + }, + "column": Object { + "type": String("integer"), + }, + }, + }, + }, + "path": Object { + "type": String("array"), + "items": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("integer"), + }, + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + "required": Array [ + String("message"), + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + }, + ), annotations: Some( ToolAnnotations { title: None, @@ -1150,6 +1614,7 @@ mod tests { MutationMode::None, false, false, + true, ) .unwrap() .unwrap(); @@ -1190,7 +1655,75 @@ mod tests { }, }, }, - output_schema: None, + output_schema: Some( + { + "type": String("object"), + "properties": Object { + "data": Object { + "type": String("object"), + "properties": Object { + "id": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("null"), + }, + ], + }, + }, + }, + "errors": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "message": Object { + "type": String("string"), + }, + "locations": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "line": Object { + "type": String("integer"), + }, + "column": Object { + "type": String("integer"), + }, + }, + }, + }, + "path": Object { + "type": String("array"), + "items": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("integer"), + }, + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + "required": Array [ + String("message"), + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + }, + ), annotations: Some( ToolAnnotations { title: None, @@ -1255,6 +1788,7 @@ mod tests { MutationMode::None, false, false, + true, ) .unwrap() .unwrap(); @@ -1293,7 +1827,75 @@ mod tests { }, }, }, - output_schema: None, + output_schema: Some( + { + "type": String("object"), + "properties": Object { + "data": Object { + "type": String("object"), + "properties": Object { + "id": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("null"), + }, + ], + }, + }, + }, + "errors": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "message": Object { + "type": String("string"), + }, + "locations": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "line": Object { + "type": String("integer"), + }, + "column": Object { + "type": String("integer"), + }, + }, + }, + }, + "path": Object { + "type": String("array"), + "items": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("integer"), + }, + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + "required": Array [ + String("message"), + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + }, + ), annotations: Some( ToolAnnotations { title: None, @@ -1328,6 +1930,7 @@ mod tests { MutationMode::None, false, false, + true, ) .unwrap() .unwrap(); @@ -1361,7 +1964,75 @@ mod tests { }, }, }, - output_schema: None, + output_schema: Some( + { + "type": String("object"), + "properties": Object { + "data": Object { + "type": String("object"), + "properties": Object { + "id": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("null"), + }, + ], + }, + }, + }, + "errors": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "message": Object { + "type": String("string"), + }, + "locations": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "line": Object { + "type": String("integer"), + }, + "column": Object { + "type": String("integer"), + }, + }, + }, + }, + "path": Object { + "type": String("array"), + "items": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("integer"), + }, + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + "required": Array [ + String("message"), + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + }, + ), annotations: Some( ToolAnnotations { title: None, @@ -1396,6 +2067,7 @@ mod tests { MutationMode::None, false, false, + true, ); insta::assert_debug_snapshot!(operation, @r#" Err( @@ -1425,6 +2097,7 @@ mod tests { MutationMode::None, false, false, + true, ); assert!(operation.unwrap().is_none()); @@ -1455,6 +2128,7 @@ mod tests { MutationMode::None, false, false, + true, ); insta::assert_debug_snapshot!(operation, @r#" Err( @@ -1482,6 +2156,7 @@ mod tests { MutationMode::None, false, false, + true, ); insta::assert_debug_snapshot!(operation, @r" Err( @@ -1508,6 +2183,7 @@ mod tests { MutationMode::None, false, false, + true, ) .unwrap() .unwrap(); @@ -1536,7 +2212,75 @@ mod tests { "id": Object {}, }, }, - output_schema: None, + output_schema: Some( + { + "type": String("object"), + "properties": Object { + "data": Object { + "type": String("object"), + "properties": Object { + "id": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("null"), + }, + ], + }, + }, + }, + "errors": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "message": Object { + "type": String("string"), + }, + "locations": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "line": Object { + "type": String("integer"), + }, + "column": Object { + "type": String("integer"), + }, + }, + }, + }, + "path": Object { + "type": String("array"), + "items": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("integer"), + }, + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + "required": Array [ + String("message"), + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + }, + ), annotations: Some( ToolAnnotations { title: None, @@ -1572,6 +2316,7 @@ mod tests { MutationMode::None, false, false, + true, ) .unwrap() .unwrap(); @@ -1607,7 +2352,75 @@ mod tests { }, }, }, - output_schema: None, + output_schema: Some( + { + "type": String("object"), + "properties": Object { + "data": Object { + "type": String("object"), + "properties": Object { + "id": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("null"), + }, + ], + }, + }, + }, + "errors": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "message": Object { + "type": String("string"), + }, + "locations": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "line": Object { + "type": String("integer"), + }, + "column": Object { + "type": String("integer"), + }, + }, + }, + }, + "path": Object { + "type": String("array"), + "items": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("integer"), + }, + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + "required": Array [ + String("message"), + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + }, + ), annotations: Some( ToolAnnotations { title: None, @@ -1643,6 +2456,7 @@ mod tests { MutationMode::None, false, false, + true, ) .unwrap() .unwrap(); @@ -1682,7 +2496,75 @@ mod tests { }, }, }, - output_schema: None, + output_schema: Some( + { + "type": String("object"), + "properties": Object { + "data": Object { + "type": String("object"), + "properties": Object { + "id": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("null"), + }, + ], + }, + }, + }, + "errors": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "message": Object { + "type": String("string"), + }, + "locations": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "line": Object { + "type": String("integer"), + }, + "column": Object { + "type": String("integer"), + }, + }, + }, + }, + "path": Object { + "type": String("array"), + "items": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("integer"), + }, + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + "required": Array [ + String("message"), + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + }, + ), annotations: Some( ToolAnnotations { title: None, @@ -1720,6 +2602,7 @@ mod tests { MutationMode::None, false, false, + true, ) .unwrap() .unwrap(); @@ -1746,7 +2629,75 @@ mod tests { }, }, }, - output_schema: None, + output_schema: Some( + { + "type": String("object"), + "properties": Object { + "data": Object { + "type": String("object"), + "properties": Object { + "id": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("null"), + }, + ], + }, + }, + }, + "errors": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "message": Object { + "type": String("string"), + }, + "locations": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "line": Object { + "type": String("integer"), + }, + "column": Object { + "type": String("integer"), + }, + }, + }, + }, + "path": Object { + "type": String("array"), + "items": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("integer"), + }, + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + "required": Array [ + String("message"), + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + }, + ), annotations: Some( ToolAnnotations { title: None, @@ -1908,6 +2859,7 @@ mod tests { MutationMode::None, false, false, + true, ) .unwrap() .unwrap(); @@ -1995,6 +2947,7 @@ mod tests { MutationMode::None, false, false, + true, ) .unwrap() .unwrap(); @@ -2028,6 +2981,7 @@ mod tests { MutationMode::None, false, false, + true, ) .unwrap() .unwrap(); @@ -2053,6 +3007,7 @@ mod tests { MutationMode::None, false, true, + true, ) .unwrap() .unwrap(); @@ -2082,6 +3037,7 @@ mod tests { MutationMode::None, true, false, + true, ) .unwrap() .unwrap(); @@ -2116,6 +3072,7 @@ mod tests { MutationMode::None, true, true, + true, ) .unwrap() .unwrap(); @@ -2165,6 +3122,7 @@ mod tests { MutationMode::None, true, true, + true, ) .unwrap() .unwrap(); @@ -2201,7 +3159,76 @@ mod tests { }, }, }, - output_schema: None, + output_schema: Some( + { + "type": String("object"), + "properties": Object { + "data": Object { + "type": String("object"), + "properties": Object { + "field": Object { + "description": String("the Query.field field"), + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("null"), + }, + ], + }, + }, + }, + "errors": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "message": Object { + "type": String("string"), + }, + "locations": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "line": Object { + "type": String("integer"), + }, + "column": Object { + "type": String("integer"), + }, + }, + }, + }, + "path": Object { + "type": String("array"), + "items": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("integer"), + }, + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + "required": Array [ + String("message"), + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + }, + ), annotations: Some( ToolAnnotations { title: None, @@ -2239,6 +3266,7 @@ mod tests { MutationMode::None, false, false, + true, ) .unwrap() .unwrap(); @@ -2259,7 +3287,75 @@ mod tests { }, }, }, - output_schema: None, + output_schema: Some( + { + "type": String("object"), + "properties": Object { + "data": Object { + "type": String("object"), + "properties": Object { + "id": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("null"), + }, + ], + }, + }, + }, + "errors": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "message": Object { + "type": String("string"), + }, + "locations": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "line": Object { + "type": String("integer"), + }, + "column": Object { + "type": String("integer"), + }, + }, + }, + }, + "path": Object { + "type": String("array"), + "items": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("integer"), + }, + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + "required": Array [ + String("message"), + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + }, + ), annotations: Some( ToolAnnotations { title: None, @@ -2295,6 +3391,7 @@ mod tests { MutationMode::None, false, false, + true, ) .unwrap() .unwrap(); @@ -2329,6 +3426,7 @@ mod tests { MutationMode::None, false, false, + true, ) .unwrap() .unwrap(); @@ -2367,6 +3465,7 @@ mod tests { MutationMode::None, false, false, + true, ) .unwrap() .unwrap(); @@ -2399,10 +3498,17 @@ mod tests { variables: None, source_path: None, }; - let operation = - Operation::from_document(raw_op, &SCHEMA, None, MutationMode::None, false, false) - .unwrap() - .unwrap(); + let operation = Operation::from_document( + raw_op, + &SCHEMA, + None, + MutationMode::None, + false, + false, + true, + ) + .unwrap() + .unwrap(); let op_details = operation.operation(Value::Null).unwrap(); assert_eq!(op_details.operation_name, Some(String::from("GetUser"))); @@ -2419,10 +3525,17 @@ mod tests { variables: None, source_path: None, }; - let operation = - Operation::from_document(raw_op, &SCHEMA, None, MutationMode::Explicit, false, false) - .unwrap() - .unwrap(); + let operation = Operation::from_document( + raw_op, + &SCHEMA, + None, + MutationMode::Explicit, + false, + false, + true, + ) + .unwrap() + .unwrap(); let op_details = operation.operation(Value::Null).unwrap(); assert_eq!(op_details.operation_name, Some(String::from("CreateUser"))); @@ -2443,6 +3556,7 @@ mod tests { MutationMode::None, false, false, + true, ) .unwrap() .unwrap(); @@ -2477,6 +3591,7 @@ mod tests { MutationMode::None, false, false, + true, ) .unwrap() .unwrap(); @@ -2511,6 +3626,7 @@ mod tests { MutationMode::None, false, false, + true, ) .unwrap() .unwrap(); @@ -2545,6 +3661,7 @@ mod tests { MutationMode::None, false, false, + true, ) .unwrap() .unwrap(); @@ -2583,6 +3700,7 @@ mod tests { MutationMode::None, false, false, + true, ) .unwrap() .unwrap(); @@ -2612,6 +3730,7 @@ mod tests { MutationMode::None, false, false, + true, ) .unwrap() .unwrap(); @@ -2650,6 +3769,7 @@ mod tests { MutationMode::None, false, false, + true, ) .unwrap() .unwrap(); @@ -2679,6 +3799,7 @@ mod tests { MutationMode::None, false, false, + true, ) .unwrap() .unwrap(); @@ -2727,7 +3848,75 @@ mod tests { }, }, }, - output_schema: None, + output_schema: Some( + { + "type": String("object"), + "properties": Object { + "data": Object { + "type": String("object"), + "properties": Object { + "id": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("null"), + }, + ], + }, + }, + }, + "errors": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "message": Object { + "type": String("string"), + }, + "locations": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "line": Object { + "type": String("integer"), + }, + "column": Object { + "type": String("integer"), + }, + }, + }, + }, + "path": Object { + "type": String("array"), + "items": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("integer"), + }, + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + "required": Array [ + String("message"), + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + }, + ), annotations: Some( ToolAnnotations { title: None, @@ -2802,6 +3991,7 @@ mod tests { MutationMode::None, false, false, + true, ) .unwrap() .unwrap(); @@ -2846,7 +4036,75 @@ mod tests { }, }, }, - output_schema: None, + output_schema: Some( + { + "type": String("object"), + "properties": Object { + "data": Object { + "type": String("object"), + "properties": Object { + "id": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("null"), + }, + ], + }, + }, + }, + "errors": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "message": Object { + "type": String("string"), + }, + "locations": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "line": Object { + "type": String("integer"), + }, + "column": Object { + "type": String("integer"), + }, + }, + }, + }, + "path": Object { + "type": String("array"), + "items": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("integer"), + }, + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + "required": Array [ + String("message"), + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + }, + ), annotations: Some( ToolAnnotations { title: None, @@ -2918,6 +4176,7 @@ mod tests { MutationMode::None, false, false, + true, ) .unwrap() .is_none() @@ -2940,6 +4199,7 @@ mod tests { MutationMode::None, false, false, + true, ) .ok() .unwrap() @@ -2962,6 +4222,7 @@ mod tests { MutationMode::Explicit, false, false, + true, ) .unwrap() .unwrap(); @@ -2978,7 +4239,75 @@ mod tests { "type": String("object"), "properties": Object {}, }, - output_schema: None, + output_schema: Some( + { + "type": String("object"), + "properties": Object { + "data": Object { + "type": String("object"), + "properties": Object { + "id": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("null"), + }, + ], + }, + }, + }, + "errors": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "message": Object { + "type": String("string"), + }, + "locations": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "line": Object { + "type": String("integer"), + }, + "column": Object { + "type": String("integer"), + }, + }, + }, + }, + "path": Object { + "type": String("array"), + "items": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("integer"), + }, + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + "required": Array [ + String("message"), + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + }, + ), annotations: Some( ToolAnnotations { title: None, @@ -3022,6 +4351,7 @@ mod tests { MutationMode::All, false, false, + true, ) .unwrap() .unwrap(); @@ -3038,7 +4368,75 @@ mod tests { "type": String("object"), "properties": Object {}, }, - output_schema: None, + output_schema: Some( + { + "type": String("object"), + "properties": Object { + "data": Object { + "type": String("object"), + "properties": Object { + "id": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("null"), + }, + ], + }, + }, + }, + "errors": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "message": Object { + "type": String("string"), + }, + "locations": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "line": Object { + "type": String("integer"), + }, + "column": Object { + "type": String("integer"), + }, + }, + }, + }, + "path": Object { + "type": String("array"), + "items": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("integer"), + }, + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + "required": Array [ + String("message"), + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + }, + ), annotations: Some( ToolAnnotations { title: None, @@ -3082,6 +4480,7 @@ mod tests { MutationMode::None, false, false, + true, ) .unwrap() .unwrap(); @@ -3098,7 +4497,75 @@ mod tests { "type": String("object"), "properties": Object {}, }, - output_schema: None, + output_schema: Some( + { + "type": String("object"), + "properties": Object { + "data": Object { + "type": String("object"), + "properties": Object { + "id": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("null"), + }, + ], + }, + }, + }, + "errors": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "message": Object { + "type": String("string"), + }, + "locations": Object { + "type": String("array"), + "items": Object { + "type": String("object"), + "properties": Object { + "line": Object { + "type": String("integer"), + }, + "column": Object { + "type": String("integer"), + }, + }, + }, + }, + "path": Object { + "type": String("array"), + "items": Object { + "oneOf": Array [ + Object { + "type": String("string"), + }, + Object { + "type": String("integer"), + }, + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + "required": Array [ + String("message"), + ], + }, + }, + "extensions": Object { + "type": String("object"), + }, + }, + }, + ), annotations: Some( ToolAnnotations { title: None, @@ -3123,4 +4590,108 @@ mod tests { } "#); } + + #[test] + fn tool_character_length_without_output_schema() { + use serde_json::Map; + + let mut input_schema = Map::new(); + input_schema.insert("type".to_string(), serde_json::json!("object")); + input_schema.insert( + "properties".to_string(), + serde_json::json!({ "id": { "type": "string" } }), + ); + + let tool = Tool::new("test_tool", "A test tool description", input_schema.clone()); + + let length = tool_character_length(&tool).unwrap(); + + let expected_input_schema_len = + serde_json::to_string_pretty(&serde_json::json!(input_schema)) + .unwrap() + .len(); + + assert_eq!( + length, + "test_tool".len() + "A test tool description".len() + expected_input_schema_len + ); + } + + #[test] + fn tool_character_length_with_output_schema() { + use serde_json::Map; + + let mut input_schema = Map::new(); + input_schema.insert("type".to_string(), serde_json::json!("object")); + input_schema.insert( + "properties".to_string(), + serde_json::json!({ "id": { "type": "string" } }), + ); + + let mut tool = Tool::new("test_tool", "A test tool description", input_schema.clone()); + + let mut output_schema = Map::new(); + output_schema.insert("type".to_string(), serde_json::json!("object")); + output_schema.insert( + "properties".to_string(), + serde_json::json!({ + "data": { + "type": "object", + "properties": { + "result": { "type": "string" } + } + } + }), + ); + tool.output_schema = Some(std::sync::Arc::new(output_schema.clone())); + + let length = tool_character_length(&tool).unwrap(); + + let expected_input_schema_len = + serde_json::to_string_pretty(&serde_json::json!(input_schema)) + .unwrap() + .len(); + + let expected_output_schema_len = + serde_json::to_string_pretty(&serde_json::json!(output_schema)) + .unwrap() + .len(); + + assert_eq!( + length, + "test_tool".len() + + "A test tool description".len() + + expected_input_schema_len + + expected_output_schema_len + ); + } + + #[test] + fn tool_character_length_output_schema_adds_to_total() { + use serde_json::Map; + + let mut input_schema = Map::new(); + input_schema.insert("type".to_string(), serde_json::json!("object")); + + let tool_without = Tool::new("test_tool", "A test tool description", input_schema.clone()); + + let mut tool_with = tool_without.clone(); + let mut output_schema = Map::new(); + output_schema.insert("type".to_string(), serde_json::json!("object")); + output_schema.insert( + "properties".to_string(), + serde_json::json!({ "data": { "type": "string" } }), + ); + tool_with.output_schema = Some(std::sync::Arc::new(output_schema.clone())); + + let length_without = tool_character_length(&tool_without).unwrap(); + let length_with = tool_character_length(&tool_with).unwrap(); + + let expected_output_schema_len = + serde_json::to_string_pretty(&serde_json::json!(output_schema)) + .unwrap() + .len(); + + assert_eq!(length_with - length_without, expected_output_schema_len); + } } diff --git a/crates/apollo-mcp-server/src/operations/raw_operation.rs b/crates/apollo-mcp-server/src/operations/raw_operation.rs index 85d25da8..c07902aa 100644 --- a/crates/apollo-mcp-server/src/operations/raw_operation.rs +++ b/crates/apollo-mcp-server/src/operations/raw_operation.rs @@ -28,6 +28,7 @@ impl RawOperation { mutation_mode: MutationMode, disable_type_description: bool, disable_schema_description: bool, + enable_output_schema: bool, ) -> Result, OperationError> { Operation::from_document( self, @@ -36,6 +37,7 @@ impl RawOperation { mutation_mode, disable_type_description, disable_schema_description, + enable_output_schema, ) } } diff --git a/crates/apollo-mcp-server/src/operations/schema_walker.rs b/crates/apollo-mcp-server/src/operations/schema_walker.rs index 0229d255..d9bb259a 100644 --- a/crates/apollo-mcp-server/src/operations/schema_walker.rs +++ b/crates/apollo-mcp-server/src/operations/schema_walker.rs @@ -10,8 +10,11 @@ use serde_json::{Map, Value}; use crate::custom_scalar_map::CustomScalarMap; mod name; +mod output; mod r#type; +pub use output::selection_set_to_schema; + /// Convert a GraphQL type into a JSON Schema. /// /// Note: This is recursive, which might cause a stack overflow if the type is diff --git a/crates/apollo-mcp-server/src/operations/schema_walker/output.rs b/crates/apollo-mcp-server/src/operations/schema_walker/output.rs new file mode 100644 index 00000000..3af40a64 --- /dev/null +++ b/crates/apollo-mcp-server/src/operations/schema_walker/output.rs @@ -0,0 +1,609 @@ +//! JSON Schema generation for GraphQL output types (selection sets) +//! +//! This module generates JSON schemas from GraphQL operation selection sets, +//! enabling MCP tools to declare their output schema. + +use std::collections::HashMap; + +use apollo_compiler::{ + Name as GraphQLName, Node, Schema as GraphQLSchema, + ast::{Field, Selection, Type as GraphQLType}, + schema::ExtendedType, +}; +use schemars::{Schema as JSONSchema, json_schema}; +use serde_json::{Map, Value}; +use tracing::warn; + +use crate::custom_scalar_map::CustomScalarMap; + +/// Generate a JSON Schema for the output of a GraphQL operation. +/// +/// This walks the selection set and generates a schema that describes +/// the expected response structure. +pub fn selection_set_to_schema( + selection_set: &[Selection], + parent_type: &ExtendedType, + graphql_schema: &GraphQLSchema, + custom_scalar_map: Option<&CustomScalarMap>, + named_fragments: &HashMap>, +) -> JSONSchema { + let mut definitions = Map::new(); + + let schema = build_selection_set_schema( + selection_set, + parent_type, + graphql_schema, + custom_scalar_map, + named_fragments, + &mut definitions, + ); + + // Wrap in standard GraphQL response envelope + let mut response_schema = json_schema!({ + "type": "object", + "properties": { + "data": schema, + "errors": { + "type": "array", + "items": { + "type": "object", + "properties": { + "message": { "type": "string" }, + "locations": { + "type": "array", + "items": { + "type": "object", + "properties": { + "line": { "type": "integer" }, + "column": { "type": "integer" } + } + } + }, + "path": { + "type": "array", + "items": { + "oneOf": [ + { "type": "string" }, + { "type": "integer" } + ] + } + }, + "extensions": { "type": "object" } + }, + "required": ["message"] + } + }, + "extensions": { "type": "object" } + } + }); + + // Add definitions if we collected any + if !definitions.is_empty() { + response_schema + .ensure_object() + .insert("definitions".to_string(), definitions.into()); + } + + response_schema +} + +/// Build a schema for a selection set (object fields) +fn build_selection_set_schema( + selection_set: &[Selection], + parent_type: &ExtendedType, + graphql_schema: &GraphQLSchema, + custom_scalar_map: Option<&CustomScalarMap>, + named_fragments: &HashMap>, + definitions: &mut Map, +) -> JSONSchema { + let mut properties = Map::new(); + let mut required = Vec::new(); + + // Always include __typename if it could be useful + let type_name = parent_type.name().to_string(); + + for selection in selection_set { + match selection { + Selection::Field(field) => { + let field_name = field.name.to_string(); + let response_key = field + .alias + .as_ref() + .map(|a| a.to_string()) + .unwrap_or_else(|| field_name.clone()); + + // Skip __typename - it's always a string + if field_name == "__typename" { + properties.insert( + response_key, + json_schema!({"type": "string", "description": "The typename of this object"}).into(), + ); + continue; + } + + // Get field definition from parent type + if let Some(field_def) = get_field_definition(parent_type, &field_name) { + let field_schema = build_field_schema( + field, + &field_def.ty, + graphql_schema, + custom_scalar_map, + named_fragments, + definitions, + field_def.description.as_ref().map(|n| n.to_string()), + ); + + properties.insert(response_key.clone(), field_schema.into()); + + // Non-null fields are required in the response + if field_def.ty.is_non_null() { + required.push(response_key); + } + } else { + warn!( + field = field_name, + parent_type = type_name, + "Field not found in parent type" + ); + } + } + Selection::FragmentSpread(fragment_spread) => { + // Merge fields from named fragment + if let Some(fragment_def) = + named_fragments.get(fragment_spread.fragment_name.as_str()) + && let Some(target_type) = graphql_schema + .types + .get(fragment_def.type_condition.as_str()) + { + let fragment_schema = build_selection_set_schema( + &fragment_def.selection_set, + target_type, + graphql_schema, + custom_scalar_map, + named_fragments, + definitions, + ); + + // Merge properties from fragment + if let Some(props) = fragment_schema + .as_object() + .and_then(|o| o.get("properties")) + .and_then(|v| v.as_object()) + { + for (key, value) in props { + properties.insert(key.clone(), value.clone()); + } + } + } + } + Selection::InlineFragment(inline_fragment) => { + // For inline fragments, we need to handle type conditions + let target_type = if let Some(type_condition) = &inline_fragment.type_condition { + graphql_schema.types.get(type_condition.as_str()) + } else { + Some(parent_type) + }; + + if let Some(target_type) = target_type { + let fragment_schema = build_selection_set_schema( + &inline_fragment.selection_set, + target_type, + graphql_schema, + custom_scalar_map, + named_fragments, + definitions, + ); + + // Merge properties from inline fragment + if let Some(props) = fragment_schema + .as_object() + .and_then(|o| o.get("properties")) + .and_then(|v| v.as_object()) + { + for (key, value) in props { + properties.insert(key.clone(), value.clone()); + } + } + } + } + } + } + + let mut schema = json_schema!({"type": "object"}); + let obj = schema.ensure_object(); + + if !properties.is_empty() { + obj.insert("properties".to_string(), properties.into()); + } + + if !required.is_empty() { + obj.insert( + "required".to_string(), + required + .into_iter() + .map(Value::String) + .collect::>() + .into(), + ); + } + + schema +} + +/// Build schema for a specific field based on its type +fn build_field_schema( + field: &Node, + field_type: &GraphQLType, + graphql_schema: &GraphQLSchema, + custom_scalar_map: Option<&CustomScalarMap>, + named_fragments: &HashMap>, + definitions: &mut Map, + description: Option, +) -> JSONSchema { + let schema = type_to_output_schema( + field_type, + &field.selection_set, + graphql_schema, + custom_scalar_map, + named_fragments, + definitions, + ); + + with_description(schema, description) +} + +/// Convert a GraphQL type to a JSON Schema for output +fn type_to_output_schema( + graphql_type: &GraphQLType, + selection_set: &[Selection], + graphql_schema: &GraphQLSchema, + custom_scalar_map: Option<&CustomScalarMap>, + named_fragments: &HashMap>, + definitions: &mut Map, +) -> JSONSchema { + match graphql_type { + // Non-null types - just unwrap + GraphQLType::NonNullNamed(name) => named_type_to_output_schema( + name, + selection_set, + graphql_schema, + custom_scalar_map, + named_fragments, + definitions, + ), + GraphQLType::NonNullList(inner) => { + let items = type_to_output_schema( + inner.as_ref(), + selection_set, + graphql_schema, + custom_scalar_map, + named_fragments, + definitions, + ); + json_schema!({ + "type": "array", + "items": items + }) + } + + // Nullable types - allow null + GraphQLType::Named(name) => { + let inner = named_type_to_output_schema( + name, + selection_set, + graphql_schema, + custom_scalar_map, + named_fragments, + definitions, + ); + json_schema!({ + "oneOf": [inner, {"type": "null"}] + }) + } + GraphQLType::List(inner) => { + let items = type_to_output_schema( + inner.as_ref(), + selection_set, + graphql_schema, + custom_scalar_map, + named_fragments, + definitions, + ); + json_schema!({ + "oneOf": [ + {"type": "array", "items": items}, + {"type": "null"} + ] + }) + } + } +} + +/// Convert a named GraphQL type to JSON Schema +fn named_type_to_output_schema( + name: &GraphQLName, + selection_set: &[Selection], + graphql_schema: &GraphQLSchema, + custom_scalar_map: Option<&CustomScalarMap>, + named_fragments: &HashMap>, + definitions: &mut Map, +) -> JSONSchema { + match name.as_str() { + // Built-in scalars + "String" => json_schema!({"type": "string"}), + "Int" => json_schema!({"type": "integer"}), + "Float" => json_schema!({"type": "number"}), + "Boolean" => json_schema!({"type": "boolean"}), + // ID can be serialized as string or integer depending on the GraphQL server + "ID" => json_schema!({"oneOf": [{"type": "string"}, {"type": "integer"}]}), + + // Check cache first + other if definitions.contains_key(other) => { + JSONSchema::new_ref(format!("#/definitions/{other}")) + } + + // Look up in schema + other => match graphql_schema.types.get(other) { + // Object types - recurse into selection set + Some(ExtendedType::Object(obj)) => { + if selection_set.is_empty() { + // No selection set - just reference the type + warn!( + type_name = other, + "Object type without selection set in output schema" + ); + json_schema!({}) + } else { + build_selection_set_schema( + selection_set, + &ExtendedType::Object(obj.clone()), + graphql_schema, + custom_scalar_map, + named_fragments, + definitions, + ) + } + } + + // Interface types - similar to objects + Some(ExtendedType::Interface(iface)) => { + if selection_set.is_empty() { + json_schema!({}) + } else { + build_selection_set_schema( + selection_set, + &ExtendedType::Interface(iface.clone()), + graphql_schema, + custom_scalar_map, + named_fragments, + definitions, + ) + } + } + + // Union types - oneOf the possible types based on inline fragments + Some(ExtendedType::Union(_union_def)) => { + if selection_set.is_empty() { + json_schema!({}) + } else { + // Collect schemas for each possible type from inline fragments + let mut type_schemas = Vec::new(); + + for selection in selection_set { + if let Selection::InlineFragment(fragment) = selection + && let Some(type_condition) = &fragment.type_condition + && let Some(member_type) = + graphql_schema.types.get(type_condition.as_str()) + { + let member_schema = build_selection_set_schema( + &fragment.selection_set, + member_type, + graphql_schema, + custom_scalar_map, + named_fragments, + definitions, + ); + type_schemas.push(member_schema); + } + } + + if type_schemas.is_empty() { + // No inline fragments - just return empty schema + json_schema!({}) + } else if type_schemas.len() == 1 { + type_schemas.remove(0) + } else { + json_schema!({"oneOf": type_schemas}) + } + } + } + + // Enum types + // Note: We only include the enum's type description (not per-value descriptions) + // to avoid token bloat with large enums. The `enum` constraint already lists + // all valid values, which is sufficient for understanding output. + Some(ExtendedType::Enum(enum_def)) => { + let values: Vec = enum_def + .values + .iter() + .map(|(_, v)| serde_json::json!(v.value)) + .collect(); + + let mut enum_schema = json_schema!({ + "type": "string", + "enum": values + }); + + // Only include the enum's type description, not per-value descriptions + if let Some(desc) = &enum_def.description { + enum_schema + .ensure_object() + .insert("description".to_string(), desc.to_string().into()); + } + + definitions.insert(other.to_string(), enum_schema.into()); + + JSONSchema::new_ref(format!("#/definitions/{other}")) + } + + // Custom scalars + Some(ExtendedType::Scalar(scalar)) => { + let description = scalar.description.as_ref().map(|n| n.to_string()); + + if let Some(custom_map) = custom_scalar_map + && let Some(custom_schema) = custom_map.get(other) + { + return with_description(custom_schema.clone(), description); + } + + // Unknown scalar - return empty schema with description + with_description(json_schema!({}), description) + } + + // InputObject shouldn't appear in output, but handle gracefully + Some(ExtendedType::InputObject(_)) => { + warn!( + type_name = other, + "InputObject type found in output schema - this is unexpected" + ); + json_schema!({}) + } + + None => { + warn!(type_name = other, "Type not found in schema"); + json_schema!({}) + } + }, + } +} + +/// Get field definition from a parent type (Object or Interface) +fn get_field_definition( + parent_type: &ExtendedType, + field_name: &str, +) -> Option> { + match parent_type { + ExtendedType::Object(obj) => obj.fields.get(field_name).map(|f| f.node.clone()), + ExtendedType::Interface(iface) => iface.fields.get(field_name).map(|f| f.node.clone()), + _ => None, + } +} + +/// Add description to a schema if provided +fn with_description(mut schema: JSONSchema, description: Option) -> JSONSchema { + if let Some(desc) = description { + schema + .ensure_object() + .entry("description") + .or_insert(desc.into()); + } + schema +} + +#[cfg(test)] +mod tests { + use super::*; + use apollo_compiler::parser::Parser; + + fn parse_schema(sdl: &str) -> GraphQLSchema { + GraphQLSchema::parse_and_validate(sdl, "schema.graphql") + .unwrap() + .into_inner() + } + + fn parse_operation(query: &str) -> (apollo_compiler::ast::Document, Vec) { + let doc = Parser::new().parse_ast(query, "query.graphql").unwrap(); + let selection_set = doc + .definitions + .iter() + .find_map(|def| match def { + apollo_compiler::ast::Definition::OperationDefinition(op) => { + Some(op.selection_set.clone()) + } + _ => None, + }) + .unwrap_or_default(); + (doc, selection_set) + } + + #[test] + fn test_simple_query_output_schema() { + let schema = parse_schema( + r#" + type Query { + "Get a user by ID" + user(id: ID!): User + } + + "A user in the system" + type User { + "The user's unique identifier" + id: ID! + "The user's display name" + name: String! + "The user's email address" + email: String + } + "#, + ); + + let (_, selection_set) = parse_operation( + r#" + query GetUser($id: ID!) { + user(id: $id) { + id + name + email + } + } + "#, + ); + + let query_type = schema.types.get("Query").unwrap(); + let output_schema = + selection_set_to_schema(&selection_set, query_type, &schema, None, &HashMap::new()); + + insta::assert_snapshot!(serde_json::to_string_pretty(&output_schema).unwrap()); + } + + #[test] + fn test_nested_object_output_schema() { + let schema = parse_schema( + r#" + type Query { + user(id: ID!): User + } + + type User { + id: ID! + profile: Profile! + } + + type Profile { + bio: String + avatar: String! + } + "#, + ); + + let (_, selection_set) = parse_operation( + r#" + query GetUser($id: ID!) { + user(id: $id) { + id + profile { + bio + avatar + } + } + } + "#, + ); + + let query_type = schema.types.get("Query").unwrap(); + let output_schema = + selection_set_to_schema(&selection_set, query_type, &schema, None, &HashMap::new()); + + insta::assert_snapshot!(serde_json::to_string_pretty(&output_schema).unwrap()); + } +} diff --git a/crates/apollo-mcp-server/src/operations/schema_walker/snapshots/apollo_mcp_server__operations__schema_walker__output__tests__nested_object_output_schema.snap b/crates/apollo-mcp-server/src/operations/schema_walker/snapshots/apollo_mcp_server__operations__schema_walker__output__tests__nested_object_output_schema.snap new file mode 100644 index 00000000..89f890fe --- /dev/null +++ b/crates/apollo-mcp-server/src/operations/schema_walker/snapshots/apollo_mcp_server__operations__schema_walker__output__tests__nested_object_output_schema.snap @@ -0,0 +1,108 @@ +--- +source: crates/apollo-mcp-server/src/operations/schema_walker/output.rs +expression: "serde_json::to_string_pretty(&output_schema).unwrap()" +--- +{ + "type": "object", + "properties": { + "data": { + "type": "object", + "properties": { + "user": { + "oneOf": [ + { + "type": "object", + "properties": { + "id": { + "oneOf": [ + { + "type": "string" + }, + { + "type": "integer" + } + ] + }, + "profile": { + "type": "object", + "properties": { + "bio": { + "oneOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ] + }, + "avatar": { + "type": "string" + } + }, + "required": [ + "avatar" + ] + } + }, + "required": [ + "id", + "profile" + ] + }, + { + "type": "null" + } + ] + } + } + }, + "errors": { + "type": "array", + "items": { + "type": "object", + "properties": { + "message": { + "type": "string" + }, + "locations": { + "type": "array", + "items": { + "type": "object", + "properties": { + "line": { + "type": "integer" + }, + "column": { + "type": "integer" + } + } + } + }, + "path": { + "type": "array", + "items": { + "oneOf": [ + { + "type": "string" + }, + { + "type": "integer" + } + ] + } + }, + "extensions": { + "type": "object" + } + }, + "required": [ + "message" + ] + } + }, + "extensions": { + "type": "object" + } + } +} diff --git a/crates/apollo-mcp-server/src/operations/schema_walker/snapshots/apollo_mcp_server__operations__schema_walker__output__tests__simple_query_output_schema.snap b/crates/apollo-mcp-server/src/operations/schema_walker/snapshots/apollo_mcp_server__operations__schema_walker__output__tests__simple_query_output_schema.snap new file mode 100644 index 00000000..4a41f40f --- /dev/null +++ b/crates/apollo-mcp-server/src/operations/schema_walker/snapshots/apollo_mcp_server__operations__schema_walker__output__tests__simple_query_output_schema.snap @@ -0,0 +1,104 @@ +--- +source: crates/apollo-mcp-server/src/operations/schema_walker/output.rs +expression: "serde_json::to_string_pretty(&output_schema).unwrap()" +--- +{ + "type": "object", + "properties": { + "data": { + "type": "object", + "properties": { + "user": { + "description": "Get a user by ID", + "oneOf": [ + { + "type": "object", + "properties": { + "id": { + "description": "The user's unique identifier", + "oneOf": [ + { + "type": "string" + }, + { + "type": "integer" + } + ] + }, + "name": { + "description": "The user's display name", + "type": "string" + }, + "email": { + "description": "The user's email address", + "oneOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ] + } + }, + "required": [ + "id", + "name" + ] + }, + { + "type": "null" + } + ] + } + } + }, + "errors": { + "type": "array", + "items": { + "type": "object", + "properties": { + "message": { + "type": "string" + }, + "locations": { + "type": "array", + "items": { + "type": "object", + "properties": { + "line": { + "type": "integer" + }, + "column": { + "type": "integer" + } + } + } + }, + "path": { + "type": "array", + "items": { + "oneOf": [ + { + "type": "string" + }, + { + "type": "integer" + } + ] + } + }, + "extensions": { + "type": "object" + } + }, + "required": [ + "message" + ] + } + }, + "extensions": { + "type": "object" + } + } +} diff --git a/crates/apollo-mcp-server/src/runtime.rs b/crates/apollo-mcp-server/src/runtime.rs index 1a6ed528..1baa6eec 100644 --- a/crates/apollo-mcp-server/src/runtime.rs +++ b/crates/apollo-mcp-server/src/runtime.rs @@ -286,6 +286,7 @@ mod test { overrides: Overrides { disable_type_description: false, disable_schema_description: false, + enable_output_schema: false, enable_explorer: false, mutation_mode: None, }, diff --git a/crates/apollo-mcp-server/src/runtime/overrides.rs b/crates/apollo-mcp-server/src/runtime/overrides.rs index 5fcbe66b..25a92e85 100644 --- a/crates/apollo-mcp-server/src/runtime/overrides.rs +++ b/crates/apollo-mcp-server/src/runtime/overrides.rs @@ -12,6 +12,9 @@ pub struct Overrides { /// Disable schema descriptions to save on context-window space pub disable_schema_description: bool, + /// Enable output schema generation for tools (adds token overhead but helps LLMs understand response structure) + pub enable_output_schema: bool, + /// Expose a tool that returns the URL to open a GraphQL operation in Apollo Explorer (requires APOLLO_GRAPH_REF) pub enable_explorer: bool, diff --git a/crates/apollo-mcp-server/src/server.rs b/crates/apollo-mcp-server/src/server.rs index 86341472..1cee33b6 100644 --- a/crates/apollo-mcp-server/src/server.rs +++ b/crates/apollo-mcp-server/src/server.rs @@ -39,6 +39,7 @@ pub struct Server { mutation_mode: MutationMode, disable_type_description: bool, disable_schema_description: bool, + enable_output_schema: bool, disable_auth_token_passthrough: bool, search_leaf_depth: usize, index_memory_bytes: usize, @@ -125,6 +126,7 @@ impl Server { mutation_mode: MutationMode, disable_type_description: bool, disable_schema_description: bool, + enable_output_schema: bool, disable_auth_token_passthrough: bool, search_leaf_depth: usize, index_memory_bytes: usize, @@ -154,6 +156,7 @@ impl Server { mutation_mode, disable_type_description, disable_schema_description, + enable_output_schema, disable_auth_token_passthrough, search_leaf_depth, index_memory_bytes, diff --git a/crates/apollo-mcp-server/src/server/states.rs b/crates/apollo-mcp-server/src/server/states.rs index b787c698..4d7c343f 100644 --- a/crates/apollo-mcp-server/src/server/states.rs +++ b/crates/apollo-mcp-server/src/server/states.rs @@ -48,6 +48,7 @@ struct Config { mutation_mode: MutationMode, disable_type_description: bool, disable_schema_description: bool, + enable_output_schema: bool, disable_auth_token_passthrough: bool, search_leaf_depth: usize, index_memory_bytes: usize, @@ -83,6 +84,7 @@ impl StateMachine { mutation_mode: server.mutation_mode, disable_type_description: server.disable_type_description, disable_schema_description: server.disable_schema_description, + enable_output_schema: server.enable_output_schema, disable_auth_token_passthrough: server.disable_auth_token_passthrough, search_leaf_depth: server.search_leaf_depth, index_memory_bytes: server.index_memory_bytes, diff --git a/crates/apollo-mcp-server/src/server/states/running.rs b/crates/apollo-mcp-server/src/server/states/running.rs index bddc06b6..5ba59aa6 100644 --- a/crates/apollo-mcp-server/src/server/states/running.rs +++ b/crates/apollo-mcp-server/src/server/states/running.rs @@ -62,6 +62,7 @@ pub(super) struct Running { pub(super) mutation_mode: MutationMode, pub(super) disable_type_description: bool, pub(super) disable_schema_description: bool, + pub(super) enable_output_schema: bool, pub(super) disable_auth_token_passthrough: bool, pub(super) health_check: Option, } @@ -90,6 +91,7 @@ impl Running { self.mutation_mode, self.disable_type_description, self.disable_schema_description, + self.enable_output_schema, ) .unwrap_or_else(|error| { error!("Invalid operation: {}", error); @@ -138,6 +140,7 @@ impl Running { self.mutation_mode, self.disable_type_description, self.disable_schema_description, + self.enable_output_schema, ) .unwrap_or_else(|error| { error!("Invalid operation: {}", error); @@ -598,6 +601,7 @@ mod tests { mutation_mode: MutationMode::None, disable_type_description: false, disable_schema_description: false, + enable_output_schema: false, disable_auth_token_passthrough: false, health_check: None, }; @@ -656,6 +660,7 @@ mod tests { mutation_mode: MutationMode::None, disable_type_description: false, disable_schema_description: false, + enable_output_schema: false, disable_auth_token_passthrough: false, health_check: None, }; @@ -695,7 +700,7 @@ mod tests { tools: vec![AppTool { operation: Arc::new( RawOperation::from(("query GetId { id }".to_string(), None)) - .into_operation(&schema, None, MutationMode::All, false, false) + .into_operation(&schema, None, MutationMode::All, false, false, true) .unwrap() .unwrap(), ), @@ -725,6 +730,7 @@ mod tests { mutation_mode: MutationMode::None, disable_type_description: false, disable_schema_description: false, + enable_output_schema: false, disable_auth_token_passthrough: false, health_check: None, } diff --git a/crates/apollo-mcp-server/src/server/states/starting.rs b/crates/apollo-mcp-server/src/server/states/starting.rs index 8dca87cd..0f490760 100644 --- a/crates/apollo-mcp-server/src/server/states/starting.rs +++ b/crates/apollo-mcp-server/src/server/states/starting.rs @@ -51,6 +51,7 @@ impl Starting { self.config.mutation_mode, self.config.disable_type_description, self.config.disable_schema_description, + self.config.enable_output_schema, ) .unwrap_or_else(|error| { error!("Invalid operation: {}", error); @@ -102,6 +103,7 @@ impl Starting { self.config.mutation_mode, self.config.disable_type_description, self.config.disable_schema_description, + self.config.enable_output_schema, ) .map_err(ServerError::Apps)? } else { @@ -168,6 +170,7 @@ impl Starting { mutation_mode: self.config.mutation_mode, disable_type_description: self.config.disable_type_description, disable_schema_description: self.config.disable_schema_description, + enable_output_schema: self.config.enable_output_schema, disable_auth_token_passthrough: self.config.disable_auth_token_passthrough, health_check: health_check.clone(), }; @@ -362,6 +365,7 @@ mod tests { custom_scalar_map: None, disable_type_description: false, disable_schema_description: false, + enable_output_schema: false, disable_auth_token_passthrough: false, search_leaf_depth: 5, index_memory_bytes: 1024 * 1024 * 1024, diff --git a/docs/source/config-file.mdx b/docs/source/config-file.mdx index 2ff34839..3c44ef52 100644 --- a/docs/source/config-file.mdx +++ b/docs/source/config-file.mdx @@ -185,6 +185,7 @@ These fields are under the top-level `overrides` key. | :--------------------------- | :---------------------------------- | :------- | :------------------------------------------------------------------------------------------------------------------------------- | | `disable_type_description` | `bool` | `false` | Disable type descriptions to save on context-window space | | `disable_schema_description` | `bool` | `false` | Disable schema descriptions to save on context-window space | +| `enable_output_schema` | `bool` | `false` | Enable output schema generation for tools. Adds token overhead but helps LLMs understand response structure | | `enable_explorer` | `bool` | `false` | Expose a tool that returns the URL to open a GraphQL operation in Apollo Explorer. Note: This requires a GraphOS graph reference | | `mutation_mode` | `oneOf ["none", "explicit", "all"]` | `"none"` | Defines the mutation access level for the MCP server |