Remove stale SSH LSP log entries after server restarts

crates/collab/tests/integration/remote_editing_collaboration_tests.rs

@@ -21,7 +21,10 @@ use node_runtime::NodeRuntime;
 use project::{
     ProjectPath,
     debugger::session::ThreadId,
-    lsp_store::{FormatTrigger, LspFormatTarget},
+    lsp_store::{
+        FormatTrigger, LspFormatTarget,
+        log_store::{self, GlobalLogStore},
+    },
     trusted_worktrees::{PathTrust, TrustedWorktrees},
 };
 use remote::RemoteClient;
@@ -837,6 +840,150 @@ async fn test_ssh_collaboration_formatting_with_prettier(
     );
 }
 
+#[gpui::test(iterations = 10)]
+async fn test_ssh_restarting_language_server_replaces_remote_status(
+    executor: BackgroundExecutor,
+    cx_a: &mut TestAppContext,
+    server_cx: &mut TestAppContext,
+) {
+    cx_a.set_name("a");
+    server_cx.set_name("server");
+
+    cx_a.update(|cx| {
+        release_channel::init(semver::Version::new(0, 0, 0), cx);
+    });
+    server_cx.update(|cx| {
+        release_channel::init(semver::Version::new(0, 0, 0), cx);
+    });
+
+    let mut server = TestServer::start(executor.clone()).await;
+    let client_a = server.create_client(cx_a, "user_a").await;
+    let log_store = cx_a.update(|cx| log_store::init(false, cx));
+
+    let (opts, server_ssh, _) = RemoteClient::fake_server(cx_a, server_cx);
+    let remote_fs = FakeFs::new(server_cx.executor());
+    remote_fs
+        .insert_tree(path!("/project"), json!({ "a.rs": "fn main() {}" }))
+        .await;
+
+    client_a.language_registry().add(rust_lang());
+
+    server_cx.update(HeadlessProject::init);
+    let languages = Arc::new(LanguageRegistry::new(server_cx.executor()));
+    languages.add(rust_lang());
+    let mut fake_language_servers = languages.register_fake_lsp(
+        "Rust",
+        FakeLspAdapter {
+            name: "the-language-server",
+            ..Default::default()
+        },
+    );
+    let _headless_project = server_cx.new(|cx| {
+        HeadlessProject::new(
+            HeadlessAppState {
+                session: server_ssh,
+                fs: remote_fs.clone(),
+                http_client: Arc::new(BlockedHttpClient),
+                node_runtime: NodeRuntime::unavailable(),
+                languages,
+                extension_host_proxy: Arc::new(ExtensionHostProxy::new()),
+                startup_time: std::time::Instant::now(),
+            },
+            false,
+            cx,
+        )
+    });
+
+    let client_ssh = RemoteClient::connect_mock(opts, cx_a).await;
+    let (project_a, worktree_id) = client_a
+        .build_ssh_project(path!("/project"), client_ssh, false, cx_a)
+        .await;
+    log_store.update(cx_a, |log_store, cx| log_store.add_project(&project_a, cx));
+
+    let (buffer, _handle) = project_a
+        .update(cx_a, |project, cx| {
+            project.open_buffer_with_lsp((worktree_id, rel_path("a.rs")), cx)
+        })
+        .await
+        .unwrap();
+
+    let first_server = fake_language_servers.next().await.unwrap();
+    let first_server_id = first_server.server.server_id();
+    executor.run_until_parked();
+
+    project_a.read_with(cx_a, |project, cx| {
+        let statuses = project.language_server_statuses(cx).collect::<Vec<_>>();
+        assert_eq!(statuses.len(), 1);
+        assert_eq!(statuses[0].0, first_server_id);
+        assert_eq!(statuses[0].1.name.0, "the-language-server");
+    });
+    cx_a.read_global::<GlobalLogStore, _>(|global, cx| {
+        let log_store = global.0.read(cx);
+        let matching_server_ids = log_store
+            .language_servers
+            .iter()
+            .filter_map(|(server_id, state)| {
+                state
+                    .name
+                    .as_ref()
+                    .is_some_and(|name| name.0 == "the-language-server")
+                    .then_some(*server_id)
+            })
+            .collect::<Vec<_>>();
+        assert_eq!(matching_server_ids, vec![first_server_id]);
+    });
+
+    project_a.update(cx_a, |project, cx| {
+        project.restart_language_servers_for_buffers(vec![buffer], HashSet::default(), cx);
+    });
+
+    let restarted_server = fake_language_servers.next().await.unwrap();
+    let restarted_server_id = restarted_server.server.server_id();
+    assert_ne!(restarted_server_id, first_server_id);
+    executor.run_until_parked();
+
+    project_a.read_with(cx_a, |project, cx| {
+        let statuses = project.language_server_statuses(cx).collect::<Vec<_>>();
+        assert_eq!(
+            statuses.len(),
+            1,
+            "restarting a remote language server should replace the previous status entry"
+        );
+        assert_eq!(
+            statuses[0].0, restarted_server_id,
+            "restarting a remote language server should publish the replacement server id"
+        );
+        assert_ne!(
+            statuses[0].0, first_server_id,
+            "restarting a remote language server should remove the previous server id"
+        );
+        assert_eq!(statuses[0].1.name.0, "the-language-server");
+    });
+    cx_a.read_global::<GlobalLogStore, _>(|global, cx| {
+        let log_store = global.0.read(cx);
+        let matching_server_ids = log_store
+            .language_servers
+            .iter()
+            .filter_map(|(server_id, state)| {
+                state
+                    .name
+                    .as_ref()
+                    .is_some_and(|name| name.0 == "the-language-server")
+                    .then_some(*server_id)
+            })
+            .collect::<Vec<_>>();
+        assert_eq!(
+            matching_server_ids,
+            vec![restarted_server_id],
+            "restarting a remote language server should replace the old log store entry"
+        );
+        assert!(
+            !log_store.language_servers.contains_key(&first_server_id),
+            "restarting a remote language server should remove the previous log store entry"
+        );
+    });
+}
+
 #[gpui::test]
 async fn test_remote_server_debugger(
     cx_a: &mut TestAppContext,

crates/project/src/lsp_store.rs

@@ -9910,6 +9910,15 @@ impl LspStore {
                     lsp_store.disk_based_diagnostics_finished(language_server_id, cx)
                 }
 
+                proto::update_language_server::Variant::Removed(_) => {
+                    lsp_store
+                        .language_server_statuses
+                        .remove(&language_server_id);
+                    lsp_store.cleanup_lsp_data(language_server_id);
+                    cx.emit(LspStoreEvent::LanguageServerRemoved(language_server_id));
+                    cx.notify();
+                }
+
                 non_lsp @ proto::update_language_server::Variant::StatusUpdate(_)
                 | non_lsp @ proto::update_language_server::Variant::RegisteredForBuffer(_)
                 | non_lsp @ proto::update_language_server::Variant::MetadataUpdated(_) => {

crates/proto/proto/lsp.proto

@@ -608,6 +608,7 @@ message UpdateLanguageServer {
     StatusUpdate status_update = 9;
     RegisteredForBuffer registered_for_buffer = 10;
     ServerMetadataUpdated metadata_updated = 11;
+    ServerRemoved removed = 12;
   }
 }
 
@@ -644,6 +645,8 @@ message LspDiskBasedDiagnosticsUpdating {}
 
 message LspDiskBasedDiagnosticsUpdated {}
 
+message ServerRemoved {}
+
 message StatusUpdate {
   optional string message = 1;
   oneof status {

crates/remote_server/src/headless_project.rs

@@ -416,6 +416,16 @@ impl HeadlessProject {
                         log_store.remove_language_server(*id, cx);
                     });
                 }
+                self.session
+                    .send(proto::UpdateLanguageServer {
+                        project_id: REMOTE_SERVER_PROJECT_ID,
+                        server_name: None,
+                        language_server_id: id.to_proto(),
+                        variant: Some(proto::update_language_server::Variant::Removed(
+                            proto::ServerRemoved {},
+                        )),
+                    })
+                    .log_err();
             }
             LspStoreEvent::LanguageServerUpdate {
                 language_server_id,