Merge pull request #176 from apollographql/GT-285

Invalid operation should not crash the MCP Server

Author: DaleSeo

crates/apollo-mcp-server/src/server/states/running.rs

@@ -17,7 +17,7 @@ use tracing::{debug, error};
 
 use crate::{
     custom_scalar_map::CustomScalarMap,
-    errors::{McpError, OperationError, ServerError},
+    errors::{McpError, ServerError},
     explorer::{EXPLORER_TOOL_NAME, Explorer},
     graphql::{self, Executable as _},
     introspection::tools::{
@@ -58,18 +58,21 @@ impl Running {
             .iter()
             .cloned()
             .map(|operation| operation.into_inner())
-            .map(|operation| {
-                operation.into_operation(
+            .filter_map(|operation| {
+                match operation.into_operation(
                     &schema,
                     self.custom_scalar_map.as_ref(),
                     self.mutation_mode,
                     self.disable_type_description,
                     self.disable_schema_description,
-                )
+                ) {
+                    Ok(operation) => operation,
+                    Err(error) => {
+                        error!("Invalid operation: {}", error);
+                        None
+                    }
+                }
             })
-            .collect::<Result<Vec<Option<Operation>>, OperationError>>()?
-            .into_iter()
-            .flatten()
             .collect();
 
         debug!(
@@ -91,23 +94,28 @@ impl Running {
         self,
         operations: Vec<RawOperation>,
     ) -> Result<Running, ServerError> {
+        debug!("Operations updated:\n{:?}", operations);
+
         // Update the operations based on the current schema
         {
             let schema = &*self.schema.lock().await;
             let updated_operations: Vec<Operation> = operations
                 .into_iter()
-                .map(|operation| {
-                    operation.into_operation(
+                .filter_map(|operation| {
+                    match operation.into_operation(
                         schema,
                         self.custom_scalar_map.as_ref(),
                         self.mutation_mode,
                         self.disable_type_description,
                         self.disable_schema_description,
-                    )
+                    ) {
+                        Ok(operation) => operation,
+                        Err(error) => {
+                            error!("Invalid operation: {}", error);
+                            None
+                        }
+                    }
                 })
-                .collect::<Result<Vec<Option<Operation>>, OperationError>>()?
-                .into_iter()
-                .flatten()
                 .collect();
 
             debug!(
@@ -268,3 +276,53 @@ fn convert_arguments<T: serde::de::DeserializeOwned>(
     serde_json::from_value(Value::from(arguments.arguments))
         .map_err(|_| McpError::new(ErrorCode::INVALID_PARAMS, "Invalid input".to_string(), None))
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn invalid_operations_should_not_crash_server() {
+        let schema = Schema::parse("type Query { id: String }", "schema.graphql")
+            .unwrap()
+            .validate()
+            .unwrap();
+
+        let running = Running {
+            schema: Arc::new(Mutex::new(schema)),
+            operations: Arc::new(Mutex::new(vec![])),
+            headers: HeaderMap::new(),
+            endpoint: "http://localhost:4000".to_string(),
+            execute_tool: None,
+            introspect_tool: None,
+            explorer_tool: None,
+            custom_scalar_map: None,
+            peers: Arc::new(RwLock::new(vec![])),
+            cancellation_token: CancellationToken::new(),
+            mutation_mode: MutationMode::None,
+            disable_type_description: false,
+            disable_schema_description: false,
+        };
+
+        let operations = vec![
+            RawOperation::from((
+                "query Valid { id }".to_string(),
+                Some("valid.graphql".to_string()),
+            )),
+            RawOperation::from((
+                "query Invalid {{ id }".to_string(),
+                Some("invalid.graphql".to_string()),
+            )),
+            RawOperation::from((
+                "query { id }".to_string(),
+                Some("unnamed.graphql".to_string()),
+            )),
+        ];
+
+        let updated_running = running.update_operations(operations).await.unwrap();
+        let updated_operations = updated_running.operations.lock().await;
+
+        assert_eq!(updated_operations.len(), 1);
+        assert_eq!(updated_operations.first().unwrap().as_ref().name, "Valid");
+    }
+}

crates/apollo-mcp-server/src/server/states/starting.rs

@@ -13,10 +13,10 @@ use tokio_util::sync::CancellationToken;
 use tracing::{debug, error, info};
 
 use crate::{
-    errors::{OperationError, ServerError},
+    errors::ServerError,
     explorer::Explorer,
     introspection::tools::{execute::Execute, introspect::Introspect},
-    operations::{MutationMode, Operation, RawOperation},
+    operations::{MutationMode, RawOperation},
     server::Transport,
 };
 
@@ -35,18 +35,21 @@ impl Starting {
         let operations: Vec<_> = self
             .operations
             .into_iter()
-            .map(|operation| {
-                operation.into_operation(
+            .filter_map(|operation| {
+                match operation.into_operation(
                     &self.schema,
                     self.config.custom_scalar_map.as_ref(),
                     self.config.mutation_mode,
                     self.config.disable_type_description,
                     self.config.disable_schema_description,
-                )
+                ) {
+                    Ok(operation) => operation,
+                    Err(error) => {
+                        error!("Invalid operation: {}", error);
+                        None
+                    }
+                }
             })
-            .collect::<Result<Vec<Option<Operation>>, OperationError>>()?
-            .into_iter()
-            .flatten()
             .collect();
 
         debug!(