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/proto/lsp.proto

@@ -753,28 +753,47 @@ message TextEdit {
     PointUtf16 lsp_range_end = 3;
 }
 
-message MultiLspQuery {
+message LspQuery {
     uint64 project_id = 1;
-    uint64 buffer_id = 2;
-    repeated VectorClockEntry version = 3;
-    oneof strategy {
-        AllLanguageServers all = 4;
-    }
+    uint64 lsp_request_id = 2;
     oneof request {
+        GetReferences get_references = 3;
+        GetDocumentColor get_document_color = 4;
         GetHover get_hover = 5;
         GetCodeActions get_code_actions = 6;
         GetSignatureHelp get_signature_help = 7;
         GetCodeLens get_code_lens = 8;
         GetDocumentDiagnostics get_document_diagnostics = 9;
-        GetDocumentColor get_document_color = 10;
-        GetDefinition get_definition = 11;
-        GetDeclaration get_declaration = 12;
-        GetTypeDefinition get_type_definition = 13;
-        GetImplementation get_implementation = 14;
-        GetReferences get_references = 15;
+        GetDefinition get_definition = 10;
+        GetDeclaration get_declaration = 11;
+        GetTypeDefinition get_type_definition = 12;
+        GetImplementation get_implementation = 13;
     }
 }
 
+message LspQueryResponse {
+    uint64 project_id = 1;
+    uint64 lsp_request_id = 2;
+    repeated LspResponse responses = 3;
+}
+
+message LspResponse {
+    oneof response {
+        GetHoverResponse get_hover_response = 1;
+        GetCodeActionsResponse get_code_actions_response = 2;
+        GetSignatureHelpResponse get_signature_help_response = 3;
+        GetCodeLensResponse get_code_lens_response = 4;
+        GetDocumentDiagnosticsResponse get_document_diagnostics_response = 5;
+        GetDocumentColorResponse get_document_color_response = 6;
+        GetDefinitionResponse get_definition_response = 8;
+        GetDeclarationResponse get_declaration_response = 9;
+        GetTypeDefinitionResponse get_type_definition_response = 10;
+        GetImplementationResponse get_implementation_response = 11;
+        GetReferencesResponse get_references_response = 12;
+    }
+    uint64 server_id = 7;
+}
+
 message AllLanguageServers {}
 
 message LanguageServerSelector {
@@ -798,27 +817,6 @@ message StopLanguageServers {
     bool all = 4;
 }
 
-message MultiLspQueryResponse {
-    repeated LspResponse responses = 1;
-}
-
-message LspResponse {
-    oneof response {
-        GetHoverResponse get_hover_response = 1;
-        GetCodeActionsResponse get_code_actions_response = 2;
-        GetSignatureHelpResponse get_signature_help_response = 3;
-        GetCodeLensResponse get_code_lens_response = 4;
-        GetDocumentDiagnosticsResponse get_document_diagnostics_response = 5;
-        GetDocumentColorResponse get_document_color_response = 6;
-        GetDefinitionResponse get_definition_response = 8;
-        GetDeclarationResponse get_declaration_response = 9;
-        GetTypeDefinitionResponse get_type_definition_response = 10;
-        GetImplementationResponse get_implementation_response = 11;
-        GetReferencesResponse get_references_response = 12;
-    }
-    uint64 server_id = 7;
-}
-
 message LspExtRunnables {
     uint64 project_id = 1;
     uint64 buffer_id = 2;
@@ -909,3 +907,30 @@ message PullWorkspaceDiagnostics {
     uint64 project_id = 1;
     uint64 server_id = 2;
 }
+
+// todo(lsp) remove after Zed Stable hits v0.204.x
+message MultiLspQuery {
+    uint64 project_id = 1;
+    uint64 buffer_id = 2;
+    repeated VectorClockEntry version = 3;
+    oneof strategy {
+        AllLanguageServers all = 4;
+    }
+    oneof request {
+        GetHover get_hover = 5;
+        GetCodeActions get_code_actions = 6;
+        GetSignatureHelp get_signature_help = 7;
+        GetCodeLens get_code_lens = 8;
+        GetDocumentDiagnostics get_document_diagnostics = 9;
+        GetDocumentColor get_document_color = 10;
+        GetDefinition get_definition = 11;
+        GetDeclaration get_declaration = 12;
+        GetTypeDefinition get_type_definition = 13;
+        GetImplementation get_implementation = 14;
+        GetReferences get_references = 15;
+    }
+}
+
+message MultiLspQueryResponse {
+    repeated LspResponse responses = 1;
+}

crates/proto/proto/zed.proto

@@ -393,7 +393,10 @@ message Envelope {
         GetCrashFilesResponse get_crash_files_response = 362;
 
         GitClone git_clone = 363;
-        GitCloneResponse git_clone_response = 364; // current max
+        GitCloneResponse git_clone_response = 364;
+
+        LspQuery lsp_query = 365;
+        LspQueryResponse lsp_query_response = 366; // current max
     }
 
     reserved 87 to 88;

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;