#
tokens: 30442/50000 2/187 files (page 5/6)
lines: off (toggle) GitHub
raw markdown copy
This is page 5 of 6. Use http://codebase.md/apollographql/apollo-mcp-server?lines=false&page={x} to view the full context.

# Directory Structure

```
├── .cargo
│   └── config.toml
├── .changesets
│   └── README.md
├── .envrc
├── .github
│   ├── CODEOWNERS
│   ├── renovate.json5
│   └── workflows
│       ├── canary-release.yml
│       ├── ci.yml
│       ├── prep-release.yml
│       ├── release-bins.yml
│       ├── release-container.yml
│       ├── sync-develop.yml
│       └── verify-changeset.yml
├── .gitignore
├── .idea
│   └── runConfigurations
│       ├── clippy.xml
│       ├── format___test___clippy.xml
│       ├── format.xml
│       ├── Run_spacedevs.xml
│       └── Test_apollo_mcp_server.xml
├── .vscode
│   ├── extensions.json
│   ├── launch.json
│   ├── settings.json
│   └── tasks.json
├── apollo.config.json
├── Cargo.lock
├── Cargo.toml
├── CHANGELOG_SECTION.md
├── CHANGELOG.md
├── clippy.toml
├── codecov.yml
├── CONTRIBUTING.md
├── crates
│   ├── apollo-mcp-registry
│   │   ├── Cargo.toml
│   │   └── src
│   │       ├── files.rs
│   │       ├── lib.rs
│   │       ├── logging.rs
│   │       ├── platform_api
│   │       │   ├── operation_collections
│   │       │   │   ├── collection_poller.rs
│   │       │   │   ├── error.rs
│   │       │   │   ├── event.rs
│   │       │   │   └── operation_collections.graphql
│   │       │   ├── operation_collections.rs
│   │       │   └── platform-api.graphql
│   │       ├── platform_api.rs
│   │       ├── testdata
│   │       │   ├── minimal_supergraph.graphql
│   │       │   └── supergraph.graphql
│   │       ├── uplink
│   │       │   ├── persisted_queries
│   │       │   │   ├── event.rs
│   │       │   │   ├── manifest_poller.rs
│   │       │   │   ├── manifest.rs
│   │       │   │   └── persisted_queries_manifest_query.graphql
│   │       │   ├── persisted_queries.rs
│   │       │   ├── schema
│   │       │   │   ├── event.rs
│   │       │   │   ├── schema_query.graphql
│   │       │   │   └── schema_stream.rs
│   │       │   ├── schema.rs
│   │       │   ├── snapshots
│   │       │   │   ├── apollo_mcp_registry__uplink__schema__tests__schema_by_url_all_fail@logs.snap
│   │       │   │   ├── apollo_mcp_registry__uplink__schema__tests__schema_by_url_fallback@logs.snap
│   │       │   │   └── apollo_mcp_registry__uplink__schema__tests__schema_by_url@logs.snap
│   │       │   └── uplink.graphql
│   │       └── uplink.rs
│   ├── apollo-mcp-server
│   │   ├── build.rs
│   │   ├── Cargo.toml
│   │   ├── src
│   │   │   ├── auth
│   │   │   │   ├── networked_token_validator.rs
│   │   │   │   ├── protected_resource.rs
│   │   │   │   ├── valid_token.rs
│   │   │   │   └── www_authenticate.rs
│   │   │   ├── auth.rs
│   │   │   ├── config_schema.rs
│   │   │   ├── cors.rs
│   │   │   ├── custom_scalar_map.rs
│   │   │   ├── errors.rs
│   │   │   ├── event.rs
│   │   │   ├── explorer.rs
│   │   │   ├── graphql.rs
│   │   │   ├── headers.rs
│   │   │   ├── health.rs
│   │   │   ├── introspection
│   │   │   │   ├── minify.rs
│   │   │   │   ├── snapshots
│   │   │   │   │   └── apollo_mcp_server__introspection__minify__tests__minify_schema.snap
│   │   │   │   ├── tools
│   │   │   │   │   ├── execute.rs
│   │   │   │   │   ├── introspect.rs
│   │   │   │   │   ├── search.rs
│   │   │   │   │   ├── snapshots
│   │   │   │   │   │   └── apollo_mcp_server__introspection__tools__search__tests__search_tool.snap
│   │   │   │   │   ├── testdata
│   │   │   │   │   │   └── schema.graphql
│   │   │   │   │   └── validate.rs
│   │   │   │   └── tools.rs
│   │   │   ├── introspection.rs
│   │   │   ├── json_schema.rs
│   │   │   ├── lib.rs
│   │   │   ├── main.rs
│   │   │   ├── meter.rs
│   │   │   ├── operations
│   │   │   │   ├── mutation_mode.rs
│   │   │   │   ├── operation_source.rs
│   │   │   │   ├── operation.rs
│   │   │   │   ├── raw_operation.rs
│   │   │   │   ├── schema_walker
│   │   │   │   │   ├── name.rs
│   │   │   │   │   └── type.rs
│   │   │   │   └── schema_walker.rs
│   │   │   ├── operations.rs
│   │   │   ├── runtime
│   │   │   │   ├── config.rs
│   │   │   │   ├── endpoint.rs
│   │   │   │   ├── filtering_exporter.rs
│   │   │   │   ├── graphos.rs
│   │   │   │   ├── introspection.rs
│   │   │   │   ├── logging
│   │   │   │   │   ├── defaults.rs
│   │   │   │   │   ├── log_rotation_kind.rs
│   │   │   │   │   └── parsers.rs
│   │   │   │   ├── logging.rs
│   │   │   │   ├── operation_source.rs
│   │   │   │   ├── overrides.rs
│   │   │   │   ├── schema_source.rs
│   │   │   │   ├── schemas.rs
│   │   │   │   ├── telemetry
│   │   │   │   │   └── sampler.rs
│   │   │   │   └── telemetry.rs
│   │   │   ├── runtime.rs
│   │   │   ├── sanitize.rs
│   │   │   ├── schema_tree_shake.rs
│   │   │   ├── server
│   │   │   │   ├── states
│   │   │   │   │   ├── configuring.rs
│   │   │   │   │   ├── operations_configured.rs
│   │   │   │   │   ├── running.rs
│   │   │   │   │   ├── schema_configured.rs
│   │   │   │   │   └── starting.rs
│   │   │   │   └── states.rs
│   │   │   ├── server.rs
│   │   │   └── telemetry_attributes.rs
│   │   └── telemetry.toml
│   └── apollo-schema-index
│       ├── Cargo.toml
│       └── src
│           ├── error.rs
│           ├── lib.rs
│           ├── path.rs
│           ├── snapshots
│           │   ├── apollo_schema_index__tests__search.snap
│           │   └── apollo_schema_index__traverse__tests__schema_traverse.snap
│           ├── testdata
│           │   └── schema.graphql
│           └── traverse.rs
├── docs
│   └── source
│       ├── _sidebar.yaml
│       ├── auth.mdx
│       ├── best-practices.mdx
│       ├── config-file.mdx
│       ├── cors.mdx
│       ├── custom-scalars.mdx
│       ├── debugging.mdx
│       ├── define-tools.mdx
│       ├── deploy.mdx
│       ├── guides
│       │   └── auth-auth0.mdx
│       ├── health-checks.mdx
│       ├── images
│       │   ├── auth0-permissions-enable.png
│       │   ├── mcp-getstarted-inspector-http.jpg
│       │   └── mcp-getstarted-inspector-stdio.jpg
│       ├── index.mdx
│       ├── licensing.mdx
│       ├── limitations.mdx
│       ├── quickstart.mdx
│       ├── run.mdx
│       └── telemetry.mdx
├── e2e
│   └── mcp-server-tester
│       ├── local-operations
│       │   ├── api.graphql
│       │   ├── config.yaml
│       │   ├── operations
│       │   │   ├── ExploreCelestialBodies.graphql
│       │   │   ├── GetAstronautDetails.graphql
│       │   │   ├── GetAstronautsCurrentlyInSpace.graphql
│       │   │   └── SearchUpcomingLaunches.graphql
│       │   └── tool-tests.yaml
│       ├── pq-manifest
│       │   ├── api.graphql
│       │   ├── apollo.json
│       │   ├── config.yaml
│       │   └── tool-tests.yaml
│       ├── run_tests.sh
│       └── server-config.template.json
├── flake.lock
├── flake.nix
├── graphql
│   ├── TheSpaceDevs
│   │   ├── .vscode
│   │   │   ├── extensions.json
│   │   │   └── tasks.json
│   │   ├── api.graphql
│   │   ├── apollo.config.json
│   │   ├── config.yaml
│   │   ├── operations
│   │   │   ├── ExploreCelestialBodies.graphql
│   │   │   ├── GetAstronautDetails.graphql
│   │   │   ├── GetAstronautsCurrentlyInSpace.graphql
│   │   │   └── SearchUpcomingLaunches.graphql
│   │   ├── persisted_queries
│   │   │   └── apollo.json
│   │   ├── persisted_queries.config.json
│   │   ├── README.md
│   │   └── supergraph.yaml
│   └── weather
│       ├── api.graphql
│       ├── config.yaml
│       ├── operations
│       │   ├── alerts.graphql
│       │   ├── all.graphql
│       │   └── forecast.graphql
│       ├── persisted_queries
│       │   └── apollo.json
│       ├── supergraph.graphql
│       ├── supergraph.yaml
│       └── weather.graphql
├── LICENSE
├── macos-entitlements.plist
├── nix
│   ├── apollo-mcp.nix
│   ├── cargo-zigbuild.patch
│   ├── mcp-server-tools
│   │   ├── default.nix
│   │   ├── node-generated
│   │   │   ├── default.nix
│   │   │   ├── node-env.nix
│   │   │   └── node-packages.nix
│   │   ├── node-mcp-servers.json
│   │   └── README.md
│   └── mcphost.nix
├── README.md
├── rust-toolchain.toml
├── scripts
│   ├── nix
│   │   └── install.sh
│   └── windows
│       └── install.ps1
└── xtask
    ├── Cargo.lock
    ├── Cargo.toml
    └── src
        ├── commands
        │   ├── changeset
        │   │   ├── matching_pull_request.graphql
        │   │   ├── matching_pull_request.rs
        │   │   ├── mod.rs
        │   │   ├── scalars.rs
        │   │   └── snapshots
        │   │       ├── xtask__commands__changeset__tests__it_templatizes_with_multiple_issues_in_title_and_multiple_prs_in_footer.snap
        │   │       ├── xtask__commands__changeset__tests__it_templatizes_with_multiple_issues_in_title.snap
        │   │       ├── xtask__commands__changeset__tests__it_templatizes_with_multiple_prs_in_footer.snap
        │   │       ├── xtask__commands__changeset__tests__it_templatizes_with_neither_issues_or_prs.snap
        │   │       ├── xtask__commands__changeset__tests__it_templatizes_with_prs_in_title_when_empty_issues.snap
        │   │       └── xtask__commands__changeset__tests__it_templatizes_without_prs_in_title_when_issues_present.snap
        │   └── mod.rs
        ├── lib.rs
        └── main.rs
```

# Files

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

```rust
//! Tree shaking for GraphQL schemas

use apollo_compiler::ast::{
    Argument, Definition, DirectiveDefinition, DirectiveList, Document, EnumTypeDefinition, Field,
    FragmentDefinition, InputObjectTypeDefinition, InterfaceTypeDefinition, ObjectTypeDefinition,
    OperationDefinition, OperationType, ScalarTypeDefinition, SchemaDefinition, Selection,
    UnionTypeDefinition,
};
use apollo_compiler::schema::ExtendedType;
use apollo_compiler::schema::InputValueDefinition;
use apollo_compiler::validation::WithErrors;
use apollo_compiler::{Name, Node, Schema};
use std::collections::HashMap;
use tracing::debug;

struct RootOperationNames {
    query: String,
    mutation: String,
    subscription: String,
}

impl RootOperationNames {
    fn operation_name(
        operation_type: OperationType,
        default_name: &str,
        schema: &Schema,
    ) -> String {
        schema
            .root_operation(operation_type)
            .map(|r| r.to_string())
            .unwrap_or(default_name.to_string())
    }

    fn new(schema: &Schema) -> Self {
        Self {
            query: Self::operation_name(OperationType::Query, "query", schema),
            mutation: Self::operation_name(OperationType::Mutation, "mutation", schema),
            subscription: Self::operation_name(OperationType::Subscription, "subscription", schema),
        }
    }

    fn name_for_operation_type(&self, operation_type: OperationType) -> &str {
        match operation_type {
            OperationType::Query => &self.query,
            OperationType::Mutation => &self.mutation,
            OperationType::Subscription => &self.subscription,
        }
    }
}

/// Limits the depth of the schema tree that is retained.
#[derive(Debug, Clone, Copy)]
pub enum DepthLimit {
    Unlimited,
    Limited(usize),
}

impl DepthLimit {
    /// Returns true if the depth limit has been reached.
    pub fn reached(&self) -> bool {
        match self {
            DepthLimit::Unlimited => false,
            DepthLimit::Limited(depth) => *depth == 0,
        }
    }

    /// Decrements the depth limit. This should be called when descending a level in the schema type tree.
    pub fn decrement(self) -> Self {
        match self {
            DepthLimit::Unlimited => self,
            DepthLimit::Limited(depth) => DepthLimit::Limited(depth - 1),
        }
    }
}

/// Tree shaker for GraphQL schemas
pub struct SchemaTreeShaker<'schema> {
    schema: &'schema Schema,
    named_type_nodes: HashMap<String, TreeTypeNode>,
    directive_nodes: HashMap<String, TreeDirectiveNode<'schema>>,
    operation_types: Vec<OperationType>,
    operation_type_names: RootOperationNames,
    named_fragments: HashMap<String, Node<FragmentDefinition>>,
    arguments_descriptions: HashMap<String, Vec<String>>,
}

struct TreeTypeNode {
    retain: bool,
    filtered_field: Option<Vec<String>>,
}

struct TreeDirectiveNode<'schema> {
    node: &'schema DirectiveDefinition,
    retain: bool,
}

impl<'schema> SchemaTreeShaker<'schema> {
    pub(crate) fn argument_descriptions(&self) -> &HashMap<String, Vec<String>> {
        &self.arguments_descriptions
    }

    pub fn new(schema: &'schema Schema) -> Self {
        let mut named_type_nodes: HashMap<String, TreeTypeNode> = HashMap::default();
        let mut directive_nodes: HashMap<String, TreeDirectiveNode> = HashMap::default();

        schema.types.iter().for_each(|(_name, extended_type)| {
            let key = extended_type.name().as_str();
            if named_type_nodes.contains_key(key) {
                tracing::error!("type {} already exists", key);
            }
            named_type_nodes.insert(
                key.to_string(),
                TreeTypeNode {
                    filtered_field: None,
                    retain: false,
                },
            );
        });

        schema
            .directive_definitions
            .iter()
            .for_each(|(name, directive_def)| {
                let key = name.as_str();
                if directive_nodes.contains_key(key) {
                    tracing::error!("directive {} already exists", key);
                }
                directive_nodes.insert(
                    key.to_string(),
                    TreeDirectiveNode {
                        node: directive_def,
                        retain: false,
                    },
                );
            });

        Self {
            schema,
            named_type_nodes,
            directive_nodes,
            operation_types: Vec::default(),
            named_fragments: HashMap::default(),
            operation_type_names: RootOperationNames::new(schema),
            arguments_descriptions: HashMap::default(),
        }
    }

    pub fn retain_operation_type(
        &mut self,
        operation_type: OperationType,
        selection_set: Option<&Vec<Selection>>,
        depth_limit: DepthLimit,
    ) {
        self.operation_types.push(operation_type);
        let operation_type_name = self
            .operation_type_names
            .name_for_operation_type(operation_type);

        if let Some(operation_type_extended_type) = self.schema.types.get(operation_type_name) {
            retain_type(
                self,
                operation_type_extended_type,
                selection_set,
                depth_limit,
            );
        } else {
            tracing::error!("root operation type {} not found in schema", operation_type);
        }
    }

    /// Retain a specific type, and recursively every type it references, up to a given depth.
    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(
        &mut self,
        operation: &OperationDefinition,
        document: &Document,
        depth_limit: DepthLimit,
    ) {
        self.named_fragments = document
            .definitions
            .iter()
            .filter_map(|def| match def {
                Definition::FragmentDefinition(fragment_def) => {
                    Some((fragment_def.name.to_string(), fragment_def.clone()))
                }
                _ => None,
            })
            .collect();
        self.retain_operation_type(
            operation.operation_type,
            Some(&operation.selection_set),
            depth_limit,
        )
    }

    /// Return the set of types retained after tree shaking.
    pub fn shaken(&mut self) -> Result<Schema, Box<WithErrors<Schema>>> {
        let root_operations = self
            .operation_types
            .iter()
            .filter_map(|operation_type| {
                self.schema
                    .root_operation(*operation_type)
                    .cloned()
                    .map(|operation_name| Node::new((*operation_type, operation_name)))
            })
            .collect();

        let schema_definition =
            Definition::SchemaDefinition(apollo_compiler::Node::new(SchemaDefinition {
                root_operations,
                description: self.schema.schema_definition.description.clone(),
                directives: DirectiveList(
                    self.schema
                        .schema_definition
                        .directives
                        .0
                        .iter()
                        .map(|directive| directive.node.clone())
                        .collect(),
                ),
            }));

        let directive_definitions = self
            .schema
            .directive_definitions
            .iter()
            .filter_map(|(directive_name, directive_def)| {
                self.directive_nodes
                    .get(directive_name.as_str())
                    .and_then(|n| {
                        (!directive_def.is_built_in() && n.retain)
                            .then_some(Definition::DirectiveDefinition(directive_def.clone()))
                    })
            })
            .collect();

        let type_definitions: Vec<_> = self
            .schema
            .types
            .iter()
            .filter_map(|(_type_name, extended_type)| {
                if extended_type.is_built_in() {
                    None
                } else {
                    match extended_type {
                        ExtendedType::Object(object_def) => self
                            .named_type_nodes
                            .get(object_def.name.as_str())
                            .and_then(|tree_node| {
                                if tree_node.retain {
                                    Some(Definition::ObjectTypeDefinition(Node::new(
                                        ObjectTypeDefinition {
                                            description: object_def.description.clone(),
                                            directives: DirectiveList(
                                                object_def
                                                    .directives
                                                    .0
                                                    .iter()
                                                    .map(|directive| directive.node.clone())
                                                    .collect(),
                                            ),
                                            name: object_def.name.clone(),
                                            implements_interfaces: object_def
                                                .implements_interfaces
                                                .iter()
                                                .map(|implemented_interface| {
                                                    implemented_interface.name.clone()
                                                })
                                                .collect(),
                                            fields: object_def
                                                .fields
                                                .clone()
                                                .into_iter()
                                                .filter_map(|(field_name, field)| {
                                                    if let Some(filtered_fields) =
                                                        &tree_node.filtered_field
                                                    {
                                                        filtered_fields
                                                            .contains(&field_name.to_string())
                                                            .then_some(field.node)
                                                    } else {
                                                        Some(field.node)
                                                    }
                                                })
                                                .collect(),
                                        },
                                    )))
                                } else if self.schema.root_operation(OperationType::Query).is_some()
                                {
                                    None
                                } else {
                                    tracing::error!("object type {} not found", object_def.name);
                                    None
                                }
                            }),
                        ExtendedType::InputObject(input_def) => self
                            .named_type_nodes
                            .get(input_def.name.as_str())
                            .and_then(|tree_node| {
                                if tree_node.retain {
                                    Some(Definition::InputObjectTypeDefinition(Node::new(
                                        InputObjectTypeDefinition {
                                            description: input_def.description.clone(),
                                            directives: DirectiveList(
                                                input_def
                                                    .directives
                                                    .0
                                                    .iter()
                                                    .map(|directive| directive.node.clone())
                                                    .collect(),
                                            ),
                                            name: input_def.name.clone(),
                                            fields: input_def
                                                .fields
                                                .clone()
                                                .into_iter()
                                                .filter_map(|(field_name, field)| {
                                                    if let Some(filtered_fields) =
                                                        &tree_node.filtered_field
                                                    {
                                                        filtered_fields
                                                            .contains(&field_name.to_string())
                                                            .then_some(field.node)
                                                    } else {
                                                        Some(field.node)
                                                    }
                                                })
                                                .collect(),
                                        },
                                    )))
                                } else {
                                    None
                                }
                            }),
                        ExtendedType::Interface(interface_def) => self
                            .named_type_nodes
                            .get(interface_def.name.as_str())
                            .and_then(|tree_node| {
                                if tree_node.retain {
                                    Some(Definition::InterfaceTypeDefinition(Node::new(
                                        InterfaceTypeDefinition {
                                            description: interface_def.description.clone(),
                                            directives: DirectiveList(
                                                interface_def
                                                    .directives
                                                    .0
                                                    .iter()
                                                    .map(|directive| directive.node.clone())
                                                    .collect(),
                                            ),
                                            name: interface_def.name.clone(),
                                            implements_interfaces: interface_def
                                                .implements_interfaces
                                                .iter()
                                                .map(|implemented_interface| {
                                                    implemented_interface.name.clone()
                                                })
                                                .collect(),
                                            fields: interface_def
                                                .fields
                                                .clone()
                                                .into_iter()
                                                .filter_map(|(field_name, field)| {
                                                    if let Some(filtered_fields) =
                                                        &tree_node.filtered_field
                                                    {
                                                        filtered_fields
                                                            .contains(&field_name.to_string())
                                                            .then_some(field.node)
                                                    } else {
                                                        Some(field.node)
                                                    }
                                                })
                                                .collect(),
                                        },
                                    )))
                                } else {
                                    None
                                }
                            }),
                        ExtendedType::Union(union_def) => self
                            .named_type_nodes
                            .get(union_def.name.as_str())
                            .is_some_and(|n| n.retain)
                            .then(|| {
                                Definition::UnionTypeDefinition(Node::new(UnionTypeDefinition {
                                    description: union_def.description.clone(),
                                    directives: DirectiveList(
                                        union_def
                                            .directives
                                            .0
                                            .iter()
                                            .map(|directive| directive.node.clone())
                                            .collect(),
                                    ),
                                    name: union_def.name.clone(),
                                    members: union_def
                                        .members
                                        .clone()
                                        .into_iter()
                                        .filter_map(|member| {
                                            if let Some(member_tree_node) =
                                                self.named_type_nodes.get(member.as_str())
                                            {
                                                member_tree_node.retain.then_some(member.name)
                                            } else {
                                                tracing::error!(
                                                    "union member {} not found",
                                                    member
                                                );
                                                None
                                            }
                                        })
                                        .collect(),
                                }))
                            }),
                        ExtendedType::Enum(enum_def) => self
                            .named_type_nodes
                            .get(enum_def.name.as_str())
                            .and_then(|tree_node| {
                                if tree_node.retain {
                                    Some(Definition::EnumTypeDefinition(Node::new(
                                        EnumTypeDefinition {
                                            description: enum_def.description.clone(),
                                            directives: DirectiveList(
                                                enum_def
                                                    .directives
                                                    .0
                                                    .iter()
                                                    .map(|directive| directive.node.clone())
                                                    .collect(),
                                            ),
                                            name: enum_def.name.clone(),
                                            values: enum_def
                                                .values
                                                .iter()
                                                .map(|(_enum_value_name, enum_value)| {
                                                    enum_value.node.clone()
                                                })
                                                .collect(),
                                        },
                                    )))
                                } else {
                                    None
                                }
                            }),
                        ExtendedType::Scalar(scalar_def) => self
                            .named_type_nodes
                            .get(scalar_def.name.as_str())
                            .and_then(|tree_node| {
                                if tree_node.retain {
                                    Some(Definition::ScalarTypeDefinition(Node::new(
                                        ScalarTypeDefinition {
                                            description: scalar_def.description.clone(),
                                            directives: DirectiveList(
                                                scalar_def
                                                    .directives
                                                    .0
                                                    .iter()
                                                    .map(|directive| directive.node.clone())
                                                    .collect(),
                                            ),
                                            name: scalar_def.name.clone(),
                                        },
                                    )))
                                } else {
                                    None
                                }
                            }),
                    }
                }
            })
            .collect();

        debug!("Tree shaking resulted in {} types", type_definitions.len());

        let mut document = Document::new();
        document.definitions = [
            // // TODO: don't push if theres no data
            vec![schema_definition],
            directive_definitions,
            type_definitions,
        ]
        .concat();

        document.to_schema().map_err(Box::new)
    }
}

fn selection_set_to_fields(
    selection_set: &Selection,
    named_fragments: &HashMap<String, Node<FragmentDefinition>>,
) -> Vec<Node<Field>> {
    match selection_set {
        Selection::Field(field) => {
            if field.name == "__typename" {
                vec![]
            } else {
                vec![field.clone()]
            }
        }
        Selection::FragmentSpread(fragment) => named_fragments
            .get(fragment.fragment_name.as_str())
            .map(|f| {
                f.selection_set
                    .iter()
                    .flat_map(|s| selection_set_to_fields(s, named_fragments))
                    .collect()
            })
            .unwrap_or_default(),
        Selection::InlineFragment(fragment) => fragment
            .selection_set
            .iter()
            .flat_map(|s| selection_set_to_fields(s, named_fragments))
            .collect(),
    }
}

fn retain_argument_descriptions(
    tree_shaker: &mut SchemaTreeShaker,
    arg: &Node<InputValueDefinition>,
    operation_arguments: &HashMap<&str, &Name>,
) {
    let operation_argument_name = operation_arguments.get(arg.name.as_str());

    if let Some(op_arg_name) = operation_argument_name
        && let Some(description) = arg.description.as_deref()
        && !description.trim().is_empty()
    {
        let descriptions = tree_shaker
            .arguments_descriptions
            .entry(op_arg_name.to_string())
            .or_default();
        descriptions.push(description.trim().to_string())
    }
}

fn build_argument_name_to_value_map(arguments: &[Node<Argument>]) -> HashMap<&str, &Name> {
    arguments
        .iter()
        .filter_map(|a| a.value.as_variable().map(|v| (a.name.as_str(), v)))
        .collect::<HashMap<_, _>>()
}

fn retain_type(
    tree_shaker: &mut SchemaTreeShaker,
    extended_type: &ExtendedType,
    selection_set: Option<&Vec<Selection>>,
    depth_limit: DepthLimit,
) {
    // Check if we've exceeded the depth limit
    if depth_limit.reached() {
        return;
    }

    let type_name = extended_type.name().as_str();
    let selected_fields = if let Some(selection_set) = selection_set {
        let selected_fields = selection_set
            .iter()
            .flat_map(|s| selection_set_to_fields(s, &tree_shaker.named_fragments))
            .collect::<Vec<_>>();

        Some(selected_fields)
    } else {
        None
    };

    if let Some(tree_node) = tree_shaker.named_type_nodes.get_mut(type_name) {
        // If we have already visited this node, early return to avoid infinite recursion.
        // depth_limit and selection_set both have inherent exit cases and may add more types with multiple passes, so never early return for them.
        if tree_node.retain
            && selection_set.is_none()
            && matches!(depth_limit, DepthLimit::Unlimited)
        {
            return;
        }

        tree_node.retain = true;
        if let Some(selected_fields) = selected_fields.as_ref() {
            let additional_fields = selected_fields
                .iter()
                .map(|f| f.name.to_string())
                .collect::<Vec<_>>();

            tree_node.filtered_field = Some(
                [
                    tree_node.filtered_field.clone().unwrap_or_default(),
                    additional_fields,
                ]
                .concat(),
            );
        }
    }

    extended_type
        .directives()
        .iter()
        .for_each(|t| retain_directive(tree_shaker, t.name.as_str(), depth_limit));

    match extended_type {
        ExtendedType::Object(def) => {
            selected_fields
                .as_ref()
                .map(|fields| {
                    fields
                        .iter()
                        .map(|field| {
                            (
                                field.name.as_str(),
                                def.fields.get(field.name.as_str()),
                                Some(&field.directives),
                                Some(&field.selection_set),
                                build_argument_name_to_value_map(&field.arguments),
                            )
                        })
                        .collect::<Vec<_>>()
                })
                .unwrap_or(
                    def.fields
                        .iter()
                        .map(|(name, field_definition)| {
                            (
                                name.as_str(),
                                Some(field_definition),
                                None,
                                None,
                                HashMap::default(),
                            )
                        })
                        .collect::<Vec<_>>(),
                )
                .into_iter()
                .for_each(
                    |(
                        field_name,
                        field_definition,
                        field_selection_directives,
                        field_selection_set,
                        field_arguments,
                    )| {
                        if let Some(field_type) = field_definition {
                            let field_type_name = field_type.ty.inner_named_type();
                            if let Some(field_type_def) =
                                tree_shaker.schema.types.get(field_type_name)
                            {
                                retain_type(
                                    tree_shaker,
                                    field_type_def,
                                    field_selection_set,
                                    depth_limit.decrement(),
                                );
                            } else {
                                tracing::error!("field type {} not found", field_type_name);
                            }

                            field_type.arguments.iter().for_each(|arg| {
                                retain_argument_descriptions(tree_shaker, arg, &field_arguments);

                                let arg_type_name = arg.ty.inner_named_type();
                                if let Some(arg_type) = tree_shaker.schema.types.get(arg_type_name)
                                {
                                    retain_type(
                                        tree_shaker,
                                        arg_type,
                                        None,
                                        depth_limit.decrement(),
                                    );
                                } else {
                                    tracing::error!(
                                        "field argument type {} not found",
                                        arg_type_name
                                    );
                                }
                            });
                        } else {
                            tracing::error!("field {} not found", field_name);
                        }

                        if let Some(field_definition_directives) =
                            field_definition.map(|f| f.directives.clone())
                        {
                            field_definition_directives.iter().for_each(|directive| {
                                retain_directive(tree_shaker, directive.name.as_str(), depth_limit);
                            })
                        }
                        if let Some(field_selection_directives) = field_selection_directives {
                            field_selection_directives.iter().for_each(|directive| {
                                if let Some(directive_definition) = tree_shaker
                                    .schema
                                    .directive_definitions
                                    .get(&directive.name)
                                {
                                    let directive_args_map =
                                        build_argument_name_to_value_map(&directive.arguments);
                                    directive_definition.arguments.iter().for_each(|arg| {
                                        retain_argument_descriptions(
                                            tree_shaker,
                                            arg,
                                            &directive_args_map,
                                        );
                                    });
                                }

                                retain_directive(tree_shaker, directive.name.as_str(), depth_limit);
                            })
                        }
                    },
                );
        }
        ExtendedType::Interface(def) => {
            selected_fields
                .as_ref()
                .map(|fields| {
                    fields
                        .iter()
                        .map(|field| {
                            (
                                field.name.as_str(),
                                def.fields.get(field.name.as_str()),
                                Some(&field.directives),
                                Some(&field.selection_set),
                            )
                        })
                        .collect::<Vec<_>>()
                })
                .unwrap_or(
                    def.fields
                        .iter()
                        .map(|(name, field_definition)| {
                            (name.as_str(), Some(field_definition), None, None)
                        })
                        .collect::<Vec<_>>(),
                )
                .into_iter()
                .for_each(
                    |(
                        field_name,
                        field_definition,
                        field_selection_directives,
                        field_selection_set,
                    )| {
                        if let Some(field_type) = field_definition {
                            let field_type_name = field_type.ty.inner_named_type();
                            if let Some(field_type_def) =
                                tree_shaker.schema.types.get(field_type_name)
                            {
                                retain_type(
                                    tree_shaker,
                                    field_type_def,
                                    field_selection_set,
                                    depth_limit.decrement(),
                                );
                            } else {
                                tracing::error!("field type {} not found", field_type_name);
                            }

                            field_type.arguments.iter().for_each(|arg| {
                                let arg_type_name = arg.ty.inner_named_type();
                                if let Some(arg_type) = tree_shaker.schema.types.get(arg_type_name)
                                {
                                    retain_type(
                                        tree_shaker,
                                        arg_type,
                                        None,
                                        depth_limit.decrement(),
                                    );
                                } else {
                                    tracing::error!(
                                        "field argument type {} not found",
                                        arg_type_name
                                    );
                                }
                            });
                        } else {
                            tracing::error!("field {} not found", field_name);
                        }

                        if let Some(field_definition_directives) =
                            field_definition.map(|f| f.directives.clone())
                        {
                            field_definition_directives.iter().for_each(|directive| {
                                retain_directive(tree_shaker, directive.name.as_str(), depth_limit);
                            })
                        }
                        if let Some(field_selection_directives) = field_selection_directives {
                            field_selection_directives.iter().for_each(|directive| {
                                retain_directive(tree_shaker, directive.name.as_str(), depth_limit);
                            })
                        }
                    },
                );
        }
        ExtendedType::Union(union_def) => union_def.members.iter().for_each(|member| {
            if let Some(member_type) = tree_shaker.schema.types.get(member.as_str()) {
                let member_selection_set = selection_set
                    .map(|selection_set| {
                        selection_set
                            .clone()
                            .into_iter()
                            .filter(|selection| match selection {
                                Selection::Field(_) => true,
                                Selection::FragmentSpread(fragment) => {
                                    if let Some(fragment_def) = &tree_shaker
                                        .named_fragments
                                        .get(fragment.fragment_name.as_str())
                                    {
                                        fragment_def.type_condition == member.as_str()
                                    } else {
                                        tracing::error!(
                                            "fragment {} not found",
                                            fragment.fragment_name
                                        );
                                        false
                                    }
                                }
                                Selection::InlineFragment(fragment) => fragment
                                    .type_condition
                                    .clone()
                                    .is_none_or(|type_condition| {
                                        type_condition.as_str() == member.as_str()
                                    }),
                            })
                            .collect::<Vec<Selection>>()
                    })
                    .and_then(|s| if s.is_empty() { None } else { Some(s) });

                if selection_set.is_none() || member_selection_set.is_some() {
                    retain_type(
                        tree_shaker,
                        member_type,
                        member_selection_set.as_ref(),
                        depth_limit.decrement(),
                    );
                }
            } else {
                tracing::error!("union member {} not found", member);
            }
        }),
        ExtendedType::Enum(def) => def.values.iter().for_each(|(_name, value)| {
            value.directives.iter().for_each(|directive| {
                retain_directive(tree_shaker, directive.name.as_str(), depth_limit);
            })
        }),
        ExtendedType::Scalar(_) => {}
        ExtendedType::InputObject(input_def) => {
            input_def
                .fields
                .iter()
                .for_each(|(_name, field_definition)| {
                    let field_type_name = field_definition.ty.inner_named_type();
                    if let Some(field_type_def) = tree_shaker.schema.types.get(field_type_name) {
                        retain_type(tree_shaker, field_type_def, None, depth_limit.decrement());
                    } else {
                        tracing::error!("field type {} not found", field_type_name);
                    }
                    field_definition.directives.iter().for_each(|directive| {
                        retain_directive(tree_shaker, directive.name.as_str(), depth_limit)
                    });
                });
        }
    }
}

fn retain_directive(
    tree_shaker: &mut SchemaTreeShaker,
    directive_name: &str,
    depth_limit: DepthLimit,
) {
    if let Some(tree_directive_node) = tree_shaker.directive_nodes.get_mut(directive_name) {
        tree_directive_node.retain = true;
        tree_directive_node.node.arguments.iter().for_each(|arg| {
            let arg_type_name = arg.ty.inner_named_type();
            if let Some(arg_type) = tree_shaker.schema.types.get(arg_type_name) {
                retain_type(tree_shaker, arg_type, None, depth_limit.decrement())
            } else {
                tracing::error!("argument type {} not found", arg_type_name);
            }
        });
    }
}

#[cfg(test)]
mod test {
    use apollo_compiler::{ast::OperationType, parser::Parser};
    use rstest::{fixture, rstest};

    use crate::{
        operations::operation_defs,
        schema_tree_shake::{DepthLimit, SchemaTreeShaker},
    };

    #[test]
    fn should_remove_type_mutation_mode_none() {
        let source_text = r#"
            type Query { id: String }
            type Mutation { id: String }
            type Subscription { id: String }
        "#;
        let document = Parser::new()
            .parse_ast(source_text, "schema.graphql")
            .unwrap();
        let schema = document.to_schema_validate().unwrap();
        let mut shaker = SchemaTreeShaker::new(&schema);
        shaker.retain_operation_type(OperationType::Query, None, DepthLimit::Unlimited);
        assert_eq!(
            shaker.shaken().unwrap().to_string(),
            "type Query {\n  id: String\n}\n"
        );
    }

    #[test]
    fn should_remove_type_mutation_mode_all() {
        let source_text = r#"
            type Query { id: String }
            type Mutation { id: String }
            type Subscription { id: String }
        "#;
        let document = Parser::new()
            .parse_ast(source_text, "schema.graphql")
            .unwrap();
        let schema = document.to_schema_validate().unwrap();
        let mut shaker = SchemaTreeShaker::new(&schema);
        shaker.retain_operation_type(OperationType::Query, None, DepthLimit::Unlimited);
        shaker.retain_operation_type(OperationType::Mutation, None, DepthLimit::Unlimited);
        assert_eq!(
            shaker.shaken().unwrap().to_string(),
            "type Query {\n  id: String\n}\n\ntype Mutation {\n  id: String\n}\n"
        );
    }

    #[test]
    fn should_remove_custom_names_mutation_mode_none() {
        let source_text = r#"
            schema {
              query: CustomQuery,
              mutation: CustomMutation,
              subscription: CustomSubscription
            }
            type CustomQuery { id: String }
            type CustomMutation { id: String }
            type CustomSubscription { id: String }
        "#;
        let document = Parser::new()
            .parse_ast(source_text, "schema.graphql")
            .unwrap();
        let schema = document.to_schema_validate().unwrap();
        let mut shaker = SchemaTreeShaker::new(&schema);
        shaker.retain_operation_type(OperationType::Query, None, DepthLimit::Unlimited);
        assert_eq!(
            shaker.shaken().unwrap().to_string(),
            "schema {\n  query: CustomQuery\n}\n\ntype CustomQuery {\n  id: String\n}\n"
        );
    }

    #[test]
    fn should_remove_custom_names_mutation_mode_all() {
        let source_text = r#"
            schema {
              query: CustomQuery,
              mutation: CustomMutation,
              subscription: CustomSubscription
            }
            type CustomQuery { id: String }
            type CustomMutation { id: String }
            type CustomSubscription { id: String }
        "#;
        let document = Parser::new()
            .parse_ast(source_text, "schema.graphql")
            .unwrap();
        let schema = document.to_schema_validate().unwrap();
        let mut shaker = SchemaTreeShaker::new(&schema);
        shaker.retain_operation_type(OperationType::Query, None, DepthLimit::Unlimited);
        shaker.retain_operation_type(OperationType::Mutation, None, DepthLimit::Unlimited);
        assert_eq!(
            shaker.shaken().unwrap().to_string(),
            "schema {\n  query: CustomQuery\n  mutation: CustomMutation\n}\n\ntype CustomQuery {\n  id: String\n}\n\ntype CustomMutation {\n  id: String\n}\n"
        );
    }

    #[test]
    fn should_remove_orphan_types() {
        let source_text = r#"
            type Query { id: UsedInQuery }
            type Mutation { id: UsedInMutation }
            type Subscription { id: UsedInSubscription }
            scalar UsedInQuery
            type UsedInMutation { id: String }
            enum UsedInSubscription { VALUE }
        "#;
        let document = Parser::new()
            .parse_ast(source_text, "schema.graphql")
            .unwrap();
        let schema = document.to_schema_validate().unwrap();
        let mut shaker = SchemaTreeShaker::new(&schema);
        shaker.retain_operation_type(OperationType::Query, None, DepthLimit::Unlimited);
        shaker.retain_operation_type(OperationType::Mutation, None, DepthLimit::Unlimited);
        assert_eq!(
            shaker.shaken().unwrap().to_string(),
            "type Query {\n  id: UsedInQuery\n}\n\ntype Mutation {\n  id: UsedInMutation\n}\n\nscalar UsedInQuery\n\ntype UsedInMutation {\n  id: String\n}\n"
        );
    }

    #[test]
    fn should_work_with_selection_set() {
        let source_text = r#"
            type Query { id: UsedInQuery unused: UsedInQueryButUnusedField }
            type Mutation { id: UsedInMutation }
            type Subscription { id: UsedInSubscription }
            scalar UsedInQuery
            type UsedInQueryButUnusedField { id: String, unused: String }
            type UsedInMutation { id: String }
            enum UsedInSubscription { VALUE }
        "#;
        let document = Parser::new()
            .parse_ast(source_text, "schema.graphql")
            .unwrap();
        let schema = document.to_schema_validate().unwrap();
        let mut shaker = SchemaTreeShaker::new(&schema);
        let (operation_document, operation_def, _comments) = operation_defs(
            "query TestQuery { id }",
            false,
            Some("operation.graphql".to_string()),
        )
        .unwrap()
        .unwrap();
        shaker.retain_operation(&operation_def, &operation_document, DepthLimit::Unlimited);
        assert_eq!(
            shaker.shaken().unwrap().to_string(),
            "type Query {\n  id: UsedInQuery\n}\n\nscalar UsedInQuery\n"
        );
    }

    #[fixture]
    fn nested_schema() -> apollo_compiler::Schema {
        Parser::new()
            .parse_ast(
                r#"
                    type Query  { level1: Level1 }
                    type Level1 { level2: Level2 }
                    type Level2 { level3: Level3 }
                    type Level3 { level4: Level4 }
                    type Level4 { id: String }
                "#,
                "schema.graphql",
            )
            .unwrap()
            .to_schema_validate()
            .unwrap()
            .into_inner()
    }

    #[rstest]
    fn should_respect_depth_limit(nested_schema: apollo_compiler::Schema) {
        let mut shaker = SchemaTreeShaker::new(&nested_schema);

        // Get the Query type to start from
        let query_type = nested_schema.types.get("Query").unwrap();

        // Test with depth limit of 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
        assert!(shaken_schema.types.contains_key("Query"));
        assert!(!shaken_schema.types.contains_key("Level1"));
        assert!(!shaken_schema.types.contains_key("Level2"));
        assert!(!shaken_schema.types.contains_key("Level3"));
        assert!(!shaken_schema.types.contains_key("Level4"));

        // Test with depth limit of 2
        let mut shaker = SchemaTreeShaker::new(&nested_schema);
        shaker.retain_type(query_type, None, DepthLimit::Limited(2));
        let shaken_schema = shaker.shaken().unwrap();

        // Should retain Query and Level1, but not deeper levels
        assert!(shaken_schema.types.contains_key("Query"));
        assert!(shaken_schema.types.contains_key("Level1"));
        assert!(!shaken_schema.types.contains_key("Level2"));
        assert!(!shaken_schema.types.contains_key("Level3"));
        assert!(!shaken_schema.types.contains_key("Level4"));

        // 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, None, DepthLimit::Limited(1));
        let shaken_schema = shaker.shaken().unwrap();

        // Should retain only Level2
        assert!(!shaken_schema.types.contains_key("Level1"));
        assert!(shaken_schema.types.contains_key("Level2"));
        assert!(!shaken_schema.types.contains_key("Level3"));
        assert!(!shaken_schema.types.contains_key("Level4"));

        // Test with depth limit of 2 starting from Level2
        let mut shaker = SchemaTreeShaker::new(&nested_schema);
        shaker.retain_type(level2_type, None, DepthLimit::Limited(2));
        let shaken_schema = shaker.shaken().unwrap();

        // Should retain Level2 and Level3
        assert!(!shaken_schema.types.contains_key("Level1"));
        assert!(shaken_schema.types.contains_key("Level2"));
        assert!(shaken_schema.types.contains_key("Level3"));
        assert!(!shaken_schema.types.contains_key("Level4"));

        // Test with depth limit of 5 starting from Level2
        let mut shaker = SchemaTreeShaker::new(&nested_schema);
        shaker.retain_type(level2_type, None, DepthLimit::Limited(5));
        let shaken_schema = shaker.shaken().unwrap();

        // Should retain Level2 and deeper types
        assert!(!shaken_schema.types.contains_key("Level1"));
        assert!(shaken_schema.types.contains_key("Level2"));
        assert!(shaken_schema.types.contains_key("Level3"));
        assert!(shaken_schema.types.contains_key("Level4"));
    }

    #[rstest]
    fn should_retain_all_types_with_unlimited_depth(nested_schema: apollo_compiler::Schema) {
        let mut shaker = SchemaTreeShaker::new(&nested_schema);

        // Get the Query type to start from
        let query_type = nested_schema.types.get("Query").unwrap();

        // Test with unlimited depth
        shaker.retain_type(query_type, None, DepthLimit::Unlimited);
        let shaken_schema = shaker.shaken().unwrap();

        // Should retain all types
        assert!(shaken_schema.types.contains_key("Query"));
        assert!(shaken_schema.types.contains_key("Level1"));
        assert!(shaken_schema.types.contains_key("Level2"));
        assert!(shaken_schema.types.contains_key("Level3"));
        assert!(shaken_schema.types.contains_key("Level4"));
    }

    #[test]
    fn should_work_with_recursive_schemas() {
        let source_text = r#"
            type Query { id: TypeA }
            type TypeA { id: TypeB }
            type TypeB { id: TypeA }
        "#;
        let document = Parser::new()
            .parse_ast(source_text, "schema.graphql")
            .unwrap();
        let schema = document.to_schema_validate().unwrap();
        let mut shaker = SchemaTreeShaker::new(&schema);
        shaker.retain_operation_type(OperationType::Query, None, DepthLimit::Unlimited);
        assert_eq!(
            shaker.shaken().unwrap().to_string(),
            "type Query {\n  id: TypeA\n}\n\ntype TypeA {\n  id: TypeB\n}\n\ntype TypeB {\n  id: TypeA\n}\n"
        );
    }

    #[test]
    fn should_work_with_recursive_and_depth() {
        let source_text = r#"
            type Query { field1: TypeA, field2: TypeB }
            type TypeA { id: TypeB }
            type TypeB { id: TypeC }
            type TypeC { id: TypeA }
        "#;
        let document = Parser::new()
            .parse_ast(source_text, "schema.graphql")
            .unwrap();
        let schema = document.to_schema_validate().unwrap();
        let mut shaker = SchemaTreeShaker::new(&schema);
        shaker.retain_operation_type(OperationType::Query, None, DepthLimit::Limited(3));
        assert_eq!(
            shaker.shaken().unwrap().to_string(),
            "type Query {\n  field1: TypeA\n  field2: TypeB\n}\n\ntype TypeA {\n  id: TypeB\n}\n\ntype TypeB {\n  id: TypeC\n}\n\ntype TypeC {\n  id: TypeA\n}\n"
        );
    }

    #[test]
    fn should_retain_builtin_directives() {
        let source_text = r#"
            type Query {
                field1: String @deprecated(reason: "Use 'field2' instead")
                field2: String
            }
        "#;
        let document = Parser::new()
            .parse_ast(source_text, "schema.graphql")
            .unwrap();
        let schema = document.to_schema_validate().unwrap();
        let mut shaker = SchemaTreeShaker::new(&schema);
        shaker.retain_operation_type(OperationType::Query, None, DepthLimit::Limited(3));
        assert_eq!(
            shaker.shaken().unwrap().to_string(),
            "type Query {\n  field1: String @deprecated(reason: \"Use 'field2' instead\")\n  field2: String\n}\n"
        );
    }

    #[test]
    fn should_retain_custom_directives() {
        let source_text = r#"
            type Query {
                field1: String @CustomDirective(arg: "Use 'field2' instead")
                field2: String
            }
            directive @CustomDirective(arg: CustomScalar) on FIELD_DEFINITION
            scalar CustomScalar
        "#;
        let document = Parser::new()
            .parse_ast(source_text, "schema.graphql")
            .unwrap();
        let schema = document.to_schema_validate().unwrap();
        let mut shaker = SchemaTreeShaker::new(&schema);
        shaker.retain_operation_type(OperationType::Query, None, DepthLimit::Limited(3));
        assert_eq!(
            shaker.shaken().unwrap().to_string(),
            "directive @CustomDirective(arg: CustomScalar) on FIELD_DEFINITION\n\ntype Query {\n  field1: String @CustomDirective(arg: \"Use 'field2' instead\")\n  field2: String\n}\n\nscalar CustomScalar\n"
        );
    }

    #[test]
    fn recursive_input() {
        let source_text = r#"
            input Filter {
                field: String
                filter: Filter
            }
            type Query {
                field(filter: Filter): String
            }
        "#;
        let document = Parser::new()
            .parse_ast(source_text, "schema.graphql")
            .unwrap();
        let schema = document.to_schema_validate().unwrap();
        let mut shaker = SchemaTreeShaker::new(&schema);
        shaker.retain_operation_type(OperationType::Query, None, DepthLimit::Unlimited);
        assert_eq!(
            shaker.shaken().unwrap().to_string(),
            "input Filter {\n  field: String\n  filter: Filter\n}\n\ntype Query {\n  field(filter: Filter): String\n}\n"
        );
    }

    #[test]
    fn should_retain_field_argument_descriptions() {
        let source_text = r#"
            type Query {
                someQuery(""" an id """ id: ID!, """ other arg """ other: String): OutoutType
            }

            type OutoutType {
                value: String
            }
        "#;
        let document = Parser::new()
            .parse_ast(source_text, "schema.graphql")
            .unwrap();
        let schema = document.to_schema_validate().unwrap();
        let mut shaker = SchemaTreeShaker::new(&schema);
        let (operation_document, operation_def, _comments) = operation_defs(
            "query TestQuery($id1: ID, $other: String) { \
                someQuery(id: $id1, other: $other, otherArg: $other) { \
                    value
                }
            }",
            false,
            Some("operation.graphql".to_string()),
        )
        .unwrap()
        .unwrap();
        shaker.retain_operation(&operation_def, &operation_document, DepthLimit::Unlimited);

        let id_description = shaker.arguments_descriptions.get("id1");
        let other_description = shaker.arguments_descriptions.get("other");

        assert_eq!(shaker.arguments_descriptions.len(), 2);
        assert!(id_description.is_some());
        assert_eq!(*id_description.unwrap(), vec!["an id"]);
        assert!(other_description.is_some());
        assert_eq!(*other_description.unwrap(), vec!["other arg"]);
    }

    #[test]
    fn should_retain_field_argument_descriptions_when_multiple_are_found() {
        let source_text = r#"
            type Query {
                someQuery(""" an id """ id: ID!, """ other arg """ other: String, """ another arg """ otherArg: String): OutoutType
                someQuery2(""" another id """ id: ID!, """ arg 2 """ other: String): OutoutType
            }

            type OutoutType {
                value: String
            }
        "#;
        let document = Parser::new()
            .parse_ast(source_text, "schema.graphql")
            .unwrap();
        let schema = document.to_schema_validate().unwrap();
        let mut shaker = SchemaTreeShaker::new(&schema);
        let (operation_document, operation_def, _comments) = operation_defs(
            "query TestQuery($id: ID, $other2: String) { \
                someQuery(id: $id, other: $other2, otherArg: $other2) { \
                    value
                }
            }",
            false,
            Some("operation.graphql".to_string()),
        )
        .unwrap()
        .unwrap();
        shaker.retain_operation(&operation_def, &operation_document, DepthLimit::Unlimited);

        let id_description = shaker.arguments_descriptions.get("id");
        let other_description = shaker.arguments_descriptions.get("other2");

        assert_eq!(shaker.arguments_descriptions.len(), 2);
        assert!(id_description.is_some());
        assert_eq!(*id_description.unwrap(), vec!["an id"]);
        assert!(other_description.is_some());
        assert_eq!(
            *other_description.unwrap(),
            vec!["other arg", "another arg"]
        );
    }

    #[test]
    fn should_retain_builtin_directive_argument_descriptions() {
        let source_text = r#"
            type Query {
                someQuery(id: ID!, other: Boolean!): OutoutType
            }

            type OutoutType {
                id: ID!
                value: String
            }
        "#;
        let document = Parser::new()
            .parse_ast(source_text, "schema.graphql")
            .unwrap();
        let schema = document.to_schema_validate().unwrap();
        let mut shaker = SchemaTreeShaker::new(&schema);
        let (operation_document, operation_def, _comments) = operation_defs(
            "query TestQuery($id: ID, $other: Boolean!) { \
                someQuery(id: $id, other: $other) { \
                    id
                    value @skip(if: $other)
                }
            }",
            false,
            Some("operation.graphql".to_string()),
        )
        .unwrap()
        .unwrap();
        shaker.retain_operation(&operation_def, &operation_document, DepthLimit::Unlimited);

        let description = shaker.arguments_descriptions.get("other");

        assert!(description.is_some());
        assert_eq!(*description.unwrap(), vec!["Skipped when true."]);
    }

    #[test]
    fn should_retain_custom_directive_argument_descriptions() {
        let source_text = r#"
            type Query {
                someQuery(id: ID!, other: Boolean!): OutoutType
            }

            directive @x(""" the value """ value: String) on FIELD_DEFINITION

            type OutoutType {
                id: ID!
                value: String
            }
        "#;
        let document = Parser::new()
            .parse_ast(source_text, "schema.graphql")
            .unwrap();
        let schema = document.to_schema_validate().unwrap();
        let mut shaker = SchemaTreeShaker::new(&schema);
        let (operation_document, operation_def, _comments) = operation_defs(
            "query TestQuery($id: ID, $other: Boolean!) { \
                someQuery(id: $id, other: $other) { \
                    id
                    value @x(value: $other)
                }
            }",
            false,
            Some("operation.graphql".to_string()),
        )
        .unwrap()
        .unwrap();
        shaker.retain_operation(&operation_def, &operation_document, DepthLimit::Unlimited);

        let description = shaker.arguments_descriptions.get("other");

        assert!(description.is_some());
        assert_eq!(*description.unwrap(), vec!["the value"]);
    }
}

```

--------------------------------------------------------------------------------
/crates/apollo-mcp-server/src/operations/operation.rs:
--------------------------------------------------------------------------------

```rust
use std::collections::HashMap;

use apollo_compiler::{
    Node, Schema as GraphqlSchema,
    ast::{Definition, Document, OperationDefinition, OperationType, Selection, Type},
    parser::Parser,
    schema::ExtendedType,
};
use http::{HeaderMap, HeaderValue};
use regex::Regex;
use rmcp::model::{ErrorCode, Tool, ToolAnnotations};
use schemars::{Schema, json_schema};
use serde::Serialize;
use serde_json::{Map, Value};
use tracing::{debug, info, warn};

use crate::{
    custom_scalar_map::CustomScalarMap,
    errors::{McpError, OperationError},
    graphql::{self, OperationDetails},
    schema_tree_shake::{DepthLimit, SchemaTreeShaker},
};

use super::{MutationMode, RawOperation, schema_walker};

/// A valid GraphQL operation
#[derive(Debug, Clone, Serialize)]
pub struct Operation {
    tool: Tool,
    inner: RawOperation,
    operation_name: String,
}

impl AsRef<Tool> for Operation {
    fn as_ref(&self) -> &Tool {
        &self.tool
    }
}

impl From<Operation> for Tool {
    fn from(value: Operation) -> Tool {
        value.tool
    }
}

impl Operation {
    pub(crate) fn into_inner(self) -> RawOperation {
        self.inner
    }

    #[tracing::instrument(skip_all, name = "load_tool")]
    pub fn from_document(
        raw_operation: RawOperation,
        graphql_schema: &GraphqlSchema,
        custom_scalar_map: Option<&CustomScalarMap>,
        mutation_mode: MutationMode,
        disable_type_description: bool,
        disable_schema_description: bool,
    ) -> Result<Option<Self>, OperationError> {
        if let Some((document, operation, comments)) = operation_defs(
            &raw_operation.source_text,
            mutation_mode != MutationMode::None,
            raw_operation.source_path.clone(),
        )? {
            let operation_name = match operation_name(&operation, raw_operation.source_path.clone())
            {
                Ok(name) => name,
                Err(OperationError::MissingName {
                    source_path,
                    operation,
                }) => {
                    if let Some(path) = source_path {
                        warn!("Skipping unnamed operation in {path}: {operation}");
                    } else {
                        warn!("Skipping unnamed operation: {operation}");
                    }
                    return Ok(None);
                }
                Err(e) => return Err(e),
            };
            let variable_description_overrides =
                variable_description_overrides(&raw_operation.source_text, &operation);
            let mut tree_shaker = SchemaTreeShaker::new(graphql_schema);
            tree_shaker.retain_operation(&operation, &document, DepthLimit::Unlimited);

            let description = Self::tool_description(
                comments,
                &mut tree_shaker,
                graphql_schema,
                &operation,
                disable_type_description,
                disable_schema_description,
            );

            let mut object = serde_json::to_value(get_json_schema(
                &operation,
                tree_shaker.argument_descriptions(),
                &variable_description_overrides,
                graphql_schema,
                custom_scalar_map,
                raw_operation.variables.as_ref(),
            ))?;

            // make sure that the properties field exists since schemas::ObjectValidation is
            // configured to skip empty maps (in the case where there are no input args)
            ensure_properties_exists(&mut object);

            let Value::Object(schema) = object else {
                return Err(OperationError::Internal(
                    "Schemars should have returned an object".to_string(),
                ));
            };

            let tool: Tool = Tool::new(operation_name.clone(), description, schema).annotate(
                ToolAnnotations::new()
                    .read_only(operation.operation_type != OperationType::Mutation),
            );
            let character_count = tool_character_length(&tool);
            match character_count {
                Ok(length) => info!(
                    "Tool {} loaded with a character count of {}. Estimated tokens: {}",
                    operation_name,
                    length,
                    length / 4 // We don't know the tokenization algorithm, so we just use 4 characters per token as a rough estimate. https://docs.anthropic.com/en/docs/resources/glossary#tokens
                ),
                Err(_) => info!(
                    "Tool {} loaded with an unknown character count",
                    operation_name
                ),
            }
            Ok(Some(Operation {
                tool,
                inner: raw_operation,
                operation_name,
            }))
        } else {
            Ok(None)
        }
    }

    /// Generate a description for an operation based on documentation in the schema
    #[tracing::instrument(skip(comments, tree_shaker, graphql_schema, operation_def), fields(operation_type = ?operation_def.operation_type, operation_id = ?operation_def.name))]
    fn tool_description(
        comments: Option<String>,
        tree_shaker: &mut SchemaTreeShaker,
        graphql_schema: &GraphqlSchema,
        operation_def: &Node<OperationDefinition>,
        disable_type_description: bool,
        disable_schema_description: bool,
    ) -> String {
        let comment_description = extract_and_format_comments(comments);

        match comment_description {
            Some(description) => description,
            None => {
                // Add the tree-shaken types to the end of the tool description
                let mut lines = vec![];
                if !disable_type_description {
                    let descriptions = operation_def
                        .selection_set
                        .iter()
                        .filter_map(|selection| {
                            match selection {
                                Selection::Field(field) => {
                                    let field_name = field.name.to_string();
                                    let operation_type = operation_def.operation_type;
                                    if let Some(root_name) =
                                        graphql_schema.root_operation(operation_type)
                                    {
                                        // Find the root field referenced by the operation
                                        let root = graphql_schema.get_object(root_name)?;
                                        let field_definition = root
                                            .fields
                                            .iter()
                                            .find(|(name, _)| {
                                                let name = name.to_string();
                                                name == field_name
                                            })
                                            .map(|(_, field_definition)| {
                                                field_definition.node.clone()
                                            });

                                        // Add the root field description to the tool description
                                        let field_description = field_definition
                                            .clone()
                                            .and_then(|field| field.description.clone())
                                            .map(|node| node.to_string());

                                        // Add information about the return type
                                        let ty = field_definition.map(|field| field.ty.clone());
                                        let type_description =
                                            ty.as_ref().map(Self::type_description);

                                        Some(
                                            vec![field_description, type_description]
                                                .into_iter()
                                                .flatten()
                                                .collect::<Vec<String>>()
                                                .join("\n"),
                                        )
                                    } else {
                                        None
                                    }
                                }
                                _ => None,
                            }
                        })
                        .collect::<Vec<String>>()
                        .join("\n---\n");

                    // Add the tree-shaken types to the end of the tool description

                    lines.push(descriptions);
                }
                if !disable_schema_description {
                    let shaken_schema =
                        tree_shaker.shaken().unwrap_or_else(|schema| schema.partial);

                    let mut types = shaken_schema
                        .types
                        .iter()
                        .filter(|(_name, extended_type)| {
                            !extended_type.is_built_in()
                                && matches!(
                                    extended_type,
                                    ExtendedType::Object(_)
                                        | ExtendedType::Scalar(_)
                                        | ExtendedType::Enum(_)
                                        | ExtendedType::Interface(_)
                                        | ExtendedType::Union(_)
                                )
                                && graphql_schema
                                    .root_operation(operation_def.operation_type)
                                    .is_none_or(|op_name| extended_type.name() != op_name)
                                && graphql_schema
                                    .root_operation(OperationType::Query)
                                    .is_none_or(|op_name| extended_type.name() != op_name)
                        })
                        .peekable();
                    if types.peek().is_some() {
                        lines.push(String::from("---"));
                    }

                    for ty in types {
                        lines.push(ty.1.serialize().to_string());
                    }
                }
                lines.join("\n")
            }
        }
    }

    fn type_description(ty: &Type) -> String {
        let type_name = ty.inner_named_type();
        let mut lines = vec![];
        let optional = if ty.is_non_null() {
            ""
        } else {
            "is optional and "
        };
        let array = if ty.is_list() {
            "is an array of type"
        } else {
            "has type"
        };
        lines.push(format!(
            "The returned value {optional}{array} `{type_name}`"
        ));

        lines.join("\n")
    }
}

impl graphql::Executable for Operation {
    fn persisted_query_id(&self) -> Option<String> {
        // TODO: id was being overridden, should we be returning? Should this be behind a flag? self.inner.persisted_query_id.clone()
        None
    }

    fn operation(&self, _input: Value) -> Result<OperationDetails, McpError> {
        Ok(OperationDetails {
            query: self.inner.source_text.clone(),
            operation_name: Some(self.operation_name.clone()),
        })
    }

    fn variables(&self, input_variables: Value) -> Result<Value, McpError> {
        if let Some(raw_variables) = self.inner.variables.as_ref() {
            let mut variables = match input_variables {
                Value::Null => Ok(serde_json::Map::new()),
                Value::Object(obj) => Ok(obj.clone()),
                _ => Err(McpError::new(
                    ErrorCode::INVALID_PARAMS,
                    "Invalid input".to_string(),
                    None,
                )),
            }?;

            raw_variables.iter().try_for_each(|(key, value)| {
                if variables.contains_key(key) {
                    Err(McpError::new(
                        ErrorCode::INVALID_PARAMS,
                        "No such parameter: {key}",
                        None,
                    ))
                } else {
                    variables.insert(key.clone(), value.clone());
                    Ok(())
                }
            })?;

            Ok(Value::Object(variables))
        } else {
            Ok(input_variables)
        }
    }

    fn headers(&self, default_headers: &HeaderMap<HeaderValue>) -> HeaderMap<HeaderValue> {
        match self.inner.headers.as_ref() {
            None => default_headers.clone(),
            Some(raw_headers) if default_headers.is_empty() => raw_headers.clone(),
            Some(raw_headers) => {
                let mut headers = default_headers.clone();
                raw_headers.iter().for_each(|(key, value)| {
                    if headers.contains_key(key) {
                        tracing::debug!(
                            "Header {} has a default value, overwriting with operation value",
                            key
                        );
                    }
                    headers.insert(key, value.clone());
                });
                headers
            }
        }
    }
}

#[allow(clippy::type_complexity)]
#[tracing::instrument(skip_all)]
pub fn operation_defs(
    source_text: &str,
    allow_mutations: bool,
    source_path: Option<String>,
) -> Result<Option<(Document, Node<OperationDefinition>, Option<String>)>, OperationError> {
    let source_path_clone = source_path.clone();
    let document = Parser::new()
        .parse_ast(
            source_text,
            source_path_clone.unwrap_or_else(|| "operation.graphql".to_string()),
        )
        .map_err(|e| OperationError::GraphQLDocument(Box::new(e)))?;
    let mut last_offset: Option<usize> = Some(0);
    let mut operation_defs = document.definitions.clone().into_iter().filter_map(|def| {
            let description = match def.location() {
                Some(source_span) => {
                    let description = last_offset
                        .map(|start_offset| &source_text[start_offset..source_span.offset()]);
                    last_offset = Some(source_span.end_offset());
                    description
                }
                None => {
                    last_offset = None;
                    None
                }
            };

            match def {
                Definition::OperationDefinition(operation_def) => {
                    Some((operation_def, description))
                }
                Definition::FragmentDefinition(_) => None,
                _ => {
                    eprintln!("Schema definitions were passed in, but only operations and fragments are allowed");
                    None
                }
            }
        });

    let (operation, comments) = match (operation_defs.next(), operation_defs.next()) {
        (None, _) => {
            return Err(OperationError::NoOperations { source_path });
        }
        (_, Some(_)) => {
            return Err(OperationError::TooManyOperations {
                source_path,
                count: 2 + operation_defs.count(),
            });
        }
        (Some(op), None) => op,
    };

    match operation.operation_type {
        OperationType::Subscription => {
            debug!(
                "Skipping subscription operation {}",
                operation_name(&operation, source_path)?
            );
            return Ok(None);
        }
        OperationType::Mutation => {
            if !allow_mutations {
                warn!(
                    "Skipping mutation operation {}",
                    operation_name(&operation, source_path)?
                );
                return Ok(None);
            }
        }
        OperationType::Query => {}
    }

    Ok(Some((document, operation, comments.map(|c| c.to_string()))))
}

pub fn operation_name(
    operation: &Node<OperationDefinition>,
    source_path: Option<String>,
) -> Result<String, OperationError> {
    Ok(operation
        .name
        .as_ref()
        .ok_or_else(|| OperationError::MissingName {
            source_path,
            operation: operation.serialize().no_indent().to_string(),
        })?
        .to_string())
}

#[tracing::instrument(skip_all, fields(operation_type = ?operation_definition.operation_type, operation_id = ?operation_definition.name))]
pub fn variable_description_overrides(
    source_text: &str,
    operation_definition: &Node<OperationDefinition>,
) -> HashMap<String, String> {
    let mut argument_overrides_map: HashMap<String, String> = HashMap::new();
    let mut last_offset = find_opening_parens_offset(source_text, operation_definition);
    operation_definition
        .variables
        .iter()
        .for_each(|v| match v.location() {
            Some(source_span) => {
                let comment = last_offset
                    .map(|start_offset| &source_text[start_offset..source_span.offset()]);

                if let Some(description) = comment.filter(|d| !d.is_empty() && d.contains('#'))
                    && let Some(description) =
                        extract_and_format_comments(Some(description.to_string()))
                {
                    argument_overrides_map.insert(v.name.to_string(), description);
                }

                last_offset = Some(source_span.end_offset());
            }
            None => {
                last_offset = None;
            }
        });

    argument_overrides_map
}

#[tracing::instrument(skip_all, fields(operation_type = ?operation_definition.operation_type, operation_id = ?operation_definition.name))]
pub fn find_opening_parens_offset(
    source_text: &str,
    operation_definition: &Node<OperationDefinition>,
) -> Option<usize> {
    let regex = match Regex::new(r"(?m)^\s*\(") {
        Ok(regex) => regex,
        Err(_) => return None,
    };

    operation_definition
        .name
        .as_ref()
        .and_then(|n| n.location())
        .map(|span| {
            regex
                .find(source_text[span.end_offset()..].as_ref())
                .map(|m| m.start() + m.len() + span.end_offset())
                .unwrap_or(0)
        })
}

pub fn extract_and_format_comments(comments: Option<String>) -> Option<String> {
    comments.and_then(|comments| {
        let content = Regex::new(r"(\n|^)(\s*,*)*#")
            .ok()?
            .replace_all(comments.as_str(), "$1");
        let trimmed = content.trim();

        if trimmed.is_empty() {
            None
        } else {
            Some(trimmed.to_string())
        }
    })
}

fn ensure_properties_exists(json_object: &mut Value) {
    if let Some(obj_type) = json_object.get("type")
        && obj_type == "object"
        && let Some(obj_map) = json_object.as_object_mut()
    {
        let props = obj_map
            .entry("properties")
            .or_insert_with(|| Value::Object(serde_json::Map::new()));
        if !props.is_object() {
            *props = Value::Object(serde_json::Map::new());
        }
    }
}

fn tool_character_length(tool: &Tool) -> Result<usize, serde_json::Error> {
    let tool_schema_string = serde_json::to_string_pretty(&serde_json::json!(tool.input_schema))?;
    Ok(tool.name.len()
        + tool.description.as_ref().map(|d| d.len()).unwrap_or(0)
        + tool_schema_string.len())
}

#[tracing::instrument(skip_all)]
fn get_json_schema(
    operation: &Node<OperationDefinition>,
    schema_argument_descriptions: &HashMap<String, Vec<String>>,
    argument_descriptions_overrides: &HashMap<String, String>,
    graphql_schema: &GraphqlSchema,
    custom_scalar_map: Option<&CustomScalarMap>,
    variable_overrides: Option<&HashMap<String, Value>>,
) -> Schema {
    // Default initialize the schema with the bare minimum needed to be a valid object
    let mut schema = json_schema!({"type": "object", "properties": {}});
    let mut definitions = Map::new();

    // TODO: Can this be unwrapped to use `schema_walker::walk` instead? This functionality is doubled
    // in some cases.
    operation.variables.iter().for_each(|variable| {
        let variable_name = variable.name.to_string();
        if !variable_overrides
            .map(|o| o.contains_key(&variable_name))
            .unwrap_or_default()
        {
            // use overridden description if there is one, otherwise use the schema description
            let description = argument_descriptions_overrides
                .get(&variable_name)
                .cloned()
                .or_else(|| {
                    schema_argument_descriptions
                        .get(&variable_name)
                        .filter(|d| !d.is_empty())
                        .map(|d| d.join("#"))
                });

            let nested = schema_walker::type_to_schema(
                variable.ty.as_ref(),
                graphql_schema,
                &mut definitions,
                custom_scalar_map,
                description,
            );
            schema
                .ensure_object()
                .entry("properties")
                .or_insert(Value::Object(Default::default()))
                .as_object_mut()
                .get_or_insert(&mut Map::default())
                .insert(variable_name.clone(), nested.into());

            if variable.ty.is_non_null() {
                schema
                    .ensure_object()
                    .entry("required")
                    .or_insert(serde_json::Value::Array(Vec::new()))
                    .as_array_mut()
                    .get_or_insert(&mut Vec::default())
                    .push(variable_name.into());
            }
        }
    });

    // Add the definitions to the overall schema if needed
    if !definitions.is_empty() {
        schema
            .ensure_object()
            .insert("definitions".to_string(), definitions.into());
    }

    schema
}

#[cfg(test)]
mod tests {
    use std::{collections::HashMap, str::FromStr as _, sync::LazyLock};

    use apollo_compiler::{Schema, parser::Parser, validation::Valid};
    use rmcp::model::Tool;
    use serde_json::Value;
    use tracing_test::traced_test;

    use crate::{
        custom_scalar_map::CustomScalarMap,
        graphql::Executable as _,
        operations::{MutationMode, Operation, RawOperation},
    };

    // Example schema for tests
    static SCHEMA: LazyLock<Valid<Schema>> = LazyLock::new(|| {
        Schema::parse(
            r#"
                type Query {
                    id: String
                    enum: RealEnum
                    customQuery(""" id description """ id: ID!, """ a flag """ flag: Boolean): OutputType
                    testOp: OpResponse
                }
                type Mutation {id: String }

                """
                RealCustomScalar exists
                """
                scalar RealCustomScalar
                input RealInputObject {
                    """
                    optional is a input field that is optional
                    """
                    optional: String

                    """
                    required is a input field that is required
                    """
                    required: String!
                }

                type OpResponse {
                  id: String
                }

                """
                the description for the enum
                """
                enum RealEnum {
                    """
                    ENUM_VALUE_1 is a value
                    """
                    ENUM_VALUE_1

                    """
                    ENUM_VALUE_2 is a value
                    """
                    ENUM_VALUE_2
                }

                """
                custom output type
                """
                type OutputType {
                    id: ID!
                }
            "#,
            "operation.graphql",
        )
        .expect("schema should parse")
        .validate()
        .expect("schema should be valid")
    });

    /// Serializes the input to JSON, sorting the object keys
    macro_rules! to_sorted_json {
        ($json:expr) => {{
            let mut j = serde_json::json!($json);
            j.sort_all_objects();

            j
        }};
    }

    #[test]
    fn nullable_named_type() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: "query QueryName($id: ID) { id }".to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &SCHEMA,
            None,
            MutationMode::None,
            false,
            false,
        )
        .unwrap()
        .unwrap();
        let tool = Tool::from(operation);

        insta::assert_debug_snapshot!(tool, @r###"
        Tool {
            name: "QueryName",
            title: None,
            description: Some(
                "The returned value is optional and has type `String`",
            ),
            input_schema: {
                "type": String("object"),
                "properties": Object {
                    "id": Object {
                        "type": String("string"),
                    },
                },
            },
            output_schema: None,
            annotations: Some(
                ToolAnnotations {
                    title: None,
                    read_only_hint: Some(
                        true,
                    ),
                    destructive_hint: None,
                    idempotent_hint: None,
                    open_world_hint: None,
                },
            ),
            icons: None,
        }
        "###);

        let json = to_sorted_json!(tool.input_schema);
        insta::assert_snapshot!(serde_json::to_string_pretty(&json).unwrap(), @r#"
        {
          "properties": {
            "id": {
              "type": "string"
            }
          },
          "type": "object"
        }
        "#);
    }

    #[test]
    fn non_nullable_named_type() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: "query QueryName($id: ID!) { id }".to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &SCHEMA,
            None,
            MutationMode::None,
            false,
            false,
        )
        .unwrap()
        .unwrap();
        let tool = Tool::from(operation);

        insta::assert_debug_snapshot!(tool, @r###"
        Tool {
            name: "QueryName",
            title: None,
            description: Some(
                "The returned value is optional and has type `String`",
            ),
            input_schema: {
                "type": String("object"),
                "properties": Object {
                    "id": Object {
                        "type": String("string"),
                    },
                },
                "required": Array [
                    String("id"),
                ],
            },
            output_schema: None,
            annotations: Some(
                ToolAnnotations {
                    title: None,
                    read_only_hint: Some(
                        true,
                    ),
                    destructive_hint: None,
                    idempotent_hint: None,
                    open_world_hint: None,
                },
            ),
            icons: None,
        }
        "###);
        insta::assert_snapshot!(serde_json::to_string_pretty(&serde_json::json!(tool.input_schema)).unwrap(), @r###"
        {
          "type": "object",
          "properties": {
            "id": {
              "type": "string"
            }
          },
          "required": [
            "id"
          ]
        }
        "###);
    }

    #[test]
    fn non_nullable_list_of_nullable_named_type() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: "query QueryName($id: [ID]!) { id }".to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &SCHEMA,
            None,
            MutationMode::None,
            false,
            false,
        )
        .unwrap()
        .unwrap();
        let tool = Tool::from(operation);

        insta::assert_debug_snapshot!(tool, @r###"
        Tool {
            name: "QueryName",
            title: None,
            description: Some(
                "The returned value is optional and has type `String`",
            ),
            input_schema: {
                "type": String("object"),
                "properties": Object {
                    "id": Object {
                        "type": String("array"),
                        "items": Object {
                            "oneOf": Array [
                                Object {
                                    "type": String("string"),
                                },
                                Object {
                                    "type": String("null"),
                                },
                            ],
                        },
                    },
                },
                "required": Array [
                    String("id"),
                ],
            },
            output_schema: None,
            annotations: Some(
                ToolAnnotations {
                    title: None,
                    read_only_hint: Some(
                        true,
                    ),
                    destructive_hint: None,
                    idempotent_hint: None,
                    open_world_hint: None,
                },
            ),
            icons: None,
        }
        "###);
        insta::assert_snapshot!(serde_json::to_string_pretty(&serde_json::json!(tool.input_schema)).unwrap(), @r###"
        {
          "type": "object",
          "properties": {
            "id": {
              "type": "array",
              "items": {
                "oneOf": [
                  {
                    "type": "string"
                  },
                  {
                    "type": "null"
                  }
                ]
              }
            }
          },
          "required": [
            "id"
          ]
        }
        "###);
    }

    #[test]
    fn non_nullable_list_of_non_nullable_named_type() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: "query QueryName($id: [ID!]!) { id }".to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &SCHEMA,
            None,
            MutationMode::None,
            false,
            false,
        )
        .unwrap()
        .unwrap();
        let tool = Tool::from(operation);

        insta::assert_debug_snapshot!(tool, @r###"
        Tool {
            name: "QueryName",
            title: None,
            description: Some(
                "The returned value is optional and has type `String`",
            ),
            input_schema: {
                "type": String("object"),
                "properties": Object {
                    "id": Object {
                        "type": String("array"),
                        "items": Object {
                            "type": String("string"),
                        },
                    },
                },
                "required": Array [
                    String("id"),
                ],
            },
            output_schema: None,
            annotations: Some(
                ToolAnnotations {
                    title: None,
                    read_only_hint: Some(
                        true,
                    ),
                    destructive_hint: None,
                    idempotent_hint: None,
                    open_world_hint: None,
                },
            ),
            icons: None,
        }
        "###);
        insta::assert_snapshot!(serde_json::to_string_pretty(&serde_json::json!(tool.input_schema)).unwrap(), @r###"
        {
          "type": "object",
          "properties": {
            "id": {
              "type": "array",
              "items": {
                "type": "string"
              }
            }
          },
          "required": [
            "id"
          ]
        }
        "###);
    }

    #[test]
    fn nullable_list_of_nullable_named_type() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: "query QueryName($id: [ID]) { id }".to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &SCHEMA,
            None,
            MutationMode::None,
            false,
            false,
        )
        .unwrap()
        .unwrap();
        let tool = Tool::from(operation);

        insta::assert_debug_snapshot!(tool, @r###"
        Tool {
            name: "QueryName",
            title: None,
            description: Some(
                "The returned value is optional and has type `String`",
            ),
            input_schema: {
                "type": String("object"),
                "properties": Object {
                    "id": Object {
                        "type": String("array"),
                        "items": Object {
                            "oneOf": Array [
                                Object {
                                    "type": String("string"),
                                },
                                Object {
                                    "type": String("null"),
                                },
                            ],
                        },
                    },
                },
            },
            output_schema: None,
            annotations: Some(
                ToolAnnotations {
                    title: None,
                    read_only_hint: Some(
                        true,
                    ),
                    destructive_hint: None,
                    idempotent_hint: None,
                    open_world_hint: None,
                },
            ),
            icons: None,
        }
        "###);
        insta::assert_snapshot!(serde_json::to_string_pretty(&serde_json::json!(tool.input_schema)).unwrap(), @r#"
        {
          "type": "object",
          "properties": {
            "id": {
              "type": "array",
              "items": {
                "oneOf": [
                  {
                    "type": "string"
                  },
                  {
                    "type": "null"
                  }
                ]
              }
            }
          }
        }
        "#);
    }

    #[test]
    fn nullable_list_of_non_nullable_named_type() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: "query QueryName($id: [ID!]) { id }".to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &SCHEMA,
            None,
            MutationMode::None,
            false,
            false,
        )
        .unwrap()
        .unwrap();
        let tool = Tool::from(operation);

        insta::assert_debug_snapshot!(tool, @r###"
        Tool {
            name: "QueryName",
            title: None,
            description: Some(
                "The returned value is optional and has type `String`",
            ),
            input_schema: {
                "type": String("object"),
                "properties": Object {
                    "id": Object {
                        "type": String("array"),
                        "items": Object {
                            "type": String("string"),
                        },
                    },
                },
            },
            output_schema: None,
            annotations: Some(
                ToolAnnotations {
                    title: None,
                    read_only_hint: Some(
                        true,
                    ),
                    destructive_hint: None,
                    idempotent_hint: None,
                    open_world_hint: None,
                },
            ),
            icons: None,
        }
        "###);
        insta::assert_snapshot!(serde_json::to_string_pretty(&serde_json::json!(tool.input_schema)).unwrap(), @r#"
        {
          "type": "object",
          "properties": {
            "id": {
              "type": "array",
              "items": {
                "type": "string"
              }
            }
          }
        }
        "#);
    }

    #[test]
    fn nullable_list_of_nullable_lists_of_nullable_named_types() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: "query QueryName($id: [[ID]]) { id }".to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &SCHEMA,
            None,
            MutationMode::None,
            false,
            false,
        )
        .unwrap()
        .unwrap();
        let tool = Tool::from(operation);

        insta::assert_debug_snapshot!(tool, @r###"
        Tool {
            name: "QueryName",
            title: None,
            description: Some(
                "The returned value is optional and has type `String`",
            ),
            input_schema: {
                "type": String("object"),
                "properties": Object {
                    "id": Object {
                        "type": String("array"),
                        "items": Object {
                            "oneOf": Array [
                                Object {
                                    "type": String("array"),
                                    "items": Object {
                                        "oneOf": Array [
                                            Object {
                                                "type": String("string"),
                                            },
                                            Object {
                                                "type": String("null"),
                                            },
                                        ],
                                    },
                                },
                                Object {
                                    "type": String("null"),
                                },
                            ],
                        },
                    },
                },
            },
            output_schema: None,
            annotations: Some(
                ToolAnnotations {
                    title: None,
                    read_only_hint: Some(
                        true,
                    ),
                    destructive_hint: None,
                    idempotent_hint: None,
                    open_world_hint: None,
                },
            ),
            icons: None,
        }
        "###);
        insta::assert_snapshot!(serde_json::to_string_pretty(&serde_json::json!(tool.input_schema)).unwrap(), @r#"
        {
          "type": "object",
          "properties": {
            "id": {
              "type": "array",
              "items": {
                "oneOf": [
                  {
                    "type": "array",
                    "items": {
                      "oneOf": [
                        {
                          "type": "string"
                        },
                        {
                          "type": "null"
                        }
                      ]
                    }
                  },
                  {
                    "type": "null"
                  }
                ]
              }
            }
          }
        }
        "#);
    }

    #[test]
    fn nullable_input_object() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: "query QueryName($id: RealInputObject) { id }".to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &SCHEMA,
            None,
            MutationMode::None,
            false,
            false,
        )
        .unwrap()
        .unwrap();
        let tool = Tool::from(operation);

        insta::assert_debug_snapshot!(tool, @r###"
        Tool {
            name: "QueryName",
            title: None,
            description: Some(
                "The returned value is optional and has type `String`",
            ),
            input_schema: {
                "type": String("object"),
                "properties": Object {
                    "id": Object {
                        "$ref": String("#/definitions/RealInputObject"),
                    },
                },
                "definitions": Object {
                    "RealInputObject": Object {
                        "type": String("object"),
                        "properties": Object {
                            "optional": Object {
                                "description": String("optional is a input field that is optional"),
                                "type": String("string"),
                            },
                            "required": Object {
                                "description": String("required is a input field that is required"),
                                "type": String("string"),
                            },
                        },
                        "required": Array [
                            String("required"),
                        ],
                    },
                },
            },
            output_schema: None,
            annotations: Some(
                ToolAnnotations {
                    title: None,
                    read_only_hint: Some(
                        true,
                    ),
                    destructive_hint: None,
                    idempotent_hint: None,
                    open_world_hint: None,
                },
            ),
            icons: None,
        }
        "###);
    }

    #[test]
    fn non_nullable_enum() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: "query QueryName($id: RealEnum!) { id }".to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &SCHEMA,
            None,
            MutationMode::None,
            false,
            false,
        )
        .unwrap()
        .unwrap();
        let tool = Tool::from(operation);

        insta::assert_debug_snapshot!(tool, @r###"
        Tool {
            name: "QueryName",
            title: None,
            description: Some(
                "The returned value is optional and has type `String`",
            ),
            input_schema: {
                "type": String("object"),
                "properties": Object {
                    "id": Object {
                        "$ref": String("#/definitions/RealEnum"),
                    },
                },
                "required": Array [
                    String("id"),
                ],
                "definitions": Object {
                    "RealEnum": Object {
                        "description": String("the description for the enum\n\nValues:\nENUM_VALUE_1: ENUM_VALUE_1 is a value\nENUM_VALUE_2: ENUM_VALUE_2 is a value"),
                        "type": String("string"),
                        "enum": Array [
                            String("ENUM_VALUE_1"),
                            String("ENUM_VALUE_2"),
                        ],
                    },
                },
            },
            output_schema: None,
            annotations: Some(
                ToolAnnotations {
                    title: None,
                    read_only_hint: Some(
                        true,
                    ),
                    destructive_hint: None,
                    idempotent_hint: None,
                    open_world_hint: None,
                },
            ),
            icons: None,
        }
        "###);
    }

    #[test]
    fn multiple_operations_should_error() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: "query QueryName { id } query QueryName { id }".to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: Some("operation.graphql".to_string()),
            },
            &SCHEMA,
            None,
            MutationMode::None,
            false,
            false,
        );
        insta::assert_debug_snapshot!(operation, @r#"
        Err(
            TooManyOperations {
                source_path: Some(
                    "operation.graphql",
                ),
                count: 2,
            },
        )
        "#);
    }

    #[test]
    #[traced_test]
    fn unnamed_operations_should_be_skipped() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: "query { id }".to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: Some("operation.graphql".to_string()),
            },
            &SCHEMA,
            None,
            MutationMode::None,
            false,
            false,
        );
        assert!(operation.unwrap().is_none());

        logs_assert(|lines: &[&str]| {
            lines
                .iter()
                .filter(|line| line.contains("WARN"))
                .any(|line| {
                    line.contains("Skipping unnamed operation in operation.graphql: { id }")
                })
                .then_some(())
                .ok_or("Expected warning about unnamed operation in logs".to_string())
        });
    }

    #[test]
    fn no_operations_should_error() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: "fragment Test on Query { id }".to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: Some("operation.graphql".to_string()),
            },
            &SCHEMA,
            None,
            MutationMode::None,
            false,
            false,
        );
        insta::assert_debug_snapshot!(operation, @r#"
        Err(
            NoOperations {
                source_path: Some(
                    "operation.graphql",
                ),
            },
        )
        "#);
    }

    #[test]
    fn schema_should_error() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: "type Query { id: String }".to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &SCHEMA,
            None,
            MutationMode::None,
            false,
            false,
        );
        insta::assert_debug_snapshot!(operation, @r"
        Err(
            NoOperations {
                source_path: None,
            },
        )
        ");
    }

    #[test]
    #[traced_test]
    fn unknown_type_should_be_any() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: "query QueryName($id: FakeType) { id }".to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &SCHEMA,
            None,
            MutationMode::None,
            false,
            false,
        )
        .unwrap()
        .unwrap();
        let tool = Tool::from(operation);

        // Verify that a warning was logged
        logs_assert(|lines: &[&str]| {
            lines
                .iter()
                .filter(|line| line.contains("WARN"))
                .any(|line| line.contains("Type not found in schema name=\"FakeType\""))
                .then_some(())
                .ok_or("Expected warning about unknown type in logs".to_string())
        });

        insta::assert_debug_snapshot!(tool, @r###"
        Tool {
            name: "QueryName",
            title: None,
            description: Some(
                "The returned value is optional and has type `String`",
            ),
            input_schema: {
                "type": String("object"),
                "properties": Object {
                    "id": Object {},
                },
            },
            output_schema: None,
            annotations: Some(
                ToolAnnotations {
                    title: None,
                    read_only_hint: Some(
                        true,
                    ),
                    destructive_hint: None,
                    idempotent_hint: None,
                    open_world_hint: None,
                },
            ),
            icons: None,
        }
        "###);
    }

    #[test]
    #[traced_test]
    fn custom_scalar_without_map_should_be_any() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: "query QueryName($id: RealCustomScalar) { id }".to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &SCHEMA,
            None,
            MutationMode::None,
            false,
            false,
        )
        .unwrap()
        .unwrap();
        let tool = Tool::from(operation);

        // Verify that a warning was logged
        logs_assert(|lines: &[&str]| {
            lines
                .iter()
                .filter(|line| line.contains("WARN"))
                .any(|line| line.contains("custom scalars aren't currently supported without a custom_scalar_map name=\"RealCustomScalar\""))
                .then_some(())
                .ok_or("Expected warning about custom scalar without map in logs".to_string())
        });

        insta::assert_debug_snapshot!(tool, @r###"
        Tool {
            name: "QueryName",
            title: None,
            description: Some(
                "The returned value is optional and has type `String`",
            ),
            input_schema: {
                "type": String("object"),
                "properties": Object {
                    "id": Object {
                        "$ref": String("#/definitions/RealCustomScalar"),
                    },
                },
                "definitions": Object {
                    "RealCustomScalar": Object {
                        "description": String("RealCustomScalar exists"),
                    },
                },
            },
            output_schema: None,
            annotations: Some(
                ToolAnnotations {
                    title: None,
                    read_only_hint: Some(
                        true,
                    ),
                    destructive_hint: None,
                    idempotent_hint: None,
                    open_world_hint: None,
                },
            ),
            icons: None,
        }
        "###);
    }

    #[test]
    #[traced_test]
    fn custom_scalar_with_map_but_not_found_should_error() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: "query QueryName($id: RealCustomScalar) { id }".to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &SCHEMA,
            Some(&CustomScalarMap::from_str("{}").unwrap()),
            MutationMode::None,
            false,
            false,
        )
        .unwrap()
        .unwrap();
        let tool = Tool::from(operation);

        // Verify that a warning was logged
        logs_assert(|lines: &[&str]| {
            lines
                .iter()
                .filter(|line| line.contains("WARN"))
                .any(|line| {
                    line.contains(
                        "custom scalar missing from custom_scalar_map name=\"RealCustomScalar\"",
                    )
                })
                .then_some(())
                .ok_or("Expected warning about custom scalar missing in logs".to_string())
        });

        insta::assert_debug_snapshot!(tool, @r###"
        Tool {
            name: "QueryName",
            title: None,
            description: Some(
                "The returned value is optional and has type `String`",
            ),
            input_schema: {
                "type": String("object"),
                "properties": Object {
                    "id": Object {
                        "$ref": String("#/definitions/RealCustomScalar"),
                    },
                },
                "definitions": Object {
                    "RealCustomScalar": Object {
                        "description": String("RealCustomScalar exists"),
                    },
                },
            },
            output_schema: None,
            annotations: Some(
                ToolAnnotations {
                    title: None,
                    read_only_hint: Some(
                        true,
                    ),
                    destructive_hint: None,
                    idempotent_hint: None,
                    open_world_hint: None,
                },
            ),
            icons: None,
        }
        "###);
    }

    #[test]
    fn custom_scalar_with_map() {
        let custom_scalar_map =
            CustomScalarMap::from_str("{ \"RealCustomScalar\": { \"type\": \"string\" }}");

        let operation = Operation::from_document(
            RawOperation {
                source_text: "query QueryName($id: RealCustomScalar) { id }".to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &SCHEMA,
            custom_scalar_map.ok().as_ref(),
            MutationMode::None,
            false,
            false,
        )
        .unwrap()
        .unwrap();
        let tool = Tool::from(operation);

        insta::assert_debug_snapshot!(tool, @r###"
        Tool {
            name: "QueryName",
            title: None,
            description: Some(
                "The returned value is optional and has type `String`",
            ),
            input_schema: {
                "type": String("object"),
                "properties": Object {
                    "id": Object {
                        "$ref": String("#/definitions/RealCustomScalar"),
                    },
                },
                "definitions": Object {
                    "RealCustomScalar": Object {
                        "description": String("RealCustomScalar exists"),
                        "type": String("string"),
                    },
                },
            },
            output_schema: None,
            annotations: Some(
                ToolAnnotations {
                    title: None,
                    read_only_hint: Some(
                        true,
                    ),
                    destructive_hint: None,
                    idempotent_hint: None,
                    open_world_hint: None,
                },
            ),
            icons: None,
        }
        "###);
    }

    #[test]
    fn test_tool_description() {
        const SCHEMA: &str = r#"
        type Query {
          """
          Get a list of A
          """
          a(input: String!): [A]!

          """
          Get a B
          """
          b: B

          """
          Get a Z
          """
          z: Z
        }

        """
        A
        """
        type A {
          c: String
          d: D
        }

        """
        B
        """
        type B {
          d: D
          u: U
        }

        """
        D
        """
        type D {
          e: E
          f: String
          g: String
        }

        """
        E
        """
        enum E {
          """
          one
          """
          ONE
          """
          two
          """
          TWO
        }

        """
        F
        """
        scalar F

        """
        U
        """
        union U = M | W

        """
        M
        """
        type M {
          m: Int
        }

        """
        W
        """
        type W {
          w: Int
        }

        """
        Z
        """
        type Z {
          z: Int
          zz: Int
          zzz: Int
        }
        "#;

        let document = Parser::new().parse_ast(SCHEMA, "schema.graphql").unwrap();
        let schema = document.to_schema().unwrap();

        let operation = Operation::from_document(
            RawOperation {
                source_text: r###"
            query GetABZ($state: String!) {
              a(input: $input) {
                d {
                  e
                }
              }
              b {
                d {
                  ...JustF
                }
                u {
                  ... on M {
                    m
                  }
                  ... on W {
                    w
                  }
                }
              }
              z {
                ...JustZZZ
              }
            }

            fragment JustF on D {
              f
            }

            fragment JustZZZ on Z {
              zzz
            }
            "###
                .to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &schema,
            None,
            MutationMode::None,
            false,
            false,
        )
        .unwrap()
        .unwrap();

        insta::assert_snapshot!(
            operation.tool.description.unwrap(),
            @r#"
        Get a list of A
        The returned value is an array of type `A`
        ---
        Get a B
        The returned value is optional and has type `B`
        ---
        Get a Z
        The returned value is optional and has type `Z`
        ---
        """A"""
        type A {
          d: D
        }

        """B"""
        type B {
          d: D
          u: U
        }

        """D"""
        type D {
          e: E
          f: String
        }

        """E"""
        enum E {
          """one"""
          ONE
          """two"""
          TWO
        }

        """U"""
        union U = M | W

        """M"""
        type M {
          m: Int
        }

        """W"""
        type W {
          w: Int
        }

        """Z"""
        type Z {
          zzz: Int
        }
        "#
        );
    }

    #[test]
    fn tool_comment_description() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: r###"
            # Overridden tool #description
            query GetABZ($state: String!) {
              b {
                d {
                  f
                }
              }
            }
            "###
                .to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &SCHEMA,
            None,
            MutationMode::None,
            false,
            false,
        )
        .unwrap()
        .unwrap();

        insta::assert_snapshot!(
            operation.tool.description.unwrap(),
            @"Overridden tool #description"
        );
    }

    #[test]
    fn tool_empty_comment_description() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: r###"
            #

            #
            query GetABZ($state: String!) {
              id
            }
            "###
                .to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &SCHEMA,
            None,
            MutationMode::None,
            false,
            false,
        )
        .unwrap()
        .unwrap();

        insta::assert_snapshot!(
            operation.tool.description.unwrap(),
            @"The returned value is optional and has type `String`"
        );
    }

    #[test]
    fn no_schema_description() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: r###"query GetABZ($state: String!) { id enum }"###.to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &SCHEMA,
            None,
            MutationMode::None,
            false,
            true,
        )
        .unwrap()
        .unwrap();

        insta::assert_snapshot!(
            operation.tool.description.unwrap(),
            @r"
        The returned value is optional and has type `String`
        ---
        The returned value is optional and has type `RealEnum`
        "
        );
    }

    #[test]
    fn no_type_description() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: r###"query GetABZ($state: String!) { id enum }"###.to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &SCHEMA,
            None,
            MutationMode::None,
            true,
            false,
        )
        .unwrap()
        .unwrap();

        insta::assert_snapshot!(
            operation.tool.description.unwrap(),
            @r#"
        ---
        """the description for the enum"""
        enum RealEnum {
          """ENUM_VALUE_1 is a value"""
          ENUM_VALUE_1
          """ENUM_VALUE_2 is a value"""
          ENUM_VALUE_2
        }
        "#
        );
    }

    #[test]
    fn no_type_description_or_schema_description() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: r###"query GetABZ($state: String!) { id enum }"###.to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &SCHEMA,
            None,
            MutationMode::None,
            true,
            true,
        )
        .unwrap()
        .unwrap();

        insta::assert_snapshot!(
            operation.tool.description.unwrap(),
            @""
        );
    }

    #[test]
    fn recursive_inputs() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: r###"query Test($filter: Filter){
                field(filter: $filter) {
                    id
                }
            }"###
                    .to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &Schema::parse(
                r#"
                """the filter input"""
                input Filter {
                """the filter.field field"""
                    field: String
                    """the filter.filter field"""
                    filter: Filter
                }
                type Query {
                """the Query.field field"""
                  field(
                    """the filter argument"""
                    filter: Filter
                  ): String
                }
            "#,
                "operation.graphql",
            )
            .unwrap(),
            None,
            MutationMode::None,
            true,
            true,
        )
        .unwrap()
        .unwrap();

        insta::assert_debug_snapshot!(operation.tool, @r###"
        Tool {
            name: "Test",
            title: None,
            description: Some(
                "",
            ),
            input_schema: {
                "type": String("object"),
                "properties": Object {
                    "filter": Object {
                        "description": String("the filter argument"),
                        "$ref": String("#/definitions/Filter"),
                    },
                },
                "definitions": Object {
                    "Filter": Object {
                        "description": String("the filter input"),
                        "type": String("object"),
                        "properties": Object {
                            "field": Object {
                                "description": String("the filter.field field"),
                                "type": String("string"),
                            },
                            "filter": Object {
                                "description": String("the filter.filter field"),
                                "$ref": String("#/definitions/Filter"),
                            },
                        },
                    },
                },
            },
            output_schema: None,
            annotations: Some(
                ToolAnnotations {
                    title: None,
                    read_only_hint: Some(
                        true,
                    ),
                    destructive_hint: None,
                    idempotent_hint: None,
                    open_world_hint: None,
                },
            ),
            icons: None,
        }
        "###);
    }

    #[test]
    fn with_variable_overrides() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: "query QueryName($id: ID, $name: String) { id }".to_string(),
                persisted_query_id: None,
                headers: None,
                variables: Some(HashMap::from([(
                    "id".to_string(),
                    serde_json::Value::String("v".to_string()),
                )])),
                source_path: None,
            },
            &SCHEMA,
            None,
            MutationMode::None,
            false,
            false,
        )
        .unwrap()
        .unwrap();
        let tool = Tool::from(operation);

        insta::assert_debug_snapshot!(tool, @r###"
        Tool {
            name: "QueryName",
            title: None,
            description: Some(
                "The returned value is optional and has type `String`",
            ),
            input_schema: {
                "type": String("object"),
                "properties": Object {
                    "name": Object {
                        "type": String("string"),
                    },
                },
            },
            output_schema: None,
            annotations: Some(
                ToolAnnotations {
                    title: None,
                    read_only_hint: Some(
                        true,
                    ),
                    destructive_hint: None,
                    idempotent_hint: None,
                    open_world_hint: None,
                },
            ),
            icons: None,
        }
        "###);
    }

    #[test]
    fn input_schema_includes_variable_descriptions() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: "query QueryName($idArg: ID) { customQuery(id: $idArg) { id } }"
                    .to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &SCHEMA,
            None,
            MutationMode::None,
            false,
            false,
        )
        .unwrap()
        .unwrap();
        let tool = Tool::from(operation);

        let json = to_sorted_json!(tool.input_schema);
        insta::assert_snapshot!(serde_json::to_string_pretty(&json).unwrap(), @r###"
        {
          "properties": {
            "idArg": {
              "description": "id description",
              "type": "string"
            }
          },
          "type": "object"
        }
        "###);
    }

    #[test]
    fn input_schema_includes_joined_variable_descriptions_if_multiple() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: "query QueryName($idArg: ID, $flag: Boolean) { customQuery(id: $idArg, flag: $flag) { id @skip(if: $flag) } }".to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &SCHEMA,
            None,
            MutationMode::None,
            false,
            false,
        )
            .unwrap()
            .unwrap();
        let tool = Tool::from(operation);

        let json = to_sorted_json!(tool.input_schema);
        insta::assert_snapshot!(serde_json::to_string_pretty(&json).unwrap(), @r###"
        {
          "properties": {
            "flag": {
              "description": "Skipped when true.#a flag",
              "type": "boolean"
            },
            "idArg": {
              "description": "id description",
              "type": "string"
            }
          },
          "type": "object"
        }
        "###);
    }

    #[test]
    fn input_schema_includes_directive_variable_descriptions() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: "query QueryName($idArg: ID, $skipArg: Boolean) { customQuery(id: $idArg) { id @skip(if: $skipArg) } }".to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &SCHEMA,
            None,
            MutationMode::None,
            false,
            false,
        )
            .unwrap()
            .unwrap();
        let tool = Tool::from(operation);

        insta::assert_snapshot!(serde_json::to_string_pretty(&serde_json::json!(tool.input_schema)).unwrap(), @r#"
        {
          "type": "object",
          "properties": {
            "idArg": {
              "description": "id description",
              "type": "string"
            },
            "skipArg": {
              "description": "Skipped when true.",
              "type": "boolean"
            }
          }
        }
        "#);
    }

    #[test]
    fn test_operation_name_with_named_query() {
        let source_text = "query GetUser($id: ID!) { user(id: $id) { name email } }";
        let raw_op = RawOperation {
            source_text: source_text.to_string(),
            persisted_query_id: None,
            headers: None,
            variables: None,
            source_path: None,
        };
        let operation =
            Operation::from_document(raw_op, &SCHEMA, None, MutationMode::None, false, false)
                .unwrap()
                .unwrap();

        let op_details = operation.operation(Value::Null).unwrap();
        assert_eq!(op_details.operation_name, Some(String::from("GetUser")));
    }

    #[test]
    fn test_operation_name_with_named_mutation() {
        let source_text =
            "mutation CreateUser($input: UserInput!) { createUser(input: $input) { id name } }";
        let raw_op = RawOperation {
            source_text: source_text.to_string(),
            persisted_query_id: None,
            headers: None,
            variables: None,
            source_path: None,
        };
        let operation =
            Operation::from_document(raw_op, &SCHEMA, None, MutationMode::Explicit, false, false)
                .unwrap()
                .unwrap();

        let op_details = operation.operation(Value::Null).unwrap();
        assert_eq!(op_details.operation_name, Some(String::from("CreateUser")));
    }

    #[test]
    fn operation_variable_comments_override_schema_descriptions() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: "# operation description\nquery QueryName(# id comment override\n$idArg: ID) { customQuery(id: $idArg) { id } }".to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &SCHEMA,
            None,
            MutationMode::None,
            false,
            false,
        )
            .unwrap()
            .unwrap();
        let tool = Tool::from(operation);

        let json = to_sorted_json!(tool.input_schema);
        insta::assert_snapshot!(serde_json::to_string_pretty(&json).unwrap(), @r###"
        {
          "properties": {
            "idArg": {
              "description": "id comment override",
              "type": "string"
            }
          },
          "type": "object"
        }
        "###);
    }

    #[test]
    fn operation_variable_comment_override_supports_multiline_comments() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: "# operation description\nquery QueryName(# id comment override\n # multi-line comment \n$idArg: ID) { customQuery(id: $idArg) { id } }".to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &SCHEMA,
            None,
            MutationMode::None,
            false,
            false,
        )
            .unwrap()
            .unwrap();
        let tool = Tool::from(operation);

        let json = to_sorted_json!(tool.input_schema);
        insta::assert_snapshot!(serde_json::to_string_pretty(&json).unwrap(), @r###"
        {
          "properties": {
            "idArg": {
              "description": "id comment override\n multi-line comment",
              "type": "string"
            }
          },
          "type": "object"
        }
        "###);
    }

    #[test]
    fn comment_with_parens_has_comments_extracted_correctly() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: "query QueryName # a comment (with parens)\n(# id comment override\n # multi-line comment \n$idArg: ID) { customQuery(id: $idArg) { id } }".to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &SCHEMA,
            None,
            MutationMode::None,
            false,
            false,
        )
            .unwrap()
            .unwrap();
        let tool = Tool::from(operation);

        let json = to_sorted_json!(tool.input_schema);
        insta::assert_snapshot!(serde_json::to_string_pretty(&json).unwrap(), @r###"
        {
          "properties": {
            "idArg": {
              "description": "id comment override\n multi-line comment",
              "type": "string"
            }
          },
          "type": "object"
        }
        "###);
    }

    #[test]
    fn multiline_comment_with_odd_spacing_and_parens_has_comments_extracted_correctly() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: "#  operation comment\n\nquery QueryName # a comment \n#     extra space\n\n\n#  blank lines (with parens)\n\n# another (paren)\n(# id comment override\n # multi-line comment \n$idArg: ID\n, \n# a flag\n$flag: Boolean) { customQuery(id: $idArg, skip: $flag) { id } }".to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &SCHEMA,
            None,
            MutationMode::None,
            false,
            false,
        )
            .unwrap()
            .unwrap();
        let tool = Tool::from(operation);

        let json = to_sorted_json!(tool.input_schema);
        insta::assert_snapshot!(serde_json::to_string_pretty(&json).unwrap(), @r###"
        {
          "properties": {
            "flag": {
              "description": "a flag",
              "type": "boolean"
            },
            "idArg": {
              "description": "id comment override\n multi-line comment",
              "type": "string"
            }
          },
          "type": "object"
        }
        "###);
    }

    #[test]
    fn operation_with_no_variables_is_handled_properly() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: "query QueryName { customQuery(id: \"123\") { id } }".to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &SCHEMA,
            None,
            MutationMode::None,
            false,
            false,
        )
        .unwrap()
        .unwrap();
        let tool = Tool::from(operation);

        let json = to_sorted_json!(tool.input_schema);
        insta::assert_snapshot!(serde_json::to_string_pretty(&json).unwrap(), @r###"
        {
          "properties": {},
          "type": "object"
        }
        "###);
    }

    #[test]
    fn commas_between_variables_are_ignored() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: "query QueryName(# id arg\n $idArg: ID,,\n,,\n # a flag\n $flag: Boolean,  ,,) { customQuery(id: $idArg, flag: $flag) { id } }".to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &SCHEMA,
            None,
            MutationMode::None,
            false,
            false,
        )
            .unwrap()
            .unwrap();
        let tool = Tool::from(operation);

        let json = to_sorted_json!(tool.input_schema);
        insta::assert_snapshot!(serde_json::to_string_pretty(&json).unwrap(), @r###"
        {
          "properties": {
            "flag": {
              "description": "a flag",
              "type": "boolean"
            },
            "idArg": {
              "description": "id arg",
              "type": "string"
            }
          },
          "type": "object"
        }
        "###);
    }

    #[test]
    fn input_schema_include_properties_field_even_when_operation_has_no_input_args() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: "query TestOp { testOp { id } }".to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &SCHEMA,
            None,
            MutationMode::None,
            false,
            false,
        )
        .unwrap()
        .unwrap();
        let tool = Tool::from(operation);

        let json = to_sorted_json!(tool.input_schema);
        insta::assert_snapshot!(serde_json::to_string_pretty(&json).unwrap(), @r#"
        {
          "properties": {},
          "type": "object"
        }
        "#);
    }

    #[test]
    fn nullable_list_of_nullable_input_objects() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: "query QueryName($objects: [RealInputObject]) { id }".to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &SCHEMA,
            None,
            MutationMode::None,
            false,
            false,
        )
        .unwrap()
        .unwrap();
        let tool = Tool::from(operation);

        insta::assert_debug_snapshot!(tool, @r###"
        Tool {
            name: "QueryName",
            title: None,
            description: Some(
                "The returned value is optional and has type `String`",
            ),
            input_schema: {
                "type": String("object"),
                "properties": Object {
                    "objects": Object {
                        "type": String("array"),
                        "items": Object {
                            "oneOf": Array [
                                Object {
                                    "$ref": String("#/definitions/RealInputObject"),
                                },
                                Object {
                                    "type": String("null"),
                                },
                            ],
                        },
                    },
                },
                "definitions": Object {
                    "RealInputObject": Object {
                        "type": String("object"),
                        "properties": Object {
                            "optional": Object {
                                "description": String("optional is a input field that is optional"),
                                "type": String("string"),
                            },
                            "required": Object {
                                "description": String("required is a input field that is required"),
                                "type": String("string"),
                            },
                        },
                        "required": Array [
                            String("required"),
                        ],
                    },
                },
            },
            output_schema: None,
            annotations: Some(
                ToolAnnotations {
                    title: None,
                    read_only_hint: Some(
                        true,
                    ),
                    destructive_hint: None,
                    idempotent_hint: None,
                    open_world_hint: None,
                },
            ),
            icons: None,
        }
        "###);

        let json = to_sorted_json!(tool.input_schema);
        insta::assert_snapshot!(serde_json::to_string_pretty(&json).unwrap(), @r###"
        {
          "definitions": {
            "RealInputObject": {
              "properties": {
                "optional": {
                  "description": "optional is a input field that is optional",
                  "type": "string"
                },
                "required": {
                  "description": "required is a input field that is required",
                  "type": "string"
                }
              },
              "required": [
                "required"
              ],
              "type": "object"
            }
          },
          "properties": {
            "objects": {
              "items": {
                "oneOf": [
                  {
                    "$ref": "#/definitions/RealInputObject"
                  },
                  {
                    "type": "null"
                  }
                ]
              },
              "type": "array"
            }
          },
          "type": "object"
        }
        "###);
    }

    #[test]
    fn non_nullable_list_of_non_nullable_input_objects() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: "query QueryName($objects: [RealInputObject!]!) { id }".to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &SCHEMA,
            None,
            MutationMode::None,
            false,
            false,
        )
        .unwrap()
        .unwrap();
        let tool = Tool::from(operation);

        insta::assert_debug_snapshot!(tool, @r###"
        Tool {
            name: "QueryName",
            title: None,
            description: Some(
                "The returned value is optional and has type `String`",
            ),
            input_schema: {
                "type": String("object"),
                "properties": Object {
                    "objects": Object {
                        "type": String("array"),
                        "items": Object {
                            "$ref": String("#/definitions/RealInputObject"),
                        },
                    },
                },
                "required": Array [
                    String("objects"),
                ],
                "definitions": Object {
                    "RealInputObject": Object {
                        "type": String("object"),
                        "properties": Object {
                            "optional": Object {
                                "description": String("optional is a input field that is optional"),
                                "type": String("string"),
                            },
                            "required": Object {
                                "description": String("required is a input field that is required"),
                                "type": String("string"),
                            },
                        },
                        "required": Array [
                            String("required"),
                        ],
                    },
                },
            },
            output_schema: None,
            annotations: Some(
                ToolAnnotations {
                    title: None,
                    read_only_hint: Some(
                        true,
                    ),
                    destructive_hint: None,
                    idempotent_hint: None,
                    open_world_hint: None,
                },
            ),
            icons: None,
        }
        "###);

        let json = to_sorted_json!(tool.input_schema);
        insta::assert_snapshot!(serde_json::to_string_pretty(&json).unwrap(), @r###"
        {
          "definitions": {
            "RealInputObject": {
              "properties": {
                "optional": {
                  "description": "optional is a input field that is optional",
                  "type": "string"
                },
                "required": {
                  "description": "required is a input field that is required",
                  "type": "string"
                }
              },
              "required": [
                "required"
              ],
              "type": "object"
            }
          },
          "properties": {
            "objects": {
              "items": {
                "$ref": "#/definitions/RealInputObject"
              },
              "type": "array"
            }
          },
          "required": [
            "objects"
          ],
          "type": "object"
        }
        "###);
    }

    #[test]
    fn subscriptions() {
        assert!(
            Operation::from_document(
                RawOperation {
                    source_text: "subscription SubscriptionName { id }".to_string(),
                    persisted_query_id: None,
                    headers: None,
                    variables: None,
                    source_path: None,
                },
                &SCHEMA,
                None,
                MutationMode::None,
                false,
                false,
            )
            .unwrap()
            .is_none()
        );
    }

    #[test]
    fn mutation_mode_none() {
        assert!(
            Operation::from_document(
                RawOperation {
                    source_text: "mutation MutationName { id }".to_string(),
                    persisted_query_id: None,
                    headers: None,
                    variables: None,
                    source_path: None,
                },
                &SCHEMA,
                None,
                MutationMode::None,
                false,
                false,
            )
            .ok()
            .unwrap()
            .is_none()
        );
    }

    #[test]
    fn mutation_mode_explicit() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: "mutation MutationName { id }".to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &SCHEMA,
            None,
            MutationMode::Explicit,
            false,
            false,
        )
        .unwrap()
        .unwrap();

        insta::assert_debug_snapshot!(operation, @r###"
        Operation {
            tool: Tool {
                name: "MutationName",
                title: None,
                description: Some(
                    "The returned value is optional and has type `String`",
                ),
                input_schema: {
                    "type": String("object"),
                    "properties": Object {},
                },
                output_schema: None,
                annotations: Some(
                    ToolAnnotations {
                        title: None,
                        read_only_hint: Some(
                            false,
                        ),
                        destructive_hint: None,
                        idempotent_hint: None,
                        open_world_hint: None,
                    },
                ),
                icons: None,
            },
            inner: RawOperation {
                source_text: "mutation MutationName { id }",
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            operation_name: "MutationName",
        }
        "###);
    }

    #[test]
    fn mutation_mode_all() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: "mutation MutationName { id }".to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &SCHEMA,
            None,
            MutationMode::All,
            false,
            false,
        )
        .unwrap()
        .unwrap();

        insta::assert_debug_snapshot!(operation, @r###"
        Operation {
            tool: Tool {
                name: "MutationName",
                title: None,
                description: Some(
                    "The returned value is optional and has type `String`",
                ),
                input_schema: {
                    "type": String("object"),
                    "properties": Object {},
                },
                output_schema: None,
                annotations: Some(
                    ToolAnnotations {
                        title: None,
                        read_only_hint: Some(
                            false,
                        ),
                        destructive_hint: None,
                        idempotent_hint: None,
                        open_world_hint: None,
                    },
                ),
                icons: None,
            },
            inner: RawOperation {
                source_text: "mutation MutationName { id }",
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            operation_name: "MutationName",
        }
        "###);
    }

    #[test]
    fn no_variables() {
        let operation = Operation::from_document(
            RawOperation {
                source_text: "query QueryName { id }".to_string(),
                persisted_query_id: None,
                headers: None,
                variables: None,
                source_path: None,
            },
            &SCHEMA,
            None,
            MutationMode::None,
            false,
            false,
        )
        .unwrap()
        .unwrap();
        let tool = Tool::from(operation);

        insta::assert_debug_snapshot!(tool, @r###"
        Tool {
            name: "QueryName",
            title: None,
            description: Some(
                "The returned value is optional and has type `String`",
            ),
            input_schema: {
                "type": String("object"),
                "properties": Object {},
            },
            output_schema: None,
            annotations: Some(
                ToolAnnotations {
                    title: None,
                    read_only_hint: Some(
                        true,
                    ),
                    destructive_hint: None,
                    idempotent_hint: None,
                    open_world_hint: None,
                },
            ),
            icons: None,
        }
        "###);
        insta::assert_snapshot!(serde_json::to_string_pretty(&serde_json::json!(tool.input_schema)).unwrap(), @r#"
        {
          "type": "object",
          "properties": {}
        }
        "#);
    }
}

```
Page 5/6FirstPrevNextLast