Stop waiting for part of LSP responses on remote Collab clients' part (#36557)

Instead of holding a connection for potentially long LSP queries (e.g.
rust-analyzer might take minutes to look up a definition), disconnect
right after sending the initial request and handle the follow-up
responses later.

As a bonus, this allows to cancel previously sent request on the local
Collab clients' side due to this, as instead of holding and serving the
old connection, local clients now can stop previous requests, if needed.

Current PR does not convert all LSP requests to the new paradigm, but
the problematic ones, deprecating `MultiLspQuery` and moving all its
requests to the new paradigm.

Release Notes:

- Improved resource usage when querying LSP over Collab

---------

Co-authored-by: David Kleingeld <git@davidsk.dev>
Co-authored-by: Mikayla Maki <mikayla@zed.dev>
Co-authored-by: David Kleingeld <davidsk@zed.dev>

crates/proto/src/macros.rs

@@ -69,3 +69,32 @@ macro_rules! entity_messages {
         })*
     };
 }
+
+#[macro_export]
+macro_rules! lsp_messages {
+    ($(($request_name:ident, $response_name:ident, $stop_previous_requests:expr)),* $(,)?) => {
+        $(impl LspRequestMessage for $request_name {
+            type Response = $response_name;
+
+            fn to_proto_query(self) -> $crate::lsp_query::Request {
+                $crate::lsp_query::Request::$request_name(self)
+            }
+
+            fn response_to_proto_query(response: Self::Response) -> $crate::lsp_response::Response {
+                $crate::lsp_response::Response::$response_name(response)
+            }
+
+            fn buffer_id(&self) -> u64 {
+                self.buffer_id
+            }
+
+            fn buffer_version(&self) -> &[$crate::VectorClockEntry] {
+                &self.version
+            }
+
+            fn stop_previous_requests() -> bool {
+                $stop_previous_requests
+            }
+        })*
+    };
+}

crates/proto/src/proto.rs

@@ -169,6 +169,9 @@ messages!(
     (MarkNotificationRead, Foreground),
     (MoveChannel, Foreground),
     (ReorderChannel, Foreground),
+    (LspQuery, Background),
+    (LspQueryResponse, Background),
+    // todo(lsp) remove after Zed Stable hits v0.204.x
     (MultiLspQuery, Background),
     (MultiLspQueryResponse, Background),
     (OnTypeFormatting, Background),
@@ -426,7 +429,10 @@ request_messages!(
     (SetRoomParticipantRole, Ack),
     (BlameBuffer, BlameBufferResponse),
     (RejoinRemoteProjects, RejoinRemoteProjectsResponse),
+    // todo(lsp) remove after Zed Stable hits v0.204.x
     (MultiLspQuery, MultiLspQueryResponse),
+    (LspQuery, Ack),
+    (LspQueryResponse, Ack),
     (RestartLanguageServers, Ack),
     (StopLanguageServers, Ack),
     (OpenContext, OpenContextResponse),
@@ -478,6 +484,20 @@ request_messages!(
     (GitClone, GitCloneResponse)
 );
 
+lsp_messages!(
+    (GetReferences, GetReferencesResponse, true),
+    (GetDocumentColor, GetDocumentColorResponse, true),
+    (GetHover, GetHoverResponse, true),
+    (GetCodeActions, GetCodeActionsResponse, true),
+    (GetSignatureHelp, GetSignatureHelpResponse, true),
+    (GetCodeLens, GetCodeLensResponse, true),
+    (GetDocumentDiagnostics, GetDocumentDiagnosticsResponse, true),
+    (GetDefinition, GetDefinitionResponse, true),
+    (GetDeclaration, GetDeclarationResponse, true),
+    (GetTypeDefinition, GetTypeDefinitionResponse, true),
+    (GetImplementation, GetImplementationResponse, true),
+);
+
 entity_messages!(
     {project_id, ShareProject},
     AddProjectCollaborator,
@@ -520,6 +540,9 @@ entity_messages!(
     LeaveProject,
     LinkedEditingRange,
     LoadCommitDiff,
+    LspQuery,
+    LspQueryResponse,
+    // todo(lsp) remove after Zed Stable hits v0.204.x
     MultiLspQuery,
     RestartLanguageServers,
     StopLanguageServers,
@@ -777,6 +800,28 @@ pub fn split_repository_update(
     }])
 }
 
+impl LspQuery {
+    pub fn query_name_and_write_permissions(&self) -> (&str, bool) {
+        match self.request {
+            Some(lsp_query::Request::GetHover(_)) => ("GetHover", false),
+            Some(lsp_query::Request::GetCodeActions(_)) => ("GetCodeActions", true),
+            Some(lsp_query::Request::GetSignatureHelp(_)) => ("GetSignatureHelp", false),
+            Some(lsp_query::Request::GetCodeLens(_)) => ("GetCodeLens", true),
+            Some(lsp_query::Request::GetDocumentDiagnostics(_)) => {
+                ("GetDocumentDiagnostics", false)
+            }
+            Some(lsp_query::Request::GetDefinition(_)) => ("GetDefinition", false),
+            Some(lsp_query::Request::GetDeclaration(_)) => ("GetDeclaration", false),
+            Some(lsp_query::Request::GetTypeDefinition(_)) => ("GetTypeDefinition", false),
+            Some(lsp_query::Request::GetImplementation(_)) => ("GetImplementation", false),
+            Some(lsp_query::Request::GetReferences(_)) => ("GetReferences", false),
+            Some(lsp_query::Request::GetDocumentColor(_)) => ("GetDocumentColor", false),
+            None => ("<unknown>", true),
+        }
+    }
+}
+
+// todo(lsp) remove after Zed Stable hits v0.204.x
 impl MultiLspQuery {
     pub fn request_str(&self) -> &str {
         match self.request {

crates/proto/src/typed_envelope.rs

@@ -31,6 +31,58 @@ pub trait RequestMessage: EnvelopedMessage {
     type Response: EnvelopedMessage;
 }
 
+/// A trait to bind LSP request and responses for the proto layer.
+/// Should be used for every LSP request that has to traverse through the proto layer.
+///
+/// `lsp_messages` macro in the same crate provides a convenient way to implement this.
+pub trait LspRequestMessage: EnvelopedMessage {
+    type Response: EnvelopedMessage;
+
+    fn to_proto_query(self) -> crate::lsp_query::Request;
+
+    fn response_to_proto_query(response: Self::Response) -> crate::lsp_response::Response;
+
+    fn buffer_id(&self) -> u64;
+
+    fn buffer_version(&self) -> &[crate::VectorClockEntry];
+
+    /// Whether to deduplicate the requests, or keep the previous ones running when another
+    /// request of the same kind is processed.
+    fn stop_previous_requests() -> bool;
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub struct LspRequestId(pub u64);
+
+/// A response from a single language server.
+/// There could be multiple responses for a single LSP request,
+/// from different servers.
+pub struct ProtoLspResponse<R> {
+    pub server_id: u64,
+    pub response: R,
+}
+
+impl ProtoLspResponse<Box<dyn AnyTypedEnvelope>> {
+    pub fn into_response<T: LspRequestMessage>(self) -> Result<ProtoLspResponse<T::Response>> {
+        let envelope = self
+            .response
+            .into_any()
+            .downcast::<TypedEnvelope<T::Response>>()
+            .map_err(|_| {
+                anyhow::anyhow!(
+                    "cannot downcast LspResponse to {} for message {}",
+                    T::Response::NAME,
+                    T::NAME,
+                )
+            })?;
+
+        Ok(ProtoLspResponse {
+            server_id: self.server_id,
+            response: envelope.payload,
+        })
+    }
+}
+
 pub trait AnyTypedEnvelope: Any + Send + Sync {
     fn payload_type_id(&self) -> TypeId;
     fn payload_type_name(&self) -> &'static str;