Move cancel to Client

MichaReiser · MichaReiser · commit f5ece05df9c2 · 2025-05-23T14:07:20.000+02:00
diff --git a/crates/ty_server/src/lib.rs b/crates/ty_server/src/lib.rs
@@ -4,7 +4,6 @@ pub use document::{DocumentKey, NotebookDocument, PositionEncoding, TextDocument
 pub use session::{ClientSettings, DocumentQuery, DocumentSnapshot, Session};
 use std::num::NonZeroUsize;
 
-mod client;
 mod document;
 mod logging;
 mod server;
diff --git a/crates/ty_server/src/server.rs b/crates/ty_server/src/server.rs
@@ -19,7 +19,7 @@ mod connection;
 mod main_loop;
 mod schedule;
 
-use crate::client::Client;
+use crate::session::client::Client;
 pub(crate) use api::Error;
 pub(crate) use connection::{ConnectionInitializer, ConnectionSender};
 pub(crate) use main_loop::{Action, Event, MainLoopReceiver, MainLoopSender};
diff --git a/crates/ty_server/src/server/api.rs b/crates/ty_server/src/server/api.rs
@@ -15,7 +15,7 @@ mod traits;
 
 use self::traits::{NotificationHandler, RequestHandler};
 use super::{Result, schedule::BackgroundSchedule};
-use crate::client::Client;
+use crate::session::client::Client;
 use crate::show_err_msg;
 use notifications as notification;
 use requests as request;
diff --git a/crates/ty_server/src/server/api/diagnostics.rs b/crates/ty_server/src/server/api/diagnostics.rs
@@ -2,8 +2,8 @@ use lsp_server::ErrorCode;
 use lsp_types::{PublishDiagnosticsParams, Url, notification::PublishDiagnostics};
 
 use super::LSPResult;
-use crate::client::Client;
 use crate::server::Result;
+use crate::session::client::Client;
 
 pub(super) fn clear_diagnostics(uri: &Url, client: &Client) -> Result<()> {
     client
diff --git a/crates/ty_server/src/server/api/notifications/cancel.rs b/crates/ty_server/src/server/api/notifications/cancel.rs
@@ -1,11 +1,11 @@
-use lsp_server::{ErrorCode, RequestId, ResponseError};
+use lsp_server::RequestId;
 use lsp_types::CancelParams;
 use lsp_types::notification::Cancel;
 
-use crate::client::Client;
 use crate::server::Result;
 use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
 use crate::session::Session;
+use crate::session::client::Client;
 
 pub(crate) struct CancelNotificationHandler;
 
@@ -20,18 +20,7 @@ impl SyncNotificationHandler for CancelNotificationHandler {
             lsp_types::NumberOrString::String(id) => id.into(),
         };
 
-        let method_name = session.request_queue_mut().incoming_mut().cancel(&id);
-
-        if let Some(method_name) = method_name {
-            tracing::debug!("Cancelled request id={id} method={method_name}");
-            let error = ResponseError {
-                code: ErrorCode::RequestCanceled as i32,
-                message: "request was cancelled by client".to_owned(),
-                data: None,
-            };
-
-            let _ = client.respond_err(id, error);
-        }
+        let _ = client.cancel(session, id);
 
         Ok(())
     }
diff --git a/crates/ty_server/src/server/api/notifications/did_change.rs b/crates/ty_server/src/server/api/notifications/did_change.rs
@@ -2,11 +2,11 @@ use lsp_server::ErrorCode;
 use lsp_types::DidChangeTextDocumentParams;
 use lsp_types::notification::DidChangeTextDocument;
 
-use crate::client::Client;
 use crate::server::Result;
 use crate::server::api::LSPResult;
 use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
 use crate::session::Session;
+use crate::session::client::Client;
 use crate::system::{AnySystemPath, url_to_any_system_path};
 use ty_project::watch::ChangeEvent;
 
diff --git a/crates/ty_server/src/server/api/notifications/did_change_watched_files.rs b/crates/ty_server/src/server/api/notifications/did_change_watched_files.rs
@@ -1,8 +1,8 @@
-use crate::client::Client;
 use crate::server::Result;
 use crate::server::api::LSPResult;
 use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
 use crate::session::Session;
+use crate::session::client::Client;
 use crate::system::{AnySystemPath, url_to_any_system_path};
 use lsp_types as types;
 use lsp_types::{FileChangeType, notification as notif};
diff --git a/crates/ty_server/src/server/api/notifications/did_close.rs b/crates/ty_server/src/server/api/notifications/did_close.rs
@@ -1,9 +1,9 @@
-use crate::client::Client;
 use crate::server::Result;
 use crate::server::api::LSPResult;
 use crate::server::api::diagnostics::clear_diagnostics;
 use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
 use crate::session::Session;
+use crate::session::client::Client;
 use crate::system::{AnySystemPath, url_to_any_system_path};
 use lsp_server::ErrorCode;
 use lsp_types::DidCloseTextDocumentParams;
diff --git a/crates/ty_server/src/server/api/notifications/did_close_notebook.rs b/crates/ty_server/src/server/api/notifications/did_close_notebook.rs
@@ -1,11 +1,11 @@
 use lsp_types::DidCloseNotebookDocumentParams;
 use lsp_types::notification::DidCloseNotebookDocument;
 
-use crate::client::Client;
 use crate::server::Result;
 use crate::server::api::LSPResult;
 use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
 use crate::session::Session;
+use crate::session::client::Client;
 use crate::system::{AnySystemPath, url_to_any_system_path};
 use ty_project::watch::ChangeEvent;
 
diff --git a/crates/ty_server/src/server/api/notifications/did_open.rs b/crates/ty_server/src/server/api/notifications/did_open.rs
@@ -2,10 +2,10 @@ use lsp_types::notification::DidOpenTextDocument;
 use lsp_types::{DidOpenTextDocumentParams, TextDocumentItem};
 
 use crate::TextDocument;
-use crate::client::Client;
 use crate::server::Result;
 use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
 use crate::session::Session;
+use crate::session::client::Client;
 use crate::system::{AnySystemPath, url_to_any_system_path};
 use ruff_db::Db;
 use ty_project::watch::ChangeEvent;
diff --git a/crates/ty_server/src/server/api/notifications/did_open_notebook.rs b/crates/ty_server/src/server/api/notifications/did_open_notebook.rs
@@ -5,12 +5,12 @@ use lsp_types::notification::DidOpenNotebookDocument;
 use ruff_db::Db;
 use ty_project::watch::ChangeEvent;
 
-use crate::client::Client;
 use crate::document::NotebookDocument;
 use crate::server::Result;
 use crate::server::api::LSPResult;
 use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
 use crate::session::Session;
+use crate::session::client::Client;
 use crate::system::{AnySystemPath, url_to_any_system_path};
 
 pub(crate) struct DidOpenNotebookHandler;
diff --git a/crates/ty_server/src/server/api/requests/completion.rs b/crates/ty_server/src/server/api/requests/completion.rs
@@ -7,9 +7,9 @@ use ty_ide::completion;
 use ty_project::ProjectDatabase;
 
 use crate::DocumentSnapshot;
-use crate::client::Client;
 use crate::document::PositionExt;
 use crate::server::api::traits::{BackgroundDocumentRequestHandler, RequestHandler};
+use crate::session::client::Client;
 
 pub(crate) struct CompletionRequestHandler;
 
diff --git a/crates/ty_server/src/server/api/requests/diagnostic.rs b/crates/ty_server/src/server/api/requests/diagnostic.rs
@@ -8,11 +8,11 @@ use lsp_types::{
 };
 
 use crate::PositionEncoding;
-use crate::client::Client;
 use crate::document::{FileRangeExt, ToRangeExt};
 use crate::server::Result;
 use crate::server::api::traits::{BackgroundDocumentRequestHandler, RequestHandler};
 use crate::session::DocumentSnapshot;
+use crate::session::client::Client;
 use ruff_db::diagnostic::{Annotation, Severity, SubDiagnostic};
 use ruff_db::files::FileRange;
 use ruff_db::source::{line_index, source_text};
diff --git a/crates/ty_server/src/server/api/requests/goto_type_definition.rs b/crates/ty_server/src/server/api/requests/goto_type_definition.rs
@@ -7,9 +7,9 @@ use ty_ide::goto_type_definition;
 use ty_project::ProjectDatabase;
 
 use crate::DocumentSnapshot;
-use crate::client::Client;
 use crate::document::{PositionExt, ToLink};
 use crate::server::api::traits::{BackgroundDocumentRequestHandler, RequestHandler};
+use crate::session::client::Client;
 
 pub(crate) struct GotoTypeDefinitionRequestHandler;
 
diff --git a/crates/ty_server/src/server/api/requests/hover.rs b/crates/ty_server/src/server/api/requests/hover.rs
@@ -1,9 +1,9 @@
 use std::borrow::Cow;
 
 use crate::DocumentSnapshot;
-use crate::client::Client;
 use crate::document::{PositionExt, ToRangeExt};
 use crate::server::api::traits::{BackgroundDocumentRequestHandler, RequestHandler};
+use crate::session::client::Client;
 use lsp_types::request::HoverRequest;
 use lsp_types::{HoverContents, HoverParams, MarkupContent, Url};
 use ruff_db::source::{line_index, source_text};
diff --git a/crates/ty_server/src/server/api/requests/inlay_hints.rs b/crates/ty_server/src/server/api/requests/inlay_hints.rs
@@ -1,9 +1,9 @@
 use std::borrow::Cow;
 
 use crate::DocumentSnapshot;
-use crate::client::Client;
 use crate::document::{RangeExt, TextSizeExt};
 use crate::server::api::traits::{BackgroundDocumentRequestHandler, RequestHandler};
+use crate::session::client::Client;
 use lsp_types::request::InlayHintRequest;
 use lsp_types::{InlayHintParams, Url};
 use ruff_db::source::{line_index, source_text};
diff --git a/crates/ty_server/src/server/api/traits.rs b/crates/ty_server/src/server/api/traits.rs
@@ -1,6 +1,6 @@
 //! A stateful LSP implementation that calls into the ty API.
 
-use crate::client::Client;
+use crate::session::client::Client;
 use crate::session::{DocumentSnapshot, Session};
 
 use lsp_types::notification::Notification as LSPNotification;
diff --git a/crates/ty_server/src/server/main_loop.rs b/crates/ty_server/src/server/main_loop.rs
@@ -1,7 +1,7 @@
 use crate::Session;
-use crate::client::Client;
 use crate::server::schedule::Scheduler;
 use crate::server::{Server, api};
+use crate::session::client::Client;
 use crossbeam::select;
 use lsp_server::Message;
 use lsp_types::{DidChangeWatchedFilesRegistrationOptions, FileSystemWatcher};
@@ -65,14 +65,8 @@ impl Server {
                 }
                 Event::Action(action) => match action {
                     Action::SendResponse(response) => {
-                        // Always send cancellation responses to the client.
-                        if response.error.as_ref().is_some_and(|err| {
-                            err.code == lsp_server::ErrorCode::RequestCanceled as i32
-                        }) {
-                            self.connection.send(Message::Response(response))?;
-                        }
                         // Filter out responses for already canceled requests.
-                        else if let Some((start_time, method)) = self
+                        if let Some((start_time, method)) = self
                             .session
                             .request_queue_mut()
                             .incoming_mut()
diff --git a/crates/ty_server/src/server/schedule.rs b/crates/ty_server/src/server/schedule.rs
@@ -9,7 +9,7 @@ use self::{
     task::{BackgroundTaskBuilder, SyncTask},
     thread::ThreadPriority,
 };
-use crate::client::Client;
+use crate::session::client::Client;
 pub(super) use task::{BackgroundSchedule, Task};
 
 /// The event loop thread is actually a secondary thread that we spawn from the
diff --git a/crates/ty_server/src/server/schedule/task.rs b/crates/ty_server/src/server/schedule/task.rs
@@ -1,8 +1,8 @@
 use lsp_server::RequestId;
 use serde::Serialize;
 
-use crate::client::Client;
 use crate::session::Session;
+use crate::session::client::Client;
 
 type LocalFn = Box<dyn FnOnce(&mut Session, &Client)>;
 
diff --git a/crates/ty_server/src/session.rs b/crates/ty_server/src/session.rs
@@ -23,6 +23,7 @@ use crate::system::{AnySystemPath, LSPSystem, url_to_any_system_path};
 use crate::{PositionEncoding, TextDocument};
 
 mod capabilities;
+pub(crate) mod client;
 pub(crate) mod index;
 mod request_queue;
 mod settings;
diff --git a/crates/ty_server/src/session/client.rs b/crates/ty_server/src/session/client.rs
@@ -2,7 +2,7 @@ use crate::Session;
 use crate::server::{Action, ConnectionSender};
 use crate::server::{Event, MainLoopSender};
 use anyhow::{Context, anyhow};
-use lsp_server::{Message, Notification, RequestId};
+use lsp_server::{ErrorCode, Message, Notification, RequestId, ResponseError};
 use serde_json::Value;
 use std::any::TypeId;
 use std::fmt::Display;
@@ -191,6 +191,30 @@ impl Client {
             .send(Event::Action(Action::RetryRequest(request)))
             .map_err(|error| anyhow!("Failed to send retry request: {error}"))
     }
+
+    pub(crate) fn cancel(&self, session: &mut Session, id: RequestId) -> crate::Result<()> {
+        let method_name = session.request_queue_mut().incoming_mut().cancel(&id);
+
+        if let Some(method_name) = method_name {
+            tracing::debug!("Cancelled request id={id} method={method_name}");
+            let error = ResponseError {
+                code: ErrorCode::RequestCanceled as i32,
+                message: "request was cancelled by client".to_owned(),
+                data: None,
+            };
+
+            // Use `client_sender` here instead of `respond_err` because
+            // `respond_err` filters out responses for canceled requests (which we just did!).
+            self.client_sender
+                .send(Message::Response(lsp_server::Response {
+                    id,
+                    result: None,
+                    error: Some(error),
+                }))?;
+        }
+
+        Ok(())
+    }
 }
 
 /// Sends an error to the client with a formatted message. The error is sent in a
diff --git a/crates/ty_server/src/session/request_queue.rs b/crates/ty_server/src/session/request_queue.rs
@@ -1,4 +1,4 @@
-use crate::client::ClientResponseHandler;
+use crate::session::client::ClientResponseHandler;
 use lsp_server::RequestId;
 use rustc_hash::FxHashMap;
 use std::cell::{Cell, OnceCell, RefCell};
@@ -66,7 +66,7 @@ impl Incoming {
     /// Cancels the pending request with the given id.
     ///
     /// Returns the method name if the request was still pending, `None` if it was already completed.
-    pub(crate) fn cancel(&mut self, request_id: &RequestId) -> Option<String> {
+    pub(super) fn cancel(&mut self, request_id: &RequestId) -> Option<String> {
         self.pending.remove(request_id).map(|mut pending| {
             if let Some(cancellation_token) = pending.cancellation_token.take() {
                 cancellation_token.cancel();
