Add support for TSP snapshot request

pyrefly/lib/commands/tsp.rs

@@ -44,14 +44,14 @@ pub fn run_tsp(connection: Arc<Connection>, args: TspArgs) -> anyhow::Result<()>
     let lsp_queue = LspQueue::new();
     let lsp_server = Box::new(crate::lsp::server::Server::new(
         connection.dupe(),
-        lsp_queue,
+        lsp_queue.dupe(),
         initialization_params.clone(),
         args.indexing_mode,
         args.workspace_indexing_limit,
     ));
 
     // Reuse the existing lsp_loop but with TSP initialization
-    tsp_loop(lsp_server, connection, initialization_params)?;
+    tsp_loop(lsp_server, connection, initialization_params, lsp_queue)?;
     Ok(())
 }
 

pyrefly/lib/lsp/server.rs

@@ -224,6 +224,9 @@ pub trait TspInterface {
     /// Get access to the state for creating transactions
     fn state(&self) -> &Arc<State>;
 
+    /// Get access to the recheck queue for async task processing
+    fn recheck_queue(&self) -> &HeavyTaskQueue;
+
     /// Process an LSP event and return the next step
     fn process_event<'a>(
         &'a self,
@@ -2066,6 +2069,10 @@ impl TspInterface for Server {
         &self.state
     }
 
+    fn recheck_queue(&self) -> &HeavyTaskQueue {
+        &self.recheck_queue
+    }
+
     fn process_event<'a>(
         &'a self,
         ide_transaction_manager: &mut TransactionManager<'a>,

pyrefly/lib/test/lsp/lsp_interaction/configuration.rs

@@ -12,13 +12,11 @@ use std::os::unix::fs::PermissionsExt;
 
 use lsp_server::Message;
 use lsp_server::Notification;
-use lsp_server::Request;
 use lsp_server::RequestId;
 use lsp_server::Response;
 use lsp_types::Url;
 use lsp_types::notification::DidChangeWorkspaceFolders;
 use lsp_types::notification::Notification as _;
-use pyrefly_util::fs_anyhow::write;
 
 use crate::test::lsp::lsp_interaction::object_model::InitializeSettings;
 use crate::test::lsp::lsp_interaction::object_model::LspInteraction;

pyrefly/lib/test/tsp/tsp_interaction/get_snapshot.rs

@@ -0,0 +1,198 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+//! Tests for TSP getSnapshot request
+
+use lsp_server::RequestId;
+use lsp_server::Response;
+use tempfile::TempDir;
+
+use crate::test::tsp::tsp_interaction::object_model::TspInteraction;
+
+#[test]
+fn test_tsp_get_snapshot() {
+    // Test retrieval of TSP snapshot version
+    let temp_dir = TempDir::new().unwrap();
+    let test_file_path = temp_dir.path().join("test.py");
+
+    let test_content = r#"# Simple test file
+print("Hello, World!")
+"#;
+
+    std::fs::write(&test_file_path, test_content).unwrap();
+
+    // Create a pyproject.toml to make this a recognized Python project
+    let pyproject_content = r#"[build-system]
+requires = ["setuptools>=45", "setuptools-scm[toml]>=6.2"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "test-project"
+version = "1.0.0"
+"#;
+    std::fs::write(temp_dir.path().join("pyproject.toml"), pyproject_content).unwrap();
+
+    let mut tsp = TspInteraction::new();
+    tsp.set_root(temp_dir.path().to_path_buf());
+    tsp.initialize(Default::default());
+
+    // Open the test file
+    tsp.server.did_open("test.py");
+
+    // Wait for any diagnostics/RecheckFinished events
+    tsp.client.expect_any_message();
+
+    // Get snapshot
+    tsp.server.get_snapshot();
+
+    // Expect snapshot response with integer (should increment after RecheckFinished from indexing)
+    tsp.client.expect_response(Response {
+        id: RequestId::from(2),
+        result: Some(serde_json::json!(1)),
+        error: None,
+    });
+
+    tsp.shutdown();
+}
+
+#[test]
+fn test_tsp_snapshot_updates_on_file_change() {
+    // Test that DidChangeWatchedFiles events trigger async recheck and update snapshots
+    // With the recheck queue thread running, async tasks should execute and generate RecheckFinished events
+    let temp_dir = TempDir::new().unwrap();
+    let test_file_path = temp_dir.path().join("changing_test.py");
+
+    let initial_content = r#"# Initial content
+x = 1
+"#;
+
+    std::fs::write(&test_file_path, initial_content).unwrap();
+
+    // Create a pyproject.toml to make this a recognized Python project
+    let pyproject_content = r#"[build-system]
+requires = ["setuptools>=45", "setuptools-scm[toml]>=6.2"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "test-project"
+version = "1.0.0"
+"#;
+    std::fs::write(temp_dir.path().join("pyproject.toml"), pyproject_content).unwrap();
+
+    let mut tsp = TspInteraction::new();
+    tsp.set_root(temp_dir.path().to_path_buf());
+    tsp.initialize(Default::default());
+
+    // Open the test file
+    tsp.server.did_open("changing_test.py");
+
+    // Wait for any diagnostics/RecheckFinished events
+    tsp.client.expect_any_message();
+
+    // Get initial snapshot
+    tsp.server.get_snapshot();
+
+    // Expect first snapshot response
+    tsp.client.expect_response(Response {
+        id: RequestId::from(2),
+        result: Some(serde_json::json!(1)),
+        error: None,
+    });
+
+    // Modify the file to trigger a state change
+    let updated_content = r#"# Updated content
+x = 2
+y = "hello"
+"#;
+
+    std::fs::write(&test_file_path, updated_content).unwrap();
+
+    // Simulate the LSP DidChangeWatchedFiles notification for the file change
+    tsp.server
+        .did_change_watched_files("changing_test.py", "changed");
+
+    // Wait for the async RecheckFinished event to be processed
+    tsp.client.expect_any_message();
+
+    // Get snapshot after async recheck completes
+    tsp.server.get_snapshot();
+
+    // Expect snapshot to be incremented to 2 after RecheckFinished from file change
+    tsp.client.expect_response(Response {
+        id: RequestId::from(3),
+        result: Some(serde_json::json!(2)), // Should be 2 after RecheckFinished from file change
+        error: None,
+    });
+
+    tsp.shutdown();
+}
+
+#[test]
+fn test_tsp_snapshot_updates_on_did_change() {
+    // Test that didChange events cause snapshot to update
+    let temp_dir = TempDir::new().unwrap();
+    let test_file_path = temp_dir.path().join("change_test.py");
+
+    let initial_content = r#"# Initial content
+x = 1
+"#;
+
+    std::fs::write(&test_file_path, initial_content).unwrap();
+
+    // Create a pyproject.toml to make this a recognized Python project
+    let pyproject_content = r#"[build-system]
+requires = ["setuptools>=45", "setuptools-scm[toml]>=6.2"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "test-project"
+version = "1.0.0"
+"#;
+    std::fs::write(temp_dir.path().join("pyproject.toml"), pyproject_content).unwrap();
+
+    let mut tsp = TspInteraction::new();
+    tsp.set_root(temp_dir.path().to_path_buf());
+    tsp.initialize(Default::default());
+
+    // Open the test file
+    tsp.server.did_open("change_test.py");
+
+    // Wait for any diagnostics/RecheckFinished events from opening
+    tsp.client.expect_any_message();
+
+    // Get initial snapshot
+    tsp.server.get_snapshot();
+
+    // Expect first snapshot response
+    tsp.client.expect_response(Response {
+        id: RequestId::from(2),
+        result: Some(serde_json::json!(1)),
+        error: None,
+    });
+
+    // Send a didChange notification with updated content
+    let changed_content = r#"# Changed content
+x = 2
+y = 'updated'
+"#;
+    tsp.server.did_change("change_test.py", changed_content, 2);
+
+    // Wait for any RecheckFinished events triggered by the change
+    tsp.client.expect_any_message();
+
+    // Get updated snapshot
+    tsp.server.get_snapshot();
+
+    // Expect second snapshot response - should be incremented due to didChange
+    tsp.client.expect_response(Response {
+        id: RequestId::from(3),
+        result: Some(serde_json::json!(2)),
+        error: None,
+    });
+
+    tsp.shutdown();
+}

pyrefly/lib/test/tsp/tsp_interaction/mod.rs

@@ -7,5 +7,6 @@
 
 //! Tests for TSP (Type Server Protocol) request handlers
 
+pub mod get_snapshot;
 pub mod get_supported_protocol_version;
 pub mod object_model;

pyrefly/lib/test/tsp/tsp_interaction/object_model.rs

@@ -119,6 +119,15 @@ impl TestTspServer {
         }));
     }
 
+    pub fn get_snapshot(&mut self) {
+        let id = self.next_request_id();
+        self.send_message(Message::Request(Request {
+            id,
+            method: "typeServer/getSnapshot".to_owned(),
+            params: serde_json::json!(null),
+        }));
+    }
+
     pub fn did_open(&self, file: &'static str) {
         let path = self.get_root_or_panic().join(file);
         self.send_message(Message::Notification(Notification {
@@ -134,6 +143,41 @@ impl TestTspServer {
         }));
     }
 
+    pub fn did_change(&self, file: &'static str, content: &str, version: i32) {
+        let path = self.get_root_or_panic().join(file);
+        self.send_message(Message::Notification(Notification {
+            method: "textDocument/didChange".to_owned(),
+            params: serde_json::json!({
+                "textDocument": {
+                    "uri": Url::from_file_path(&path).unwrap().to_string(),
+                    "version": version
+                },
+                "contentChanges": [{
+                    "text": content
+                }]
+            }),
+        }));
+    }
+
+    pub fn did_change_watched_files(&self, file: &'static str, change_type: &str) {
+        let path = self.get_root_or_panic().join(file);
+        let file_change_type = match change_type {
+            "created" => 1, // FileChangeType::CREATED
+            "changed" => 2, // FileChangeType::CHANGED
+            "deleted" => 3, // FileChangeType::DELETED
+            _ => 2,         // Default to changed
+        };
+        self.send_message(Message::Notification(Notification {
+            method: "workspace/didChangeWatchedFiles".to_owned(),
+            params: serde_json::json!({
+                "changes": [{
+                    "uri": Url::from_file_path(&path).unwrap().to_string(),
+                    "type": file_change_type
+                }]
+            }),
+        }));
+    }
+
     pub fn get_initialize_params(&self, settings: &InitializeSettings) -> Value {
         let mut params: Value = serde_json::json!({
             "rootPath": "/",
@@ -251,6 +295,21 @@ impl TestTspClient {
             }
         }
     }
+
+    pub fn receive_any_message(&self) -> Message {
+        match self.receiver.recv_timeout(self.timeout) {
+            Ok(msg) => {
+                eprintln!("client<---server {}", serde_json::to_string(&msg).unwrap());
+                msg
+            }
+            Err(RecvTimeoutError::Timeout) => {
+                panic!("Timeout waiting for response");
+            }
+            Err(RecvTimeoutError::Disconnected) => {
+                panic!("Channel disconnected");
+            }
+        }
+    }
 }
 
 pub struct TspInteraction {
@@ -266,7 +325,7 @@ impl TspInteraction {
         let (language_server_sender, language_server_receiver) = bounded::<Message>(0);
 
         let args = TspArgs {
-            indexing_mode: IndexingMode::None,
+            indexing_mode: IndexingMode::LazyBlocking,
             workspace_indexing_limit: 0,
         };
         let connection = Connection {

pyrefly/lib/tsp/requests/get_snapshot.rs

@@ -0,0 +1,25 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+//! Implementation of the getSnapshot TSP request
+
+use crate::tsp::server::TspServer;
+
+impl TspServer {
+    /// Get the current snapshot version
+    ///
+    /// The snapshot represents the current epoch of the global state.
+    /// It changes whenever files are modified, configuration changes,
+    /// or any other event that would trigger a recomputation.
+    pub fn get_snapshot(&self) -> i32 {
+        *self.current_snapshot.lock().unwrap_or_else(|poisoned| {
+            // In case of poisoned mutex, recover and return the value
+            eprintln!("TSP: Warning - snapshot mutex was poisoned, recovering");
+            poisoned.into_inner()
+        })
+    }
+}

pyrefly/lib/tsp/requests/mod.rs

@@ -7,4 +7,5 @@
 
 //! TSP request implementations
 
+pub mod get_snapshot;
 pub mod get_supported_protocol_version;