Add default collection support (#151)

* wip

* platform-api

* update operations

* update

* update schema

* add changeset

* update for the non null

* make graph ref required for default

* change to debug

* Update crates/apollo-mcp-registry/src/platform_api/operation_collections/collection_poller.rs

Co-authored-by: Matthew Hawkins <[email protected]>

* use ref

---------

Co-authored-by: Matthew Hawkins <[email protected]>

Author: Jephuff

.changesets/feat_jeffrey_add_default_collection_support.md

@@ -0,0 +1,3 @@
+### Pass in --collection default to use default collection - @Jephuff PR #151
+
+--collection default will use the configured default collection on the graph variant specified by the --apollo-graph-ref arg

crates/apollo-mcp-registry/src/platform_api/operation_collections/collection_poller.rs

@@ -9,6 +9,15 @@ use tokio_stream::wrappers::ReceiverStream;
 
 use super::{error::CollectionError, event::CollectionEvent};
 use crate::platform_api::PlatformApiConfig;
+use operation_collection_default_polling_query::{
+    OperationCollectionDefaultPollingQueryVariant as PollingDefaultGraphVariant,
+    OperationCollectionDefaultPollingQueryVariantOnGraphVariantMcpDefaultCollection as PollingDefaultCollection,
+};
+use operation_collection_default_query::{
+    OperationCollectionDefaultQueryVariant,
+    OperationCollectionDefaultQueryVariantOnGraphVariantMcpDefaultCollection as DefaultCollectionResult,
+    OperationCollectionDefaultQueryVariantOnGraphVariantMcpDefaultCollectionOnOperationCollectionOperations as OperationCollectionDefaultEntry,
+};
 use operation_collection_entries_query::OperationCollectionEntriesQueryOperationCollectionEntries;
 use operation_collection_polling_query::{
     OperationCollectionPollingQueryOperationCollection as PollingOperationCollectionResult,
@@ -55,6 +64,24 @@ struct OperationCollectionPollingQuery;
 )]
 struct OperationCollectionQuery;
 
+#[derive(GraphQLQuery)]
+#[graphql(
+    query_path = "src/platform_api/operation_collections/operation_collections.graphql",
+    schema_path = "src/platform_api/platform-api.graphql",
+    request_derives = "Debug",
+    response_derives = "PartialEq, Debug, Deserialize"
+)]
+struct OperationCollectionDefaultQuery;
+
+#[derive(GraphQLQuery)]
+#[graphql(
+    query_path = "src/platform_api/operation_collections/operation_collections.graphql",
+    schema_path = "src/platform_api/platform-api.graphql",
+    request_derives = "Debug",
+    response_derives = "PartialEq, Debug, Deserialize"
+)]
+struct OperationCollectionDefaultPollingQuery;
+
 async fn handle_poll_result(
     previous_updated_at: &mut HashMap<String, OperationData>,
     poll: Vec<(String, String)>,
@@ -89,7 +116,7 @@ async fn handle_poll_result(
     }
 
     if !changed_ids.is_empty() {
-        tracing::info!("changed operation ids: {:?}", changed_ids);
+        tracing::debug!("changed operation ids: {:?}", changed_ids);
         let full_response = graphql_request::<OperationCollectionEntriesQuery>(
             &OperationCollectionEntriesQuery::build_query(
                 operation_collection_entries_query::Variables {
@@ -176,10 +203,40 @@ impl From<&OperationCollectionEntriesQueryOperationCollectionEntries> for Operat
         }
     }
 }
+impl From<&OperationCollectionDefaultEntry> for OperationData {
+    fn from(operation: &OperationCollectionDefaultEntry) -> Self {
+        Self {
+            id: operation.id.clone(),
+            last_updated_at: operation.last_updated_at.clone(),
+            source_text: operation
+                .operation_data
+                .current_operation_revision
+                .body
+                .clone(),
+            headers: operation
+                .operation_data
+                .current_operation_revision
+                .headers
+                .as_ref()
+                .map(|headers| {
+                    headers
+                        .iter()
+                        .map(|h| (h.name.clone(), h.value.clone()))
+                        .collect()
+                }),
+            variables: operation
+                .operation_data
+                .current_operation_revision
+                .variables
+                .clone(),
+        }
+    }
+}
 
 #[derive(Clone)]
 pub enum CollectionSource {
     Id(String, PlatformApiConfig),
+    Default(String, PlatformApiConfig),
 }
 
 async fn write_init_response(
@@ -213,11 +270,14 @@ async fn write_init_response(
     }
 }
 impl CollectionSource {
-    pub fn into_stream(&self) -> Pin<Box<dyn Stream<Item = CollectionEvent> + Send>> {
+    pub fn into_stream(self) -> Pin<Box<dyn Stream<Item = CollectionEvent> + Send>> {
         match self {
-            CollectionSource::Id(id, platform_api_config) => {
+            CollectionSource::Id(ref id, ref platform_api_config) => {
                 self.collection_id_stream(id.clone(), platform_api_config.clone())
             }
+            CollectionSource::Default(ref graph_ref, ref platform_api_config) => {
+                self.default_collection_stream(graph_ref.clone(), platform_api_config.clone())
+            }
         }
     }
 
@@ -318,6 +378,134 @@ impl CollectionSource {
         });
         Box::pin(ReceiverStream::new(receiver))
     }
+
+    pub fn default_collection_stream(
+        &self,
+        graph_ref: String,
+        platform_api_config: PlatformApiConfig,
+    ) -> Pin<Box<dyn Stream<Item = CollectionEvent> + Send>> {
+        let (sender, receiver) = channel(2);
+        tokio::task::spawn(async move {
+            let mut previous_updated_at = HashMap::new();
+            match graphql_request::<OperationCollectionDefaultQuery>(
+                &OperationCollectionDefaultQuery::build_query(
+                    operation_collection_default_query::Variables {
+                        graph_ref: graph_ref.clone(),
+                    },
+                ),
+                &platform_api_config,
+            )
+            .await
+            {
+                Ok(response) => match response.variant {
+                    Some(OperationCollectionDefaultQueryVariant::GraphVariant(variant)) => {
+                        match variant.mcp_default_collection {
+                            DefaultCollectionResult::OperationCollection(collection) => {
+                                let should_poll = write_init_response(
+                                    &sender,
+                                    &mut previous_updated_at,
+                                    collection.operations.iter().map(OperationData::from),
+                                )
+                                .await;
+                                if !should_poll {
+                                    return;
+                                }
+                            }
+                            DefaultCollectionResult::PermissionError(error) => {
+                                if let Err(e) = sender
+                                    .send(CollectionEvent::CollectionError(
+                                        CollectionError::Response(error.message),
+                                    ))
+                                    .await
+                                {
+                                    tracing::debug!(
+                                        "failed to send error to collection stream. This is likely to be because the server is shutting down: {e}"
+                                    );
+                                    return;
+                                }
+                            }
+                        }
+                    }
+                    Some(OperationCollectionDefaultQueryVariant::InvalidRefFormat(err)) => {
+                        if let Err(e) = sender
+                            .send(CollectionEvent::CollectionError(CollectionError::Response(
+                                err.message,
+                            )))
+                            .await
+                        {
+                            tracing::debug!(
+                                "failed to send error to collection stream. This is likely to be because the server is shutting down: {e}"
+                            );
+                            return;
+                        }
+                    }
+                    None => {
+                        if let Err(e) = sender
+                            .send(CollectionEvent::CollectionError(CollectionError::Response(
+                                format!("{graph_ref} not found"),
+                            )))
+                            .await
+                        {
+                            tracing::debug!(
+                                "failed to send error to collection stream. This is likely to be because the server is shutting down: {e}"
+                            );
+                        }
+                        return;
+                    }
+                },
+                Err(err) => {
+                    if let Err(e) = sender.send(CollectionEvent::CollectionError(err)).await {
+                        tracing::debug!(
+                            "failed to send error to collection stream. This is likely to be because the server is shutting down: {e}"
+                        );
+                    }
+                    return;
+                }
+            };
+
+            loop {
+                tokio::time::sleep(platform_api_config.poll_interval).await;
+
+                match poll_operation_collection_default(
+                    graph_ref.clone(),
+                    &platform_api_config,
+                    &mut previous_updated_at,
+                )
+                .await
+                {
+                    Ok(Some(operations)) => {
+                        let operations_count = operations.len();
+                        if let Err(e) = sender
+                            .send(CollectionEvent::UpdateOperationCollection(operations))
+                            .await
+                        {
+                            tracing::debug!(
+                                "failed to push to stream. This is likely to be because the server is shutting down: {e}"
+                            );
+                            break;
+                        } else if operations_count > MAX_COLLECTION_SIZE_FOR_POLLING {
+                            tracing::warn!(
+                                "Operation Collection polling disabled. Collection has {operations_count} operations which exceeds the maximum of {MAX_COLLECTION_SIZE_FOR_POLLING}."
+                            );
+                            break;
+                        }
+                    }
+                    Ok(None) => {
+                        tracing::debug!("Operation collection unchanged");
+                    }
+                    Err(err) => {
+                        if let Err(e) = sender.send(CollectionEvent::CollectionError(err)).await {
+                            tracing::debug!(
+                                "failed to send error to collection stream. This is likely to be because the server is shutting down: {e}"
+                            );
+                            break;
+                        }
+                    }
+                }
+            }
+        });
+        Box::pin(ReceiverStream::new(receiver))
+    }
 }
 
 async fn poll_operation_collection_id(
@@ -356,12 +544,56 @@ async fn poll_operation_collection_id(
     }
 }
 
+async fn poll_operation_collection_default(
+    graph_ref: String,
+    platform_api_config: &PlatformApiConfig,
+    previous_updated_at: &mut HashMap<String, OperationData>,
+) -> Result<Option<Vec<OperationData>>, CollectionError> {
+    let response = graphql_request::<OperationCollectionDefaultPollingQuery>(
+        &OperationCollectionDefaultPollingQuery::build_query(
+            operation_collection_default_polling_query::Variables { graph_ref },
+        ),
+        platform_api_config,
+    )
+    .await?;
+
+    match response.variant {
+        Some(PollingDefaultGraphVariant::GraphVariant(variant)) => {
+            match variant.mcp_default_collection {
+                PollingDefaultCollection::OperationCollection(collection) => {
+                    handle_poll_result(
+                        previous_updated_at,
+                        collection
+                            .operations
+                            .into_iter()
+                            .map(|operation| (operation.id, operation.last_updated_at))
+                            .collect(),
+                        platform_api_config,
+                    )
+                    .await
+                }
+
+                PollingDefaultCollection::PermissionError(error) => {
+                    Err(CollectionError::Response(error.message))
+                }
+            }
+        }
+        Some(PollingDefaultGraphVariant::InvalidRefFormat(err)) => {
+            Err(CollectionError::Response(err.message))
+        }
+        None => Err(CollectionError::Response(
+            "Default collection not found".to_string(),
+        )),
+    }
+}
+
 async fn graphql_request<Query>(
     request_body: &graphql_client::QueryBody<Query::Variables>,
     platform_api_config: &PlatformApiConfig,
 ) -> Result<Query::ResponseData, CollectionError>
 where
     Query: graphql_client::GraphQLQuery,
+    <Query as graphql_client::GraphQLQuery>::ResponseData: std::fmt::Debug,
 {
     let res = reqwest::Client::new()
         .post(platform_api_config.registry_url.clone())

crates/apollo-mcp-registry/src/platform_api/operation_collections/operation_collections.graphql

@@ -53,6 +53,53 @@ query OperationCollectionPollingQuery($operationCollectionId: ID!) {
     }
 }
 
+query OperationCollectionDefaultQuery($graphRef: ID!) {
+    variant(ref: $graphRef) {
+        __typename
+        ... on GraphVariant {
+          mcpDefaultCollection {
+            __typename
+            ... on OperationCollection {
+              operations {
+                lastUpdatedAt
+                id
+                ...OperationData
+              }
+            }
+            ... on PermissionError {
+              message
+            }
+          }
+        }
+        ... on InvalidRefFormat {
+            message
+        }
+    }
+}
+
+query OperationCollectionDefaultPollingQuery($graphRef: ID!) {
+    variant(ref: $graphRef) {
+        __typename
+        ... on GraphVariant {
+          mcpDefaultCollection {
+            __typename
+            ... on OperationCollection {
+              operations {
+                id
+                lastUpdatedAt
+              }
+            }
+            ... on PermissionError {
+              message
+            }
+          }
+        }
+        ... on InvalidRefFormat {
+            message
+        }
+    }
+}
+
 query OperationCollectionEntriesQuery($collectionEntryIds: [ID!]!) {
   operationCollectionEntries(collectionEntryIds: $collectionEntryIds) {
     id