[ty] Send a single request for registrations/unregistrations (#19822)

## Summary

This is a small refactor to update the server to send a single request
to perform registrations and unregistrations of dynamic capabilities.

## Test Plan

Existing E2E test cases pass, add a new test case to verify multiple
registrations.

crates/ty_server/src/session.rs

@@ -570,8 +570,7 @@ impl Session {
             self.global_settings = Arc::new(global_settings);
         }
 
-        self.register_diagnostic_capability(client);
+        self.register_capabilities(client);
-        self.register_rename_capability(client);
 
         assert!(
             self.workspaces.all_initialized(),
@@ -587,45 +586,44 @@ impl Session {
         }
     }
 
-    // TODO: Merge the following two methods as `register_capability` and `unregister_capability`
+    /// Registers the dynamic capabilities with the client as per the resolved global settings.
-
-    /// Sends a registration notification to the client to enable / disable workspace diagnostics
-    /// as per the `ty.diagnosticMode` global setting.
     ///
-    /// This method is a no-op if the client doesn't support dynamic registration of diagnostic
+    /// ## Diagnostic capability
-    /// capability.
+    ///
-    fn register_diagnostic_capability(&mut self, client: &Client) {
+    /// This capability is used to enable / disable workspace diagnostics as per the
+    /// `ty.diagnosticMode` global setting.
+    ///
+    /// ## Rename capability
+    ///
+    /// This capability is used to enable / disable rename functionality as per the
+    /// `ty.experimental.rename` global setting.
+    fn register_capabilities(&mut self, client: &Client) {
         static DIAGNOSTIC_REGISTRATION_ID: &str = "ty/textDocument/diagnostic";
+        static RENAME_REGISTRATION_ID: &str = "ty/textDocument/rename";
 
-        if !self
+        let mut registrations = vec![];
+        let mut unregistrations = vec![];
+
+        if self
             .resolved_client_capabilities
             .supports_diagnostic_dynamic_registration()
         {
-            return;
+            if self
-        }
-
-        let registered = self
                 .registrations
-            .contains(DocumentDiagnosticRequest::METHOD);
+                .contains(DocumentDiagnosticRequest::METHOD)
-
+            {
-        if registered {
+                unregistrations.push(Unregistration {
-            client.send_request::<UnregisterCapability>(
-                self,
-                UnregistrationParams {
-                    unregisterations: vec![Unregistration {
                     id: DIAGNOSTIC_REGISTRATION_ID.into(),
                     method: DocumentDiagnosticRequest::METHOD.into(),
-                    }],
+                });
-                },
-                |_: &Client, ()| {
-                    tracing::debug!("Unregistered diagnostic capability");
-                },
-            );
             }
 
             let diagnostic_mode = self.global_settings.diagnostic_mode;
 
-        let registration = Registration {
+            tracing::debug!(
+                "Registering diagnostic capability with {diagnostic_mode:?} diagnostic mode"
+            );
+            registrations.push(Registration {
                 id: DIAGNOSTIC_REGISTRATION_ID.into(),
                 method: DocumentDiagnosticRequest::METHOD.into(),
                 register_options: Some(
@@ -639,82 +637,86 @@ impl Session {
                     ))
                     .unwrap(),
                 ),
-        };
+            });
-
-        client.send_request::<RegisterCapability>(
-            self,
-            RegistrationParams {
-                registrations: vec![registration],
-            },
-            move |_: &Client, ()| {
-                tracing::debug!(
-                    "Registered diagnostic capability in {diagnostic_mode:?} diagnostic mode"
-                );
-            },
-        );
-
-        if !registered {
-            self.registrations
-                .insert(DocumentDiagnosticRequest::METHOD.to_string());
-        }
         }
 
-    /// Sends a registration notification to the client to enable / disable rename capability as
+        if self
-    /// per the `ty.experimental.rename` global setting.
-    ///
-    /// This method is a no-op if the client doesn't support dynamic registration of rename
-    /// capability.
-    fn register_rename_capability(&mut self, client: &Client) {
-        static RENAME_REGISTRATION_ID: &str = "ty/textDocument/rename";
-
-        if !self
             .resolved_client_capabilities
             .supports_rename_dynamic_registration()
         {
-            return;
+            let is_rename_enabled = self.global_settings.is_rename_enabled();
-        }
 
-        let registered = self.registrations.contains(Rename::METHOD);
+            if !is_rename_enabled {
-
+                tracing::debug!("Rename capability is disabled in the resolved global settings");
-        if registered {
+                if self.registrations.contains(Rename::METHOD) {
-            client.send_request::<UnregisterCapability>(
+                    unregistrations.push(Unregistration {
-                self,
-                UnregistrationParams {
-                    unregisterations: vec![Unregistration {
                         id: RENAME_REGISTRATION_ID.into(),
                         method: Rename::METHOD.into(),
-                    }],
+                    });
-                },
+                }
-                move |_: &Client, ()| {
-                    tracing::debug!("Unregistered rename capability");
-                },
-            );
             }
 
-        if !self.global_settings.experimental.rename {
+            if is_rename_enabled {
-            tracing::debug!("Rename capability is disabled in the client settings");
+                registrations.push(Registration {
-            return;
-        }
-
-        let registration = Registration {
                     id: RENAME_REGISTRATION_ID.into(),
                     method: Rename::METHOD.into(),
                     register_options: Some(serde_json::to_value(server_rename_options()).unwrap()),
-        };
+                });
+            }
+        }
+
+        // First, unregister any existing capabilities and then register or re-register them.
+        self.unregister_dynamic_capability(client, unregistrations);
+        self.register_dynamic_capability(client, registrations);
+    }
+
+    /// Registers a list of dynamic capabilities with the client.
+    fn register_dynamic_capability(&mut self, client: &Client, registrations: Vec<Registration>) {
+        if registrations.is_empty() {
+            return;
+        }
+
+        for registration in &registrations {
+            self.registrations.insert(registration.method.clone());
+        }
 
         client.send_request::<RegisterCapability>(
             self,
-            RegistrationParams {
+            RegistrationParams { registrations },
-                registrations: vec![registration],
+            |_: &Client, ()| {
-            },
+                tracing::debug!("Registered dynamic capabilities");
-            move |_: &Client, ()| {
-                tracing::debug!("Registered rename capability");
             },
         );
-
-        if !registered {
-            self.registrations.insert(Rename::METHOD.to_string());
     }
+
+    /// Unregisters a list of dynamic capabilities with the client.
+    fn unregister_dynamic_capability(
+        &mut self,
+        client: &Client,
+        unregistrations: Vec<Unregistration>,
+    ) {
+        if unregistrations.is_empty() {
+            return;
+        }
+
+        for unregistration in &unregistrations {
+            if !self.registrations.remove(&unregistration.method) {
+                tracing::debug!(
+                    "Unregistration for `{}` was requested, but it was not registered",
+                    unregistration.method
+                );
+            }
+        }
+
+        client.send_request::<UnregisterCapability>(
+            self,
+            UnregistrationParams {
+                unregisterations: unregistrations,
+            },
+            |_: &Client, ()| {
+                tracing::debug!("Unregistered dynamic capabilities");
+            },
+        );
     }
 
     /// Creates a document snapshot with the URL referencing the document to snapshot.

crates/ty_server/tests/e2e/initialize.rs

@@ -508,3 +508,53 @@ fn not_register_rename_capability_when_disabled() -> Result<()> {
 
     Ok(())
 }
+
+/// Tests that the server can register multiple capabilities at once.
+///
+/// This test would need to be updated when the server supports additional capabilities in the
+/// future.
+#[test]
+fn register_multiple_capabilities() -> Result<()> {
+    let workspace_root = SystemPath::new("foo");
+    let mut server = TestServerBuilder::new()?
+        .with_workspace(workspace_root, None)?
+        .with_initialization_options(
+            ClientOptions::default()
+                .with_experimental_rename(true)
+                .with_diagnostic_mode(DiagnosticMode::Workspace),
+        )
+        .enable_rename_dynamic_registration(true)
+        .enable_diagnostic_dynamic_registration(true)
+        .build()?
+        .wait_until_workspaces_are_initialized()?;
+
+    let (_, params) = server.await_request::<RegisterCapability>()?;
+    let registrations = params.registrations;
+
+    assert_eq!(registrations.len(), 2);
+
+    insta::assert_json_snapshot!(registrations, @r#"
+    [
+      {
+        "id": "ty/textDocument/diagnostic",
+        "method": "textDocument/diagnostic",
+        "registerOptions": {
+          "documentSelector": null,
+          "identifier": "ty",
+          "interFileDependencies": true,
+          "workDoneProgress": true,
+          "workspaceDiagnostics": true
+        }
+      },
+      {
+        "id": "ty/textDocument/rename",
+        "method": "textDocument/rename",
+        "registerOptions": {
+          "prepareProvider": true
+        }
+      }
+    ]
+    "#);
+
+    Ok(())
+}