Limit search results to just fields on root path (#195)

Author: Matthew Hawkins

clippy.toml

@@ -1,3 +1,4 @@
 allow-expect-in-tests = true
 allow-panic-in-tests = true
 allow-unwrap-in-tests = true
+allow-indexing-slicing-in-tests = true

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

@@ -30,7 +30,7 @@ pub struct Input {
     type_name: String,
     /// How far to recurse the type hierarchy. Use 0 for no limit. Defaults to 1.
     #[serde(default = "default_depth")]
-    depth: u32,
+    depth: usize,
 }
 
 impl Introspect {
@@ -65,6 +65,7 @@ impl Introspect {
         match schema.types.get(type_name) {
             Some(extended_type) => tree_shaker.retain_type(
                 extended_type,
+                None,
                 if input.depth > 0 {
                     DepthLimit::Limited(input.depth)
                 } else {
@@ -106,6 +107,6 @@ impl Introspect {
 }
 
 /// The default depth to recurse the type hierarchy.
-fn default_depth() -> u32 {
-    1u32
+fn default_depth() -> usize {
+    1
 }

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

@@ -3,9 +3,9 @@
 use crate::errors::McpError;
 use crate::schema_from_type;
 use crate::schema_tree_shake::{DepthLimit, SchemaTreeShaker};
-use apollo_compiler::Schema;
-use apollo_compiler::ast::OperationType as AstOperationType;
+use apollo_compiler::ast::{Field, OperationType as AstOperationType, Selection};
 use apollo_compiler::validation::Valid;
+use apollo_compiler::{Name, Node, Schema};
 use apollo_schema_index::{OperationType, Options, SchemaIndex};
 use rmcp::model::{CallToolResult, Content, ErrorCode, Tool};
 use rmcp::schemars::JsonSchema;
@@ -20,15 +20,16 @@ use tracing::debug;
 /// The name of the tool to search a GraphQL schema.
 pub const SEARCH_TOOL_NAME: &str = "search";
 
-/// The depth of nested types to include for leaf nodes on matching root paths.
-pub const LEAF_DEPTH: DepthLimit = DepthLimit::Limited(1);
+/// The maximum number of search results to consider.
+const MAX_SEARCH_RESULTS: usize = 5;
 
 /// A tool to search a GraphQL schema.
 #[derive(Clone)]
 pub struct Search {
     schema: Arc<Mutex<Valid<Schema>>>,
     index: SchemaIndex,
     allow_mutations: bool,
+    leaf_depth: usize,
     pub tool: Tool,
 }
 
@@ -53,6 +54,8 @@ impl Search {
     pub fn new(
         schema: Arc<Mutex<Valid<Schema>>>,
         allow_mutations: bool,
+        leaf_depth: usize,
+        index_memory_bytes: usize,
     ) -> Result<Self, IndexingError> {
         let root_types = if allow_mutations {
             OperationType::Query | OperationType::Mutation
@@ -62,8 +65,9 @@ impl Search {
         let locked = &schema.try_lock()?;
         Ok(Self {
             schema: schema.clone(),
-            index: SchemaIndex::new(locked, root_types)?,
+            index: SchemaIndex::new(locked, root_types, index_memory_bytes)?,
             allow_mutations,
+            leaf_depth,
             tool: Tool::new(
                 SEARCH_TOOL_NAME,
                 "Search a GraphQL schema",
@@ -84,7 +88,7 @@ impl Search {
                 )
             })?;
 
-        root_paths.truncate(5);
+        root_paths.truncate(MAX_SEARCH_RESULTS);
         debug!(
             "Root paths for search terms: {}\n{}",
             input.terms.join(", "),
@@ -98,16 +102,32 @@ impl Search {
         let schema = self.schema.lock().await;
         let mut tree_shaker = SchemaTreeShaker::new(&schema);
         for root_path in root_paths {
-            let types = root_path.inner.types.clone();
-            let path_len = types.len();
-            for (i, type_name) in types.into_iter().enumerate() {
-                if let Some(extended_type) = schema.types.get(type_name.as_ref()) {
-                    let depth = if i == path_len - 1 {
-                        LEAF_DEPTH
+            let path_len = root_path.inner.len();
+            for (i, path_node) in root_path.inner.into_iter().enumerate() {
+                if let Some(extended_type) = schema.types.get(path_node.node_type.as_str()) {
+                    let (selection_set, depth) = if i == path_len - 1 {
+                        (None, DepthLimit::Limited(self.leaf_depth))
                     } else {
-                        DepthLimit::Limited(1)
+                        (
+                            path_node.field_name.as_ref().map(|field_name| {
+                                vec![Selection::Field(Node::from(Field {
+                                    alias: Default::default(),
+                                    name: Name::new_unchecked(field_name),
+                                    arguments: Default::default(),
+                                    selection_set: Default::default(),
+                                    directives: Default::default(),
+                                }))]
+                            }),
+                            DepthLimit::Limited(1),
+                        )
                     };
-                    tree_shaker.retain_type(extended_type, depth)
+                    tree_shaker.retain_type(extended_type, selection_set.as_ref(), depth)
+                }
+                for field_arg in path_node.field_args {
+                    if let Some(extended_type) = schema.types.get(field_arg.as_str()) {
+                        // Retain input types with unlimited depth because all input must be given
+                        tree_shaker.retain_type(extended_type, None, DepthLimit::Unlimited);
+                    }
                 }
             }
         }
@@ -171,7 +191,8 @@ mod tests {
     #[tokio::test]
     async fn test_search_tool(schema: Valid<Schema>) {
         let schema = Arc::new(Mutex::new(schema));
-        let search = Search::new(schema.clone(), false).expect("Failed to create search tool");
+        let search = Search::new(schema.clone(), false, 1, 15_000_000)
+            .expect("Failed to create search tool");
 
         let result = search
             .execute(Input {
@@ -188,7 +209,8 @@ mod tests {
     #[tokio::test]
     async fn test_referencing_types_are_collected(schema: Valid<Schema>) {
         let schema = Arc::new(Mutex::new(schema));
-        let search = Search::new(schema.clone(), true).expect("Failed to create search tool");
+        let search =
+            Search::new(schema.clone(), true, 1, 15_000_000).expect("Failed to create search tool");
 
         // Search for a type that should have references
         let result = search

crates/apollo-mcp-server/src/introspection/tools/snapshots/apollo_mcp_server__introspection__tools__search__tests__search_tool.snap

@@ -2,50 +2,44 @@
 source: crates/apollo-mcp-server/src/introspection/tools/search.rs
 expression: content_to_snapshot(result)
 ---
+scalar DateTime
+
 enum UserRole {
   ADMIN
   MODERATOR
   USER
   GUEST
 }
 
+enum ContentStatus {
+  DRAFT
+  PUBLISHED
+  ARCHIVED
+  DELETED
+}
+
 type User implements Node {
-  id: ID!
-  createdAt: DateTime!
-  updatedAt: DateTime!
-  username: String!
-  email: String!
   role: UserRole!
-  profile: UserProfile
-  posts: [Post!]!
-  comments: [Comment!]!
-  notifications: [Notification!]!
-  preferences: UserPreferences!
 }
 
 type Post implements Node & Content {
-  id: ID!
-  createdAt: DateTime!
-  updatedAt: DateTime!
-  title: String!
-  content: String!
-  status: ContentStatus!
   author: User!
-  metadata: JSON
-  comments: [Comment!]!
-  media: [Media!]!
-  tags: [Tag!]!
-  analytics: PostAnalytics!
 }
 
 type Query {
-  node(id: ID!): Node
   user(id: ID!): User
   post(id: ID!): Post
   posts(filter: PostFilter): [Post!]!
-  comments(postId: ID!): [Comment!]!
-  notifications(filter: NotificationFilter): [Notification!]!
-  search(query: String!): SearchResult!
 }
 
-union SearchResult = User | Post
+input PostFilter {
+  status: ContentStatus
+  authorId: ID
+  tags: [String!]
+  dateRange: DateRangeInput
+}
+
+input DateRangeInput {
+  start: DateTime!
+  end: DateTime!
+}

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

@@ -149,6 +149,8 @@ async fn main() -> anyhow::Result<()> {
                 .map(|custom_scalars_config| CustomScalarMap::try_from(&custom_scalars_config))
                 .transpose()?,
         )
+        .search_leaf_depth(config.introspection.search.leaf_depth)
+        .index_memory_bytes(config.introspection.search.index_memory_bytes)
         .build()
         .start()
         .await?)

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

@@ -32,11 +32,29 @@ pub struct IntrospectConfig {
 }
 
 /// Search tool configuration
-#[derive(Debug, Default, Deserialize, JsonSchema)]
+#[derive(Debug, Deserialize, JsonSchema)]
 #[serde(default)]
 pub struct SearchConfig {
     /// Enable search tool
     pub enabled: bool,
+
+    /// The amount of memory used for indexing (in bytes)
+    pub index_memory_bytes: usize,
+
+    /// The depth of subtype information to include from matching types
+    /// (1 is just the matching type, 2 is the matching type plus the types it references, etc.
+    /// Defaults to 1.)
+    pub leaf_depth: usize,
+}
+
+impl Default for SearchConfig {
+    fn default() -> Self {
+        Self {
+            enabled: false,
+            index_memory_bytes: 50_000_000,
+            leaf_depth: 1,
+        }
+    }
 }
 
 impl Introspection {

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

@@ -52,7 +52,7 @@ impl RootOperationNames {
 #[derive(Debug, Clone, Copy)]
 pub enum DepthLimit {
     Unlimited,
-    Limited(u32),
+    Limited(usize),
 }
 
 impl DepthLimit {
@@ -169,8 +169,13 @@ impl<'schema> SchemaTreeShaker<'schema> {
     }
 
     /// Retain a specific type, and recursively every type it references, up to a given depth.
-    pub fn retain_type(&mut self, retain: &ExtendedType, depth_limit: DepthLimit) {
-        retain_type(self, retain, None, depth_limit);
+    pub fn retain_type(
+        &mut self,
+        retain: &ExtendedType,
+        selection_set: Option<&Vec<Selection>>,
+        depth_limit: DepthLimit,
+    ) {
+        retain_type(self, retain, selection_set, depth_limit);
     }
 
     pub fn retain_operation(
@@ -1062,7 +1067,7 @@ mod test {
         let query_type = nested_schema.types.get("Query").unwrap();
 
         // Test with depth limit of 1
-        shaker.retain_type(query_type, DepthLimit::Limited(1));
+        shaker.retain_type(query_type, None, DepthLimit::Limited(1));
         let shaken_schema = shaker.shaken().unwrap();
 
         // Should retain only Query, not Level1, Level2, Level3, or Level4
@@ -1074,7 +1079,7 @@ mod test {
 
         // Test with depth limit of 2
         let mut shaker = SchemaTreeShaker::new(&nested_schema);
-        shaker.retain_type(query_type, DepthLimit::Limited(2));
+        shaker.retain_type(query_type, None, DepthLimit::Limited(2));
         let shaken_schema = shaker.shaken().unwrap();
 
         // Should retain Query and Level1, but not deeper levels
@@ -1087,7 +1092,7 @@ mod test {
         // Test with depth limit of 1 starting from Level2
         let mut shaker = SchemaTreeShaker::new(&nested_schema);
         let level2_type = nested_schema.types.get("Level2").unwrap();
-        shaker.retain_type(level2_type, DepthLimit::Limited(1));
+        shaker.retain_type(level2_type, None, DepthLimit::Limited(1));
         let shaken_schema = shaker.shaken().unwrap();
 
         // Should retain only Level2
@@ -1098,7 +1103,7 @@ mod test {
 
         // Test with depth limit of 2 starting from Level2
         let mut shaker = SchemaTreeShaker::new(&nested_schema);
-        shaker.retain_type(level2_type, DepthLimit::Limited(2));
+        shaker.retain_type(level2_type, None, DepthLimit::Limited(2));
         let shaken_schema = shaker.shaken().unwrap();
 
         // Should retain Level2 and Level3
@@ -1109,7 +1114,7 @@ mod test {
 
         // Test with depth limit of 5 starting from Level2
         let mut shaker = SchemaTreeShaker::new(&nested_schema);
-        shaker.retain_type(level2_type, DepthLimit::Limited(5));
+        shaker.retain_type(level2_type, None, DepthLimit::Limited(5));
         let shaken_schema = shaker.shaken().unwrap();
 
         // Should retain Level2 and deeper types
@@ -1127,7 +1132,7 @@ mod test {
         let query_type = nested_schema.types.get("Query").unwrap();
 
         // Test with unlimited depth
-        shaker.retain_type(query_type, DepthLimit::Unlimited);
+        shaker.retain_type(query_type, None, DepthLimit::Unlimited);
         let shaken_schema = shaker.shaken().unwrap();
 
         // Should retain all types

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

@@ -31,6 +31,8 @@ pub struct Server {
     mutation_mode: MutationMode,
     disable_type_description: bool,
     disable_schema_description: bool,
+    search_leaf_depth: usize,
+    index_memory_bytes: usize,
 }
 
 #[derive(Debug, Clone, Deserialize, Default, JsonSchema)]
@@ -93,6 +95,8 @@ impl Server {
         mutation_mode: MutationMode,
         disable_type_description: bool,
         disable_schema_description: bool,
+        search_leaf_depth: usize,
+        index_memory_bytes: usize,
     ) -> Self {
         let headers = {
             let mut headers = headers.clone();
@@ -113,6 +117,8 @@ impl Server {
             mutation_mode,
             disable_type_description,
             disable_schema_description,
+            search_leaf_depth,
+            index_memory_bytes,
         }
     }
 

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

@@ -40,6 +40,8 @@ struct Config {
     mutation_mode: MutationMode,
     disable_type_description: bool,
     disable_schema_description: bool,
+    search_leaf_depth: usize,
+    index_memory_bytes: usize,
 }
 
 impl StateMachine {
@@ -66,6 +68,8 @@ impl StateMachine {
                 mutation_mode: server.mutation_mode,
                 disable_type_description: server.disable_type_description,
                 disable_schema_description: server.disable_schema_description,
+                search_leaf_depth: server.search_leaf_depth,
+                index_memory_bytes: server.index_memory_bytes,
             },
         });
 

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

@@ -94,6 +94,8 @@ impl Starting {
             Some(Search::new(
                 schema.clone(),
                 matches!(self.config.mutation_mode, MutationMode::All),
+                self.config.search_leaf_depth,
+                self.config.index_memory_bytes,
             )?)
         } else {
             None