diff --git a/pyrefly/lib/commands/tsp.rs b/pyrefly/lib/commands/tsp.rs index 1ef0f9a74..5330bdd54 100644 --- a/pyrefly/lib/commands/tsp.rs +++ b/pyrefly/lib/commands/tsp.rs @@ -44,14 +44,14 @@ pub fn run_tsp(connection: Arc, 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(()) } diff --git a/pyrefly/lib/lsp/server.rs b/pyrefly/lib/lsp/server.rs index 016ee4216..08aebe987 100644 --- a/pyrefly/lib/lsp/server.rs +++ b/pyrefly/lib/lsp/server.rs @@ -227,6 +227,9 @@ pub trait TspInterface { /// Get access to the state for creating transactions fn state(&self) -> &Arc; + /// 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, @@ -2087,6 +2090,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>, diff --git a/pyrefly/lib/test/tsp/tsp_interaction/get_snapshot.rs b/pyrefly/lib/test/tsp/tsp_interaction/get_snapshot.rs new file mode 100644 index 000000000..39adcee90 --- /dev/null +++ b/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(); +} diff --git a/pyrefly/lib/test/tsp/tsp_interaction/mod.rs b/pyrefly/lib/test/tsp/tsp_interaction/mod.rs index fbeb421e2..f168b22dd 100644 --- a/pyrefly/lib/test/tsp/tsp_interaction/mod.rs +++ b/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; diff --git a/pyrefly/lib/test/tsp/tsp_interaction/object_model.rs b/pyrefly/lib/test/tsp/tsp_interaction/object_model.rs index 4d9e038be..3a960b089 100644 --- a/pyrefly/lib/test/tsp/tsp_interaction/object_model.rs +++ b/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::(0); let args = TspArgs { - indexing_mode: IndexingMode::None, + indexing_mode: IndexingMode::LazyBlocking, workspace_indexing_limit: 0, }; let connection = Connection { diff --git a/pyrefly/lib/tsp/requests/get_snapshot.rs b/pyrefly/lib/tsp/requests/get_snapshot.rs new file mode 100644 index 000000000..b5098c377 --- /dev/null +++ b/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() + }) + } +} diff --git a/pyrefly/lib/tsp/requests/mod.rs b/pyrefly/lib/tsp/requests/mod.rs index be221560b..32f8e9395 100644 --- a/pyrefly/lib/tsp/requests/mod.rs +++ b/pyrefly/lib/tsp/requests/mod.rs @@ -7,4 +7,5 @@ //! TSP request implementations +pub mod get_snapshot; pub mod get_supported_protocol_version; diff --git a/pyrefly/lib/tsp/server.rs b/pyrefly/lib/tsp/server.rs index b74062d60..291bd15df 100644 --- a/pyrefly/lib/tsp/server.rs +++ b/pyrefly/lib/tsp/server.rs @@ -7,6 +7,7 @@ use std::collections::HashSet; use std::sync::Arc; +use std::sync::Mutex; use dupe::Dupe; use lsp_server::Connection; @@ -29,11 +30,16 @@ use crate::lsp::transaction_manager::TransactionManager; /// TSP server that delegates to LSP server infrastructure while handling only TSP requests pub struct TspServer { pub inner: Box, + /// Current snapshot version, updated on RecheckFinished events + pub(crate) current_snapshot: Arc>, } impl TspServer { pub fn new(lsp_server: Box) -> Self { - Self { inner: lsp_server } + Self { + inner: lsp_server, + current_snapshot: Arc::new(Mutex::new(0)), // Start at 0, increments on RecheckFinished + } } pub fn process_event<'a>( @@ -43,6 +49,18 @@ impl TspServer { subsequent_mutation: bool, event: LspEvent, ) -> anyhow::Result { + // Remember if this event should increment the snapshot after processing + let should_increment_snapshot = match &event { + LspEvent::RecheckFinished => true, + // Increment on DidChange since it affects type checker state via synchronous validation + LspEvent::DidChangeTextDocument(_) => true, + // Don't increment on DidChangeWatchedFiles directly since it triggers RecheckFinished + // LspEvent::DidChangeWatchedFiles(_) => true, + // Don't increment on DidOpen since it triggers RecheckFinished events that will increment + // LspEvent::DidOpenTextDocument(_) => true, + _ => false, + }; + // For TSP requests, handle them specially if let LspEvent::LspRequest(ref request) = event { if self.handle_tsp_request(ide_transaction_manager, request)? { @@ -58,12 +76,19 @@ impl TspServer { } // For all other events (notifications, responses, etc.), delegate to inner server - self.inner.process_event( + let result = self.inner.process_event( ide_transaction_manager, canceled_requests, subsequent_mutation, event, - ) + )?; + + // Increment snapshot after the inner server has processed the event + if should_increment_snapshot && let Ok(mut current) = self.current_snapshot.lock() { + *current += 1; + } + + Ok(result) } fn handle_tsp_request<'a>( @@ -94,6 +119,12 @@ impl TspServer { ide_transaction_manager.save(transaction); Ok(true) } + TSPRequests::GetSnapshotRequest { .. } => { + // Get snapshot doesn't need a transaction since it just returns the cached value + self.inner + .send_response(new_response(request.id.clone(), Ok(self.get_snapshot()))); + Ok(true) + } _ => { // Other TSP requests not yet implemented Ok(false) @@ -106,13 +137,19 @@ pub fn tsp_loop( lsp_server: Box, connection: Arc, _initialization_params: InitializeParams, + lsp_queue: LspQueue, ) -> anyhow::Result<()> { eprintln!("Reading TSP messages"); let connection_for_dispatcher = connection.dupe(); - let lsp_queue = LspQueue::new(); let server = TspServer::new(lsp_server); + // Start the recheck queue thread to process async tasks + let recheck_queue = server.inner.recheck_queue().dupe(); + std::thread::spawn(move || { + recheck_queue.run_until_stopped(); + }); + let lsp_queue2 = lsp_queue.dupe(); std::thread::spawn(move || { dispatch_lsp_events(&connection_for_dispatcher, lsp_queue2);