diff --git a/crates/ty_server/src/session.rs b/crates/ty_server/src/session.rs index 9e08c026e3..1a21a3aec5 100644 --- a/crates/ty_server/src/session.rs +++ b/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_rename_capability(client); + self.register_capabilities(client); assert!( self.workspaces.all_initialized(), @@ -587,134 +586,137 @@ impl Session { } } - // TODO: Merge the following two methods as `register_capability` and `unregister_capability` - - /// Sends a registration notification to the client to enable / disable workspace diagnostics - /// as per the `ty.diagnosticMode` global setting. + /// Registers the dynamic capabilities with the client as per the resolved global settings. /// - /// This method is a no-op if the client doesn't support dynamic registration of diagnostic - /// capability. - fn register_diagnostic_capability(&mut self, client: &Client) { + /// ## Diagnostic capability + /// + /// 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 + .registrations + .contains(DocumentDiagnosticRequest::METHOD) + { + unregistrations.push(Unregistration { + id: DIAGNOSTIC_REGISTRATION_ID.into(), + method: DocumentDiagnosticRequest::METHOD.into(), + }); + } - let registered = self - .registrations - .contains(DocumentDiagnosticRequest::METHOD); + let diagnostic_mode = self.global_settings.diagnostic_mode; - if registered { - client.send_request::( - self, - UnregistrationParams { - unregisterations: vec![Unregistration { - id: DIAGNOSTIC_REGISTRATION_ID.into(), - method: DocumentDiagnosticRequest::METHOD.into(), - }], - }, - |_: &Client, ()| { - tracing::debug!("Unregistered diagnostic capability"); - }, + 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( + serde_json::to_value(DiagnosticServerCapabilities::RegistrationOptions( + DiagnosticRegistrationOptions { + diagnostic_options: server_diagnostic_options( + diagnostic_mode.is_workspace(), + ), + ..Default::default() + }, + )) + .unwrap(), + ), + }); } - let diagnostic_mode = self.global_settings.diagnostic_mode; - - let registration = Registration { - id: DIAGNOSTIC_REGISTRATION_ID.into(), - method: DocumentDiagnosticRequest::METHOD.into(), - register_options: Some( - serde_json::to_value(DiagnosticServerCapabilities::RegistrationOptions( - DiagnosticRegistrationOptions { - diagnostic_options: server_diagnostic_options( - diagnostic_mode.is_workspace(), - ), - ..Default::default() - }, - )) - .unwrap(), - ), - }; - - client.send_request::( - 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 - /// 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 + 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 registered { - client.send_request::( - self, - UnregistrationParams { - unregisterations: vec![Unregistration { + if !is_rename_enabled { + tracing::debug!("Rename capability is disabled in the resolved global settings"); + if self.registrations.contains(Rename::METHOD) { + unregistrations.push(Unregistration { id: RENAME_REGISTRATION_ID.into(), method: Rename::METHOD.into(), - }], - }, - move |_: &Client, ()| { - tracing::debug!("Unregistered rename capability"); - }, - ); + }); + } + } + + if is_rename_enabled { + registrations.push(Registration { + id: RENAME_REGISTRATION_ID.into(), + method: Rename::METHOD.into(), + register_options: Some(serde_json::to_value(server_rename_options()).unwrap()), + }); + } } - if !self.global_settings.experimental.rename { - tracing::debug!("Rename capability is disabled in the client settings"); + // 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) { + if registrations.is_empty() { return; } - let registration = Registration { - id: RENAME_REGISTRATION_ID.into(), - method: Rename::METHOD.into(), - register_options: Some(serde_json::to_value(server_rename_options()).unwrap()), - }; + for registration in ®istrations { + self.registrations.insert(registration.method.clone()); + } client.send_request::( self, - RegistrationParams { - registrations: vec![registration], - }, - move |_: &Client, ()| { - tracing::debug!("Registered rename capability"); + RegistrationParams { registrations }, + |_: &Client, ()| { + tracing::debug!("Registered dynamic capabilities"); }, ); + } - 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, + ) { + 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::( + self, + UnregistrationParams { + unregisterations: unregistrations, + }, + |_: &Client, ()| { + tracing::debug!("Unregistered dynamic capabilities"); + }, + ); } /// Creates a document snapshot with the URL referencing the document to snapshot. diff --git a/crates/ty_server/tests/e2e/initialize.rs b/crates/ty_server/tests/e2e/initialize.rs index ce696d717c..1a754ea029 100644 --- a/crates/ty_server/tests/e2e/initialize.rs +++ b/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::()?; + 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(()) +}