Merge pull request #33 from apollographql/jeffrey/better-subscription-and-mutation-handling

Subscription and mutation handling

Author: Jephuff

crates/mcp-apollo-server/src/errors.rs

@@ -1,9 +1,15 @@
-use apollo_compiler::{Schema, ast::Document, validation::WithErrors};
+use apollo_compiler::{
+    Node, Schema,
+    ast::{Document, OperationDefinition},
+    validation::WithErrors,
+};
 use apollo_federation::error::FederationError;
 use reqwest::header::{InvalidHeaderName, InvalidHeaderValue};
 use rmcp::serde_json;
 use tokio::task::JoinError;
 
+use crate::operations::MutationMode;
+
 /// An error in operation parsing
 #[derive(Debug, thiserror::Error)]
 pub enum OperationError {
@@ -27,6 +33,12 @@ pub enum OperationError {
 
     #[error(transparent)]
     File(#[from] std::io::Error),
+
+    #[error("Mutation {0} not allowed while with --allow-mutations {1:?}")]
+    MutationNotAllowed(Node<OperationDefinition>, MutationMode),
+
+    #[error("Subscriptions are not allowed: {0}")]
+    SubscriptionNotAllowed(Node<OperationDefinition>),
 }
 
 /// An error in server initialization
@@ -38,6 +50,9 @@ pub enum ServerError {
     #[error("Could not parse GraphQL schema: {0}")]
     GraphQLSchema(Box<WithErrors<Schema>>),
 
+    #[error("Could not parse GraphQL schema: {0}")]
+    GraphQLDocumentSchema(Box<WithErrors<Document>>),
+
     #[error("Federation error in GraphQL schema: {0}")]
     Federation(FederationError),
 

crates/mcp-apollo-server/src/introspection.rs

@@ -2,6 +2,7 @@
 
 use crate::errors::McpError;
 use crate::graphql;
+use crate::operations::{MutationMode, operation_defs};
 use crate::schema_from_type;
 use apollo_compiler::Schema;
 use apollo_compiler::validation::Valid;
@@ -41,6 +42,7 @@ impl GetSchema {
 #[derive(Clone)]
 pub struct Execute {
     pub tool: Tool,
+    mutation_mode: MutationMode,
 }
 
 #[derive(JsonSchema, Deserialize)]
@@ -53,8 +55,9 @@ pub struct Input {
 }
 
 impl Execute {
-    pub fn new() -> Self {
+    pub fn new(mutation_mode: MutationMode) -> Self {
         Self {
+            mutation_mode,
             tool: Tool::new(
                 EXECUTE_TOOL_NAME,
                 "Execute a GraphQL operation. Use the `schema` tool to get the GraphQL schema. Always use the schema to create operations - do not try arbitrary operations. DO NOT try to execute introspection queries.",
@@ -70,11 +73,19 @@ impl graphql::Executable for Execute {
     }
 
     fn operation(&self, input: Value) -> Result<String, McpError> {
-        serde_json::from_value::<Input>(input)
-            .map(|input| input.query)
-            .map_err(|_| {
-                McpError::new(ErrorCode::INVALID_PARAMS, "Invalid input".to_string(), None)
-            })
+        let input = serde_json::from_value::<Input>(input).map_err(|_| {
+            McpError::new(ErrorCode::INVALID_PARAMS, "Invalid input".to_string(), None)
+        })?;
+
+        // validate the operation
+        operation_defs(
+            &input.query,
+            self.mutation_mode == MutationMode::All,
+            self.mutation_mode,
+        )
+        .map_err(|e| McpError::new(ErrorCode::INVALID_PARAMS, e.to_string(), None))?;
+
+        Ok(input.query)
     }
 
     fn variables(&self, input: Value) -> Result<Value, McpError> {

crates/mcp-apollo-server/src/lib.rs

@@ -6,5 +6,5 @@ mod introspection;
 pub mod json_schema;
 pub mod operations;
 pub mod sanitize;
+pub(crate) mod schema_tree_shake;
 pub mod server;
-pub(crate) mod tree_shake;

crates/mcp-apollo-server/src/main.rs

@@ -7,8 +7,10 @@ use mcp_apollo_registry::uplink::schema::SchemaSource;
 use mcp_apollo_registry::uplink::{SecretString, UplinkConfig};
 use mcp_apollo_server::custom_scalar_map::CustomScalarMap;
 use mcp_apollo_server::errors::ServerError;
+use mcp_apollo_server::operations::MutationMode;
 use mcp_apollo_server::operations::OperationSource;
-use mcp_apollo_server::server::{Server, Transport};
+use mcp_apollo_server::server::Server;
+use mcp_apollo_server::server::Transport;
 use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
 use std::env;
 use std::path::PathBuf;
@@ -74,6 +76,10 @@ struct Args {
     /// The path to the persisted query manifest containing operations
     #[arg(long)]
     manifest: Option<PathBuf>,
+
+    // Configure when to allow mutations
+    #[clap(long, short = 'm', default_value_t, value_enum)]
+    allow_mutations: MutationMode,
 }
 
 #[tokio::main]
@@ -134,6 +140,7 @@ async fn main() -> anyhow::Result<()> {
         .explorer(args.explorer)
         .headers(default_headers)
         .introspection(args.introspection)
+        .mutation_mode(args.allow_mutations)
         .and_custom_scalar_map(
             args.custom_scalars_config
                 .map(|custom_scalars_config| CustomScalarMap::try_from(&custom_scalars_config))