ROVER-245 Responding to PR comments

src/command/lsp/errors.rs

@@ -0,0 +1,21 @@
+use crate::composition::supergraph::config::resolver::{
+    LoadRemoteSubgraphsError, LoadSupergraphConfigError, ResolveSupergraphConfigError,
+};
+use anyhow::Error;
+use std::path::PathBuf;
+
+#[derive(thiserror::Error, Debug)]
+pub enum SupergraphConfigLazyResolutionError {
+    #[error("Could not instantiate Studio Client")]
+    StudioClientInitialisationFailed(#[from] Error),
+    #[error("Could not load remote subgraphs")]
+    LoadRemoteSubgraphsFailed(#[from] LoadRemoteSubgraphsError),
+    #[error("Could not load supergraph config from local file")]
+    LoadLocalSupergraphConfigFailed(#[from] LoadSupergraphConfigError),
+    #[error("Could not resolve local and remote elements into complete SupergraphConfig")]
+    ResolveSupergraphConfigFailed(#[from] ResolveSupergraphConfigError),
+    #[error("Path `{0}` does not point to a file")]
+    PathDoesNotPointToAFile(PathBuf),
+}
+#[derive(thiserror::Error, Debug)]
+pub enum CompositionError {}

src/command/lsp/mod.rs

@@ -1,6 +1,11 @@
+mod errors;
+
+use crate::command::lsp::errors::SupergraphConfigLazyResolutionError;
+use crate::command::lsp::errors::SupergraphConfigLazyResolutionError::PathDoesNotPointToAFile;
 use crate::composition::events::CompositionEvent;
 use crate::composition::runner::Runner;
 use crate::composition::supergraph::binary::OutputTarget;
+use crate::composition::supergraph::config::lazy::LazilyResolvedSupergraphConfig;
 use crate::composition::supergraph::config::resolver::{
     ResolveSupergraphConfigError, SupergraphConfigResolver,
 };
@@ -23,13 +28,15 @@ use apollo_language_server::{ApolloLanguageServer, Config};
 use camino::Utf8PathBuf;
 use clap::Parser;
 use futures::StreamExt;
+use rover_client::blocking::StudioClient;
 use serde::Serialize;
 use std::collections::HashMap;
 use std::env::temp_dir;
 use std::io::stdin;
 use tower_lsp::lsp_types::{Diagnostic, Range};
 use tower_lsp::Server;
 use tracing::debug;
+use url::Url;
 
 #[derive(Debug, Serialize, Parser)]
 pub struct Lsp {
@@ -60,57 +67,116 @@ impl Lsp {
 }
 
 async fn run_lsp(client_config: StudioClientConfig, lsp_opts: LspOpts) -> RoverResult<()> {
-    //TODO: Check this error handling is right i.e. if we don't get a supergraph.yaml passed
-    // this should fail rather than doing something else.
-    let supergraph_yaml_path = lsp_opts
-        .supergraph_yaml
-        .as_ref()
-        .and_then(|path| {
-            if path.is_relative() {
-                Some(
-                    Utf8PathBuf::try_from(std::env::current_dir().ok()?)
-                        .ok()?
-                        .join(path),
+    let supergraph_yaml_path = lsp_opts.supergraph_yaml.as_ref().and_then(|path| {
+        if path.is_relative() {
+            Some(
+                Utf8PathBuf::try_from(std::env::current_dir().ok()?)
+                    .ok()?
+                    .join(path),
+            )
+        } else {
+            Some(path.clone())
+        }
+    });
+
+    // Early return if there is no `supergraph.yaml` given as there is no further need to construct
+    // anything
+    let (service, socket) = match supergraph_yaml_path {
+        None => {
+            let (service, socket, _receiver) = ApolloLanguageServer::build_service(
+                Config {
+                    // TODO Do we need to worry about these now?
+                    root_uri: String::default(),
+                    enable_auto_composition: false,
+                    force_federation: false,
+                    disable_telemetry: false,
+                },
+                HashMap::new(),
+            );
+            (service, socket)
+        }
+        Some(supergraph_yaml_path) => {
+            let studio_client =
+                client_config.get_authenticated_client(&lsp_opts.plugin_opts.profile)?;
+            // Resolve Supergraph Config -> Lazy
+            let (lazily_resolved_supergraph_config, supergraph_content_root) =
+                generate_lazily_resolved_supergraph_config(
+                    &studio_client,
+                    supergraph_yaml_path.clone(),
                 )
-            } else {
-                Some(path.clone())
-            }
-        })
-        .ok_or_else(|| anyhow!("Could not find supergraph.yaml file."))?;
-    let supergraph_content_root = supergraph_yaml_path.parent().unwrap().to_path_buf();
+                .await?;
+            // Generate the config needed to spin up the Language Server
+            let (service, socket, _receiver) = ApolloLanguageServer::build_service(
+                Config {
+                    root_uri: supergraph_content_root,
+                    enable_auto_composition: false,
+                    force_federation: false,
+                    disable_telemetry: false,
+                },
+                HashMap::from_iter(
+                    lazily_resolved_supergraph_config
+                        .subgraphs()
+                        .iter()
+                        .map(|(a, b)| (a.to_string(), b.schema().clone())),
+                ),
+            );
+            let supergraph_yaml_url = Url::from_file_path(supergraph_yaml_path)
+                .map_err(|_| anyhow!("Failed to convert supergraph yaml path to url"))?;
+            // Start running composition
+            start_composition(
+                lazily_resolved_supergraph_config,
+                supergraph_yaml_url,
+                client_config,
+                studio_client,
+                lsp_opts,
+                service.inner().to_owned(),
+            )
+            .await;
+            (service, socket)
+        }
+    };
 
-    let studio_client = client_config.get_authenticated_client(&lsp_opts.plugin_opts.profile)?;
+    let stdin = tokio::io::stdin();
+    let stdout = tokio::io::stdout();
+    let server = Server::new(stdin, stdout, socket);
+    server.serve(service).await;
+    Ok(())
+}
 
+async fn generate_lazily_resolved_supergraph_config(
+    studio_client: &StudioClient,
+    supergraph_yaml_path: Utf8PathBuf,
+) -> Result<(LazilyResolvedSupergraphConfig, String), SupergraphConfigLazyResolutionError> {
     // Get the SupergraphConfig in a form we can use
     let supergraph_config = SupergraphConfigResolver::default()
-        .load_remote_subgraphs(&studio_client, None)
+        .load_remote_subgraphs(studio_client, None)
         .await?
         .load_from_file_descriptor(
             &mut stdin(),
             Some(&FileDescriptorType::File(supergraph_yaml_path.clone())),
         )?;
-    let lazily_resolved_supergraph_config = supergraph_config
-        .lazily_resolve_subgraphs(&supergraph_content_root)
-        .await?;
-
-    // Build the service to spin up the language server
-    let (service, socket, _receiver) = ApolloLanguageServer::build_service(
-        Config {
-            root_uri: supergraph_content_root.to_string(),
-            enable_auto_composition: false,
-            force_federation: false,
-            disable_telemetry: false,
-        },
-        HashMap::from_iter(
-            lazily_resolved_supergraph_config
-                .subgraphs()
-                .iter()
-                .map(|(a, b)| (a.to_string(), b.schema().clone())),
-        ),
-    );
-
-    let language_server = service.inner().clone();
+    if let Some(parent) = supergraph_yaml_path.parent() {
+        Ok((
+            supergraph_config
+                .lazily_resolve_subgraphs(&parent.to_owned())
+                .await?,
+            parent.to_string(),
+        ))
+    } else {
+        Err(PathDoesNotPointToAFile(
+            supergraph_yaml_path.into_std_path_buf(),
+        ))
+    }
+}
 
+async fn start_composition(
+    lazily_resolved_supergraph_config: LazilyResolvedSupergraphConfig,
+    supergraph_yaml_url: Url,
+    client_config: StudioClientConfig,
+    studio_client: StudioClient,
+    lsp_opts: LspOpts,
+    language_server: ApolloLanguageServer,
+) {
     // Spawn a separate thread to handle composition and passing that data to the language server
     tokio::spawn(async move {
         // Create a supergraph binary
@@ -153,13 +219,13 @@ async fn run_lsp(client_config: StudioClientConfig, lsp_opts: LspOpts) -> RoverR
             )
             .run();
 
-        let supergraph_yaml_url =
-            tower_lsp::lsp_types::Url::from_file_path(supergraph_yaml_path)
-                .map_err(|_| anyhow!("Failed to convert supergraph yaml path to url"))?;
-
         while let Some(event) = stream.next().await {
             match event {
                 CompositionEvent::Started => {
+                    // Even though it's hidden by library calls, this function emits a WorkDoneProgressBegin event,
+                    // which is paired with a WorkDoneProgressEnd event, sent by the `composition_did_update` function.
+                    // Any refactoring needs to ensure that we don't break this ordering, otherwise the LSP may well
+                    // cease to function in a useful way.
                     language_server.composition_did_start().await;
                 }
                 CompositionEvent::Success(CompositionSuccess {
@@ -199,7 +265,6 @@ async fn run_lsp(client_config: StudioClientConfig, lsp_opts: LspOpts) -> RoverR
                 }
                 CompositionEvent::Error(err) => {
                     debug!("Composition {federation_version} failed: {err}");
-                    // TODO: we could highlight the version of federation, since it failed.
                     let message = format!("Failed run composition {federation_version}: {err}",);
                     let diagnostic = Diagnostic::new_simple(Range::default(), message);
                     language_server
@@ -221,10 +286,4 @@ async fn run_lsp(client_config: StudioClientConfig, lsp_opts: LspOpts) -> RoverR
         }
         Ok::<(), Error>(())
     });
-
-    let stdin = tokio::io::stdin();
-    let stdout = tokio::io::stdout();
-    let server = Server::new(stdin, stdout, socket);
-    server.serve(service).await;
-    Ok(())
 }

src/composition/supergraph/config/full/subgraphs.rs

@@ -24,8 +24,8 @@ impl FullyResolvedSubgraphs {
     }
 
     /// Used to upsert a fully resolved subgraph into this object's definitions
-    pub fn upsert_subgraph(&mut self, name: String, schema: String) -> bool {
-        self.subgraphs.insert(name, schema).is_none()
+    pub fn upsert_subgraph(&mut self, name: String, schema: String) -> Option<String> {
+        self.subgraphs.insert(name, schema)
     }
 
     /// Removes a subgraph from this object's definitions

src/composition/supergraph/config/lazy/supergraph.rs

@@ -59,7 +59,11 @@ impl LazilyResolvedSupergraphConfig {
     }
 
     /// Fully resolves a [`LazilyResolvedSupergraphConfig`] into a [`FullyResolvedSubgraphs`]
-    /// by retrieving all the schemas as strings
+    /// by retrieving all the schemas as strings that represent the SDL.
+    ///
+    /// This is a one-way conversion by design as FullyResolvedSubgraphs are just SDL representations
+    /// of whatever source was used to produce them. As such it's impossible to go backwards once
+    /// you are left with the bare SDL representation of the subgraph.
     pub async fn fully_resolve_subgraphs(
         self,
         introspect_subgraph_impl: &impl IntrospectSubgraph,

src/composition/watchers/composition.rs

@@ -56,7 +56,10 @@ where
                         SubgraphEvent::SubgraphChanged(subgraph_schema_changed) => {
                             let name = subgraph_schema_changed.name();
                             let sdl = subgraph_schema_changed.sdl();
-                            if subgraphs.upsert_subgraph(name.to_string(), sdl.to_string()) {
+                            if subgraphs
+                                .upsert_subgraph(name.to_string(), sdl.to_string())
+                                .is_none()
+                            {
                                 let _ = sender
                                     .send(CompositionEvent::SubgraphAdded(
                                         CompositionSubgraphAdded {