[ty] Add panic-by-default await methods to `TestServer` (#21451)

Cargo.lock

@@ -4521,6 +4521,7 @@ dependencies = [
  "serde",
  "serde_json",
  "shellexpand",
+ "smallvec",
  "tempfile",
  "thiserror 2.0.17",
  "tracing",

crates/ty_server/Cargo.toml

@@ -34,6 +34,7 @@ salsa = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
 shellexpand = { workspace = true }
+smallvec = { workspace=true }
 thiserror = { workspace = true }
 tracing = { workspace = true }
 tracing-subscriber = { workspace = true, features = ["chrono"] }

crates/ty_server/tests/e2e/commands.rs

@@ -9,7 +9,7 @@ fn execute_command(
     server: &mut TestServer,
     command: String,
     arguments: Vec<serde_json::Value>,
-) -> anyhow::Result<Option<serde_json::Value>> {
+) -> Option<serde_json::Value> {
     let params = ExecuteCommandParams {
         command,
         arguments,
@@ -32,10 +32,10 @@ return 42
         .with_workspace(workspace_root, None)?
         .with_file(foo, foo_content)?
         .enable_pull_diagnostics(false)
-        .build()?
+        .build()
-        .wait_until_workspaces_are_initialized()?;
+        .wait_until_workspaces_are_initialized();
 
-    let response = execute_command(&mut server, "ty.printDebugInformation".to_string(), vec![])?;
+    let response = execute_command(&mut server, "ty.printDebugInformation".to_string(), vec![]);
     let response = response.expect("expect server response");
 
     let response = response

crates/ty_server/tests/e2e/initialize.rs

@@ -9,8 +9,8 @@ use crate::TestServerBuilder;
 #[test]
 fn empty_workspace_folders() -> Result<()> {
     let server = TestServerBuilder::new()?
-        .build()?
+        .build()
-        .wait_until_workspaces_are_initialized()?;
+        .wait_until_workspaces_are_initialized();
 
     let initialization_result = server.initialization_result().unwrap();
 
@@ -24,8 +24,8 @@ fn single_workspace_folder() -> Result<()> {
     let workspace_root = SystemPath::new("foo");
     let server = TestServerBuilder::new()?
         .with_workspace(workspace_root, None)?
-        .build()?
+        .build()
-        .wait_until_workspaces_are_initialized()?;
+        .wait_until_workspaces_are_initialized();
 
     let initialization_result = server.initialization_result().unwrap();
 
@@ -47,12 +47,12 @@ fn workspace_diagnostic_registration_without_configuration() -> Result<()> {
         .with_workspace(workspace_root, None)?
         .enable_workspace_configuration(false)
         .enable_diagnostic_dynamic_registration(true)
-        .build()?;
+        .build();
 
     // No need to wait for workspaces to initialize as the client does not support workspace
     // configuration.
 
-    let (_, params) = server.await_request::<RegisterCapability>()?;
+    let (_, params) = server.await_request::<RegisterCapability>();
     let [registration] = params.registrations.as_slice() else {
         panic!(
             "Expected a single registration, got: {:#?}",
@@ -90,12 +90,12 @@ fn open_files_diagnostic_registration_without_configuration() -> Result<()> {
         .with_workspace(workspace_root, None)?
         .enable_workspace_configuration(false)
         .enable_diagnostic_dynamic_registration(true)
-        .build()?;
+        .build();
 
     // No need to wait for workspaces to initialize as the client does not support workspace
     // configuration.
 
-    let (_, params) = server.await_request::<RegisterCapability>()?;
+    let (_, params) = server.await_request::<RegisterCapability>();
     let [registration] = params.registrations.as_slice() else {
         panic!(
             "Expected a single registration, got: {:#?}",
@@ -131,10 +131,10 @@ fn workspace_diagnostic_registration_via_initialization() -> Result<()> {
         )
         .with_workspace(workspace_root, None)?
         .enable_diagnostic_dynamic_registration(true)
-        .build()?
+        .build()
-        .wait_until_workspaces_are_initialized()?;
+        .wait_until_workspaces_are_initialized();
 
-    let (_, params) = server.await_request::<RegisterCapability>()?;
+    let (_, params) = server.await_request::<RegisterCapability>();
     let [registration] = params.registrations.as_slice() else {
         panic!(
             "Expected a single registration, got: {:#?}",
@@ -170,10 +170,10 @@ fn open_files_diagnostic_registration_via_initialization() -> Result<()> {
         )
         .with_workspace(workspace_root, None)?
         .enable_diagnostic_dynamic_registration(true)
-        .build()?
+        .build()
-        .wait_until_workspaces_are_initialized()?;
+        .wait_until_workspaces_are_initialized();
 
-    let (_, params) = server.await_request::<RegisterCapability>()?;
+    let (_, params) = server.await_request::<RegisterCapability>();
     let [registration] = params.registrations.as_slice() else {
         panic!(
             "Expected a single registration, got: {:#?}",
@@ -209,10 +209,10 @@ fn workspace_diagnostic_registration() -> Result<()> {
             Some(ClientOptions::default().with_diagnostic_mode(DiagnosticMode::Workspace)),
         )?
         .enable_diagnostic_dynamic_registration(true)
-        .build()?
+        .build()
-        .wait_until_workspaces_are_initialized()?;
+        .wait_until_workspaces_are_initialized();
 
-    let (_, params) = server.await_request::<RegisterCapability>()?;
+    let (_, params) = server.await_request::<RegisterCapability>();
     let [registration] = params.registrations.as_slice() else {
         panic!(
             "Expected a single registration, got: {:#?}",
@@ -248,10 +248,10 @@ fn open_files_diagnostic_registration() -> Result<()> {
             Some(ClientOptions::default().with_diagnostic_mode(DiagnosticMode::OpenFilesOnly)),
         )?
         .enable_diagnostic_dynamic_registration(true)
-        .build()?
+        .build()
-        .wait_until_workspaces_are_initialized()?;
+        .wait_until_workspaces_are_initialized();
 
-    let (_, params) = server.await_request::<RegisterCapability>()?;
+    let (_, params) = server.await_request::<RegisterCapability>();
     let [registration] = params.registrations.as_slice() else {
         panic!(
             "Expected a single registration, got: {:#?}",
@@ -291,11 +291,11 @@ def foo() -> str:
         .with_workspace(workspace_root, None)?
         .enable_pull_diagnostics(true)
         .with_file(foo, foo_content)?
-        .build()?
+        .build()
-        .wait_until_workspaces_are_initialized()?;
+        .wait_until_workspaces_are_initialized();
 
     server.open_text_document(foo, &foo_content, 1);
-    let hover = server.hover_request(foo, Position::new(0, 5))?;
+    let hover = server.hover_request(foo, Position::new(0, 5));
 
     assert!(
         hover.is_none(),
@@ -323,11 +323,11 @@ def foo() -> str:
         )?
         .enable_pull_diagnostics(true)
         .with_file(foo, foo_content)?
-        .build()?
+        .build()
-        .wait_until_workspaces_are_initialized()?;
+        .wait_until_workspaces_are_initialized();
 
     server.open_text_document(foo, &foo_content, 1);
-    let hover = server.hover_request(foo, Position::new(0, 5))?;
+    let hover = server.hover_request(foo, Position::new(0, 5));
 
     assert!(
         hover.is_none(),
@@ -364,18 +364,18 @@ def bar() -> str:
         .enable_pull_diagnostics(true)
         .with_file(foo, foo_content)?
         .with_file(bar, bar_content)?
-        .build()?
+        .build()
-        .wait_until_workspaces_are_initialized()?;
+        .wait_until_workspaces_are_initialized();
 
     server.open_text_document(foo, &foo_content, 1);
-    let hover_foo = server.hover_request(foo, Position::new(0, 5))?;
+    let hover_foo = server.hover_request(foo, Position::new(0, 5));
     assert!(
         hover_foo.is_none(),
         "Expected no hover information for workspace A, got: {hover_foo:?}"
     );
 
     server.open_text_document(bar, &bar_content, 1);
-    let hover_bar = server.hover_request(bar, Position::new(0, 5))?;
+    let hover_bar = server.hover_request(bar, Position::new(0, 5));
     assert!(
         hover_bar.is_some(),
         "Expected hover information for workspace B, got: {hover_bar:?}"
@@ -394,10 +394,10 @@ fn unknown_initialization_options() -> Result<()> {
         .with_initialization_options(
             ClientOptions::default().with_unknown([("bar".to_string(), Value::Null)].into()),
         )
-        .build()?
+        .build()
-        .wait_until_workspaces_are_initialized()?;
+        .wait_until_workspaces_are_initialized();
 
-    let show_message_params = server.await_notification::<ShowMessage>()?;
+    let show_message_params = server.await_notification::<ShowMessage>();
 
     insta::assert_json_snapshot!(show_message_params, @r#"
     {
@@ -419,10 +419,10 @@ fn unknown_options_in_workspace_configuration() -> Result<()> {
             workspace_root,
             Some(ClientOptions::default().with_unknown([("bar".to_string(), Value::Null)].into())),
         )?
-        .build()?
+        .build()
-        .wait_until_workspaces_are_initialized()?;
+        .wait_until_workspaces_are_initialized();
 
-    let show_message_params = server.await_notification::<ShowMessage>()?;
+    let show_message_params = server.await_notification::<ShowMessage>();
 
     insta::assert_json_snapshot!(show_message_params, @r#"
     {
@@ -443,10 +443,10 @@ fn register_rename_capability_when_enabled() -> Result<()> {
         .with_workspace(workspace_root, None)?
         .with_initialization_options(ClientOptions::default().with_experimental_rename(true))
         .enable_rename_dynamic_registration(true)
-        .build()?
+        .build()
-        .wait_until_workspaces_are_initialized()?;
+        .wait_until_workspaces_are_initialized();
 
-    let (_, params) = server.await_request::<RegisterCapability>()?;
+    let (_, params) = server.await_request::<RegisterCapability>();
     let [registration] = params.registrations.as_slice() else {
         panic!(
             "Expected a single registration, got: {:#?}",
@@ -477,8 +477,8 @@ fn rename_available_without_dynamic_registration() -> Result<()> {
         .with_workspace(workspace_root, None)?
         .with_initialization_options(ClientOptions::default().with_experimental_rename(true))
         .enable_rename_dynamic_registration(false)
-        .build()?
+        .build()
-        .wait_until_workspaces_are_initialized()?;
+        .wait_until_workspaces_are_initialized();
 
     let initialization_result = server.initialization_result().unwrap();
     insta::assert_json_snapshot!(initialization_result.capabilities.rename_provider, @r#"
@@ -500,8 +500,8 @@ fn not_register_rename_capability_when_disabled() -> Result<()> {
         .with_workspace(workspace_root, None)?
         .with_initialization_options(ClientOptions::default().with_experimental_rename(false))
         .enable_rename_dynamic_registration(true)
-        .build()?
+        .build()
-        .wait_until_workspaces_are_initialized()?;
+        .wait_until_workspaces_are_initialized();
 
     // The `Drop` implementation will make sure that the client did not receive any registration
     // request.
@@ -525,10 +525,10 @@ fn register_multiple_capabilities() -> Result<()> {
         )
         .enable_rename_dynamic_registration(true)
         .enable_diagnostic_dynamic_registration(true)
-        .build()?
+        .build()
-        .wait_until_workspaces_are_initialized()?;
+        .wait_until_workspaces_are_initialized();
 
-    let (_, params) = server.await_request::<RegisterCapability>()?;
+    let (_, params) = server.await_request::<RegisterCapability>();
     let registrations = params.registrations;
 
     assert_eq!(registrations.len(), 2);

crates/ty_server/tests/e2e/inlay_hints.rs

@@ -25,14 +25,14 @@ y = foo(1)
         .with_workspace(workspace_root, None)?
         .with_file(foo, foo_content)?
         .enable_inlay_hints(true)
-        .build()?
+        .build()
-        .wait_until_workspaces_are_initialized()?;
+        .wait_until_workspaces_are_initialized();
 
     server.open_text_document(foo, &foo_content, 1);
-    let _ = server.await_notification::<PublishDiagnostics>()?;
+    let _ = server.await_notification::<PublishDiagnostics>();
 
     let hints = server
-        .inlay_hints_request(foo, Range::new(Position::new(0, 0), Position::new(6, 0)))?
+        .inlay_hints_request(foo, Range::new(Position::new(0, 0), Position::new(6, 0)))
         .unwrap();
 
     insta::assert_json_snapshot!(hints, @r#"
@@ -87,14 +87,14 @@ fn variable_inlay_hints_disabled() -> Result<()> {
         .with_workspace(workspace_root, None)?
         .with_file(foo, foo_content)?
         .enable_inlay_hints(true)
-        .build()?
+        .build()
-        .wait_until_workspaces_are_initialized()?;
+        .wait_until_workspaces_are_initialized();
 
     server.open_text_document(foo, &foo_content, 1);
-    let _ = server.await_notification::<PublishDiagnostics>()?;
+    let _ = server.await_notification::<PublishDiagnostics>();
 
     let hints = server
-        .inlay_hints_request(foo, Range::new(Position::new(0, 0), Position::new(0, 5)))?
+        .inlay_hints_request(foo, Range::new(Position::new(0, 0), Position::new(0, 5)))
         .unwrap();
 
     assert!(

crates/ty_server/tests/e2e/main.rs

@@ -34,7 +34,6 @@ mod notebook;
 mod publish_diagnostics;
 mod pull_diagnostics;
 
-use std::collections::hash_map::Entry;
 use std::collections::{BTreeMap, HashMap, VecDeque};
 use std::num::NonZeroUsize;
 use std::sync::{Arc, OnceLock};
@@ -88,31 +87,68 @@ fn setup_tracing() {
     });
 }
 
-/// Errors that can occur during testing
+/// Errors when receiving a notification or request from the server.
 #[derive(thiserror::Error, Debug)]
-pub(crate) enum TestServerError {
+pub(crate) enum ServerMessageError {
+    #[error("waiting for message timed out")]
+    Timeout,
+
+    #[error("server disconnected")]
+    ServerDisconnected,
+
+    #[error("Failed to deserialize message body: {0}")]
+    DeserializationError(#[from] serde_json::Error),
+}
+
+impl From<ReceiveError> for ServerMessageError {
+    fn from(value: ReceiveError) -> Self {
+        match value {
+            ReceiveError::Timeout => Self::Timeout,
+            ReceiveError::ServerDisconnected => Self::ServerDisconnected,
+        }
+    }
+}
+
+/// Errors when receiving a response from the server.
+#[derive(thiserror::Error, Debug)]
+pub(crate) enum AwaitResponseError {
     /// The response came back, but was an error response, not a successful one.
-    #[error("Response error: {0:?}")]
+    #[error("request failed because the server replied with an error: {0:?}")]
-    ResponseError(ResponseError),
+    RequestFailed(ResponseError),
 
-    #[error("Invalid response message for request {0}: {1:?}")]
+    #[error("malformed response message with both result and error: {0:#?}")]
-    InvalidResponse(RequestId, Box<Response>),
+    MalformedResponse(Box<Response>),
 
-    #[error("Got a duplicate response for request ID {0}: {1:?}")]
+    #[error("received multiple responses for the same request ID: {0:#?}")]
-    DuplicateResponse(RequestId, Box<Response>),
+    MultipleResponses(Box<[Response]>),
 
-    #[error("Failed to receive message from server: {0}")]
+    #[error("waiting for response timed out")]
-    RecvTimeoutError(RecvTimeoutError),
+    Timeout,
+
+    #[error("server disconnected")]
+    ServerDisconnected,
+
+    #[error("failed to deserialize response result: {0}")]
+    DeserializationError(#[from] serde_json::Error),
 }
 
-impl TestServerError {
+impl From<ReceiveError> for AwaitResponseError {
-    fn is_disconnected(&self) -> bool {
+    fn from(err: ReceiveError) -> Self {
-        matches!(
+        match err {
-            self,
+            ReceiveError::Timeout => Self::Timeout,
-            TestServerError::RecvTimeoutError(RecvTimeoutError::Disconnected)
+            ReceiveError::ServerDisconnected => Self::ServerDisconnected,
-        )
         }
     }
+}
+
+#[derive(thiserror::Error, Debug)]
+pub(crate) enum ReceiveError {
+    #[error("waiting for message timed out")]
+    Timeout,
+
+    #[error("server disconnected")]
+    ServerDisconnected,
+}
 
 /// A test server for the ty language server that provides helpers for sending requests,
 /// correlating responses, and handling notifications.
@@ -137,8 +173,12 @@ pub(crate) struct TestServer {
     /// Incrementing counter to automatically generate request IDs
     request_counter: i32,
 
-    /// A mapping of request IDs to responses received from the server
+    /// A mapping of request IDs to responses received from the server.
-    responses: FxHashMap<RequestId, Response>,
+    ///
+    /// Valid responses contain exactly one response but may contain multiple responses
+    /// when the server sends multiple responses for a single request.
+    /// The responses are guaranteed to never be empty.
+    responses: FxHashMap<RequestId, smallvec::SmallVec<[Response; 1]>>,
 
     /// An ordered queue of all the notifications received from the server
     notifications: VecDeque<lsp_server::Notification>,
@@ -164,7 +204,7 @@ impl TestServer {
         test_context: TestContext,
         capabilities: ClientCapabilities,
         initialization_options: Option<ClientOptions>,
-    ) -> Result<Self> {
+    ) -> Self {
         setup_tracing();
 
         tracing::debug!("Starting test client with capabilities {:#?}", capabilities);
@@ -227,7 +267,7 @@ impl TestServer {
         workspace_folders: Vec<WorkspaceFolder>,
         capabilities: ClientCapabilities,
         initialization_options: Option<ClientOptions>,
-    ) -> Result<Self> {
+    ) -> Self {
         let init_params = InitializeParams {
             capabilities,
             workspace_folders: Some(workspace_folders),
@@ -240,10 +280,10 @@ impl TestServer {
         };
 
         let init_request_id = self.send_request::<Initialize>(init_params);
-        self.initialize_response = Some(self.await_response::<Initialize>(&init_request_id)?);
+        self.initialize_response = Some(self.await_response::<Initialize>(&init_request_id));
         self.send_notification::<Initialized>(InitializedParams {});
 
-        Ok(self)
+        self
     }
 
     /// Wait until the server has initialized all workspaces.
@@ -252,28 +292,18 @@ impl TestServer {
     /// server, and handles the request.
     ///
     /// This should only be called if the server is expected to send this request.
-    pub(crate) fn wait_until_workspaces_are_initialized(mut self) -> Result<Self> {
+    #[track_caller]
-        let (request_id, params) = self.await_request::<WorkspaceConfiguration>()?;
+    pub(crate) fn wait_until_workspaces_are_initialized(mut self) -> Self {
-        self.handle_workspace_configuration_request(request_id, &params)?;
+        let (request_id, params) = self.await_request::<WorkspaceConfiguration>();
-        Ok(self)
+        self.handle_workspace_configuration_request(request_id, &params);
+        self
     }
 
     /// Drain all messages from the server.
     fn drain_messages(&mut self) {
-        loop {
         // Don't wait too long to drain the messages, as this is called in the `Drop`
         // implementation which happens everytime the test ends.
-            match self.receive(Some(Duration::from_millis(10))) {
+        while let Ok(()) = self.receive(Some(Duration::from_millis(10))) {}
-                Ok(()) => {}
-                Err(TestServerError::RecvTimeoutError(_)) => {
-                    // Only break if we have no more messages to process.
-                    break;
-                }
-                Err(err) => {
-                    tracing::error!("Error while draining messages: {err:?}");
-                }
-            }
-        }
     }
 
     /// Validate that there are no pending messages from the server.
@@ -372,13 +402,48 @@ impl TestServer {
     /// This method will remove the response from the internal data structure, so it can only be
     /// called once per request ID.
     ///
+    /// # Panics
+    ///
+    /// If the server didn't send a response, the response failed with an error code, failed to deserialize,
+    /// or the server responded twice. Use [`Self::try_await_response`] if you want a non-panicking version.
+    ///
     /// [`send_request`]: TestServer::send_request
-    pub(crate) fn await_response<R>(&mut self, id: &RequestId) -> Result<R::Result>
+    #[track_caller]
+    pub(crate) fn await_response<R>(&mut self, id: &RequestId) -> R::Result
+    where
+        R: Request,
+    {
+        self.try_await_response::<R>(id, None)
+            .unwrap_or_else(|err| panic!("Failed to receive response for request {id}: {err}"))
+    }
+
+    /// Wait for a server response corresponding to the given request ID.
+    ///
+    /// This should only be called if a request was already sent to the server via [`send_request`]
+    /// which returns the request ID that should be used here.
+    ///
+    /// This method will remove the response from the internal data structure, so it can only be
+    /// called once per request ID.
+    ///
+    /// [`send_request`]: TestServer::send_request
+    pub(crate) fn try_await_response<R>(
+        &mut self,
+        id: &RequestId,
+        timeout: Option<Duration>,
+    ) -> Result<R::Result, AwaitResponseError>
     where
         R: Request,
     {
         loop {
-            if let Some(response) = self.responses.remove(id) {
+            if let Some(mut responses) = self.responses.remove(id) {
+                if responses.len() > 1 {
+                    return Err(AwaitResponseError::MultipleResponses(
+                        responses.into_boxed_slice(),
+                    ));
+                }
+
+                let response = responses.pop().unwrap();
+
                 match response {
                     Response {
                         error: None,
@@ -392,19 +457,15 @@ impl TestServer {
                         result: None,
                         ..
                     } => {
-                        return Err(TestServerError::ResponseError(err).into());
+                        return Err(AwaitResponseError::RequestFailed(err));
                     }
                     response => {
-                        return Err(TestServerError::InvalidResponse(
+                        return Err(AwaitResponseError::MalformedResponse(Box::new(response)));
-                            id.clone(),
-                            Box::new(response),
-                        )
-                        .into());
                     }
                 }
             }
 
-            self.receive_or_panic()?;
+            self.receive(timeout)?;
         }
     }
 
@@ -417,7 +478,31 @@ impl TestServer {
     ///
     /// This method will remove the notification from the internal data structure, so it should
     /// only be called if the notification is expected to be sent by the server.
-    pub(crate) fn await_notification<N: Notification>(&mut self) -> Result<N::Params> {
+    ///
+    /// # Panics
+    ///
+    /// If the server doesn't send the notification within the default timeout or
+    /// the notification failed to deserialize. Use [`Self::try_await_notification`] for
+    /// a panic-free alternative.
+    #[track_caller]
+    pub(crate) fn await_notification<N: Notification>(&mut self) -> N::Params {
+        self.try_await_notification::<N>(None)
+            .unwrap_or_else(|err| panic!("Failed to receive notification `{}`: {err}", N::METHOD))
+    }
+
+    /// Wait for a notification of the specified type from the server and return its parameters.
+    ///
+    /// The caller should ensure that the server is expected to send this notification type. It
+    /// will keep polling the server for this notification up to 10 times before giving up after
+    /// which it will return an error. It will also return an error if the notification is not
+    /// received within `recv_timeout` duration.
+    ///
+    /// This method will remove the notification from the internal data structure, so it should
+    /// only be called if the notification is expected to be sent by the server.
+    pub(crate) fn try_await_notification<N: Notification>(
+        &mut self,
+        timeout: Option<Duration>,
+    ) -> Result<N::Params, ServerMessageError> {
         for retry_count in 0..RETRY_COUNT {
             if retry_count > 0 {
                 tracing::info!("Retrying to receive `{}` notification", N::METHOD);
@@ -428,29 +513,30 @@ impl TestServer {
                 .position(|notification| N::METHOD == notification.method)
                 .and_then(|index| self.notifications.remove(index));
             if let Some(notification) = notification {
-                return Ok(serde_json::from_value(notification.params)?);
+                let params = serde_json::from_value(notification.params)?;
+                return Ok(params);
             }
-            self.receive_or_panic()?;
+
+            self.receive(timeout)?;
         }
-        Err(anyhow::anyhow!(
+
-            "Failed to receive `{}` notification after {RETRY_COUNT} retries",
+        Err(ServerMessageError::Timeout)
-            N::METHOD
-        ))
     }
 
     /// Collects `N` publish diagnostic notifications into a map, indexed by the document url.
     ///
     /// ## Panics
     /// If there are multiple publish diagnostics notifications for the same document.
+    #[track_caller]
     pub(crate) fn collect_publish_diagnostic_notifications(
         &mut self,
         count: usize,
-    ) -> Result<BTreeMap<lsp_types::Url, Vec<lsp_types::Diagnostic>>> {
+    ) -> BTreeMap<lsp_types::Url, Vec<lsp_types::Diagnostic>> {
         let mut results = BTreeMap::default();
 
         for _ in 0..count {
             let notification =
-                self.await_notification::<lsp_types::notification::PublishDiagnostics>()?;
+                self.await_notification::<lsp_types::notification::PublishDiagnostics>();
 
             if let Some(existing) =
                 results.insert(notification.uri.clone(), notification.diagnostics)
@@ -462,7 +548,7 @@ impl TestServer {
             }
         }
 
-        Ok(results)
+        results
     }
 
     /// Wait for a request of the specified type from the server and return the request ID and
@@ -475,7 +561,31 @@ impl TestServer {
     ///
     /// This method will remove the request from the internal data structure, so it should only be
     /// called if the request is expected to be sent by the server.
-    pub(crate) fn await_request<R: Request>(&mut self) -> Result<(RequestId, R::Params)> {
+    ///
+    /// # Panics
+    ///
+    /// If receiving the request fails.
+    #[track_caller]
+    pub(crate) fn await_request<R: Request>(&mut self) -> (RequestId, R::Params) {
+        self.try_await_request::<R>(None)
+            .unwrap_or_else(|err| panic!("Failed to receive server request `{}`: {err}", R::METHOD))
+    }
+
+    /// Wait for a request of the specified type from the server and return the request ID and
+    /// parameters.
+    ///
+    /// The caller should ensure that the server is expected to send this request type. It will
+    /// keep polling the server for this request up to 10 times before giving up after which it
+    /// will return an error. It can also return an error if the request is not received within
+    /// `recv_timeout` duration.
+    ///
+    /// This method will remove the request from the internal data structure, so it should only be
+    /// called if the request is expected to be sent by the server.
+    #[track_caller]
+    pub(crate) fn try_await_request<R: Request>(
+        &mut self,
+        timeout: Option<Duration>,
+    ) -> Result<(RequestId, R::Params), ServerMessageError> {
         for retry_count in 0..RETRY_COUNT {
             if retry_count > 0 {
                 tracing::info!("Retrying to receive `{}` request", R::METHOD);
@@ -489,12 +599,10 @@ impl TestServer {
                 let params = serde_json::from_value(request.params)?;
                 return Ok((request.id, params));
             }
-            self.receive_or_panic()?;
+
+            self.receive(timeout)?;
         }
-        Err(anyhow::anyhow!(
+        Err(ServerMessageError::Timeout)
-            "Failed to receive `{}` request after {RETRY_COUNT} retries",
-            R::METHOD
-        ))
     }
 
     /// Receive a message from the server.
@@ -503,45 +611,33 @@ impl TestServer {
     /// within that time, it will return an error.
     ///
     /// If `timeout` is `None`, it will use a default timeout of 10 second.
-    fn receive(&mut self, timeout: Option<Duration>) -> Result<(), TestServerError> {
+    fn receive(&mut self, timeout: Option<Duration>) -> Result<(), ReceiveError> {
         static DEFAULT_TIMEOUT: Duration = Duration::from_secs(10);
 
         let receiver = self.client_connection.as_ref().unwrap().receiver.clone();
         let message = receiver
             .recv_timeout(timeout.unwrap_or(DEFAULT_TIMEOUT))
-            .map_err(TestServerError::RecvTimeoutError)?;
+            .map_err(|err| match err {
+                RecvTimeoutError::Disconnected => ReceiveError::ServerDisconnected,
+                RecvTimeoutError::Timeout => ReceiveError::Timeout,
+            })?;
 
-        self.handle_message(message)?;
+        self.handle_message(message);
 
         for message in receiver.try_iter() {
-            self.handle_message(message)?;
+            self.handle_message(message);
         }
 
         Ok(())
     }
 
-    /// This is a convenience method that's same as [`receive`], but panics if the server got
-    /// disconnected. It will pass other errors as is.
-    ///
-    /// [`receive`]: TestServer::receive
-    fn receive_or_panic(&mut self) -> Result<(), TestServerError> {
-        if let Err(err) = self.receive(None) {
-            if err.is_disconnected() {
-                self.panic_on_server_disconnect();
-            } else {
-                return Err(err);
-            }
-        }
-        Ok(())
-    }
-
     /// Handle the incoming message from the server.
     ///
     /// This method will store the message as follows:
     /// - Requests are stored in `self.requests`
     /// - Responses are stored in `self.responses` with the request ID as the key
     /// - Notifications are stored in `self.notifications`
-    fn handle_message(&mut self, message: Message) -> Result<(), TestServerError> {
+    fn handle_message(&mut self, message: Message) {
         match message {
             Message::Request(request) => {
                 tracing::debug!("Received server request `{}`", &request.method);
@@ -549,24 +645,16 @@ impl TestServer {
             }
             Message::Response(response) => {
                 tracing::debug!("Received server response for request {}", &response.id);
-                match self.responses.entry(response.id.clone()) {
+                self.responses
-                    Entry::Occupied(existing) => {
+                    .entry(response.id.clone())
-                        return Err(TestServerError::DuplicateResponse(
+                    .or_default()
-                            response.id,
+                    .push(response);
-                            Box::new(existing.get().clone()),
-                        ));
-                    }
-                    Entry::Vacant(entry) => {
-                        entry.insert(response);
-                    }
-                }
             }
             Message::Notification(notification) => {
                 tracing::debug!("Received notification `{}`", &notification.method);
                 self.notifications.push_back(notification);
             }
         }
-        Ok(())
     }
 
     #[track_caller]
@@ -599,11 +687,12 @@ impl TestServer {
     /// Use the [`get_request`] method to wait for the server to send this request.
     ///
     /// [`get_request`]: TestServer::get_request
+    #[track_caller]
     fn handle_workspace_configuration_request(
         &mut self,
         request_id: RequestId,
         params: &ConfigurationParams,
-    ) -> Result<()> {
+    ) {
         let mut results = Vec::new();
 
         for item in &params.items {
@@ -618,7 +707,12 @@ impl TestServer {
                 // > If the client can't provide a configuration setting for a given scope
                 // > then null needs to be present in the returned array.
                 match item.section.as_deref() {
-                    Some("ty") => serde_json::to_value(options)?,
+                    Some("ty") => match serde_json::to_value(options) {
+                        Ok(value) => value,
+                        Err(err) => {
+                            panic!("Failed to deserialize workspace configuration options: {err}",)
+                        }
+                    },
                     Some(section) => {
                         tracing::debug!("Unrecognized section `{section}` for {scope_uri}");
                         serde_json::Value::Null
@@ -639,8 +733,6 @@ impl TestServer {
 
         let response = Response::new_ok(request_id, results);
         self.send(Message::Response(response));
-
-        Ok(())
     }
 
     /// Get the initialization result
@@ -711,7 +803,7 @@ impl TestServer {
         &mut self,
         path: impl AsRef<SystemPath>,
         previous_result_id: Option<String>,
-    ) -> Result<DocumentDiagnosticReportResult> {
+    ) -> DocumentDiagnosticReportResult {
         let params = DocumentDiagnosticParams {
             text_document: TextDocumentIdentifier {
                 uri: self.file_uri(path),
@@ -730,7 +822,7 @@ impl TestServer {
         &mut self,
         work_done_token: Option<lsp_types::NumberOrString>,
         previous_result_ids: Option<Vec<PreviousResultId>>,
-    ) -> Result<WorkspaceDiagnosticReportResult> {
+    ) -> WorkspaceDiagnosticReportResult {
         let params = WorkspaceDiagnosticParams {
             identifier: Some("ty".to_string()),
             previous_result_ids: previous_result_ids.unwrap_or_default(),
@@ -747,7 +839,7 @@ impl TestServer {
         &mut self,
         path: impl AsRef<SystemPath>,
         position: Position,
-    ) -> Result<Option<Hover>> {
+    ) -> Option<Hover> {
         let params = HoverParams {
             text_document_position_params: TextDocumentPositionParams {
                 text_document: TextDocumentIdentifier {
@@ -766,7 +858,7 @@ impl TestServer {
         &mut self,
         path: impl AsRef<SystemPath>,
         range: Range,
-    ) -> Result<Option<Vec<InlayHint>>> {
+    ) -> Option<Vec<InlayHint>> {
         let params = InlayHintParams {
             text_document: TextDocumentIdentifier {
                 uri: self.file_uri(path),
@@ -803,7 +895,7 @@ impl Drop for TestServer {
         // it dropped the client connection.
         let shutdown_error = if self.server_thread.is_some() && !self.shutdown_requested {
             let shutdown_id = self.send_request::<Shutdown>(());
-            match self.await_response::<Shutdown>(&shutdown_id) {
+            match self.try_await_response::<Shutdown>(&shutdown_id, None) {
                 Ok(()) => {
                     self.send_notification::<Exit>(());
 
@@ -834,10 +926,7 @@ impl Drop for TestServer {
                         );
                     }
                     Ok(message) => {
-                        // Ignore any errors: A duplicate pending message
+                        self.handle_message(message);
-                        // won't matter that much because `assert_no_pending_messages` will
-                        // panic anyway.
-                        let _ = self.handle_message(message);
                     }
                 }
             }
@@ -1048,7 +1137,7 @@ impl TestServerBuilder {
     }
 
     /// Build the test server
-    pub(crate) fn build(self) -> Result<TestServer> {
+    pub(crate) fn build(self) -> TestServer {
         TestServer::new(
             self.workspaces,
             self.test_context,

crates/ty_server/tests/e2e/notebook.rs

@@ -8,8 +8,8 @@ use crate::{TestServer, TestServerBuilder};
 #[test]
 fn publish_diagnostics_open() -> anyhow::Result<()> {
     let mut server = TestServerBuilder::new()?
-        .build()?
+        .build()
-        .wait_until_workspaces_are_initialized()?;
+        .wait_until_workspaces_are_initialized();
 
     server.initialization_result().unwrap();
 
@@ -42,11 +42,11 @@ type Style = Literal["italic", "bold", "underline"]"#,
     builder.open(&mut server);
 
     let cell1_diagnostics =
-        server.await_notification::<lsp_types::notification::PublishDiagnostics>()?;
+        server.await_notification::<lsp_types::notification::PublishDiagnostics>();
     let cell2_diagnostics =
-        server.await_notification::<lsp_types::notification::PublishDiagnostics>()?;
+        server.await_notification::<lsp_types::notification::PublishDiagnostics>();
     let cell3_diagnostics =
-        server.await_notification::<lsp_types::notification::PublishDiagnostics>()?;
+        server.await_notification::<lsp_types::notification::PublishDiagnostics>();
 
     assert_json_snapshot!([cell1_diagnostics, cell2_diagnostics, cell3_diagnostics]);
 
@@ -56,8 +56,8 @@ type Style = Literal["italic", "bold", "underline"]"#,
 #[test]
 fn diagnostic_end_of_file() -> anyhow::Result<()> {
     let mut server = TestServerBuilder::new()?
-        .build()?
+        .build()
-        .wait_until_workspaces_are_initialized()?;
+        .wait_until_workspaces_are_initialized();
 
     server.initialization_result().unwrap();
 
@@ -90,7 +90,7 @@ IOError"#,
 
     let notebook_url = builder.open(&mut server);
 
-    server.collect_publish_diagnostic_notifications(3)?;
+    server.collect_publish_diagnostic_notifications(3);
 
     server.send_notification::<lsp_types::notification::DidChangeNotebookDocument>(
         lsp_types::DidChangeNotebookDocumentParams {
@@ -122,7 +122,7 @@ IOError"#,
         },
     );
 
-    let diagnostics = server.collect_publish_diagnostic_notifications(3)?;
+    let diagnostics = server.collect_publish_diagnostic_notifications(3);
     assert_json_snapshot!(diagnostics);
 
     Ok(())
@@ -131,8 +131,8 @@ IOError"#,
 #[test]
 fn semantic_tokens() -> anyhow::Result<()> {
     let mut server = TestServerBuilder::new()?
-        .build()?
+        .build()
-        .wait_until_workspaces_are_initialized()?;
+        .wait_until_workspaces_are_initialized();
 
     server.initialization_result().unwrap();
 
@@ -164,13 +164,13 @@ type Style = Literal["italic", "bold", "underline"]"#,
 
     builder.open(&mut server);
 
-    let cell1_tokens = semantic_tokens_full_for_cell(&mut server, &first_cell)?;
+    let cell1_tokens = semantic_tokens_full_for_cell(&mut server, &first_cell);
-    let cell2_tokens = semantic_tokens_full_for_cell(&mut server, &second_cell)?;
+    let cell2_tokens = semantic_tokens_full_for_cell(&mut server, &second_cell);
-    let cell3_tokens = semantic_tokens_full_for_cell(&mut server, &third_cell)?;
+    let cell3_tokens = semantic_tokens_full_for_cell(&mut server, &third_cell);
 
     assert_json_snapshot!([cell1_tokens, cell2_tokens, cell3_tokens]);
 
-    server.collect_publish_diagnostic_notifications(3)?;
+    server.collect_publish_diagnostic_notifications(3);
 
     Ok(())
 }
@@ -178,8 +178,8 @@ type Style = Literal["italic", "bold", "underline"]"#,
 #[test]
 fn swap_cells() -> anyhow::Result<()> {
     let mut server = TestServerBuilder::new()?
-        .build()?
+        .build()
-        .wait_until_workspaces_are_initialized()?;
+        .wait_until_workspaces_are_initialized();
 
     server.initialization_result().unwrap();
 
@@ -196,7 +196,7 @@ fn swap_cells() -> anyhow::Result<()> {
 
     let notebook = builder.open(&mut server);
 
-    let diagnostics = server.collect_publish_diagnostic_notifications(3)?;
+    let diagnostics = server.collect_publish_diagnostic_notifications(3);
     assert_json_snapshot!(diagnostics, @r###"
     {
       "vscode-notebook-cell://src/test.ipynb#0": [
@@ -265,7 +265,7 @@ fn swap_cells() -> anyhow::Result<()> {
         },
     );
 
-    let diagnostics = server.collect_publish_diagnostic_notifications(3)?;
+    let diagnostics = server.collect_publish_diagnostic_notifications(3);
 
     assert_json_snapshot!(diagnostics, @r###"
     {
@@ -285,8 +285,8 @@ fn auto_import() -> anyhow::Result<()> {
             SystemPath::new("src"),
             Some(ClientOptions::default().with_experimental_auto_import(true)),
         )?
-        .build()?
+        .build()
-        .wait_until_workspaces_are_initialized()?;
+        .wait_until_workspaces_are_initialized();
 
     server.initialization_result().unwrap();
 
@@ -305,9 +305,9 @@ b: Litera
 
     builder.open(&mut server);
 
-    server.collect_publish_diagnostic_notifications(2)?;
+    server.collect_publish_diagnostic_notifications(2);
 
-    let completions = literal_completions(&mut server, &second_cell, Position::new(1, 9))?;
+    let completions = literal_completions(&mut server, &second_cell, Position::new(1, 9));
 
     assert_json_snapshot!(completions);
 
@@ -321,8 +321,8 @@ fn auto_import_same_cell() -> anyhow::Result<()> {
             SystemPath::new("src"),
             Some(ClientOptions::default().with_experimental_auto_import(true)),
         )?
-        .build()?
+        .build()
-        .wait_until_workspaces_are_initialized()?;
+        .wait_until_workspaces_are_initialized();
 
     server.initialization_result().unwrap();
 
@@ -336,9 +336,9 @@ b: Litera
 
     builder.open(&mut server);
 
-    server.collect_publish_diagnostic_notifications(1)?;
+    server.collect_publish_diagnostic_notifications(1);
 
-    let completions = literal_completions(&mut server, &first_cell, Position::new(1, 9))?;
+    let completions = literal_completions(&mut server, &first_cell, Position::new(1, 9));
 
     assert_json_snapshot!(completions);
 
@@ -352,8 +352,8 @@ fn auto_import_from_future() -> anyhow::Result<()> {
             SystemPath::new("src"),
             Some(ClientOptions::default().with_experimental_auto_import(true)),
         )?
-        .build()?
+        .build()
-        .wait_until_workspaces_are_initialized()?;
+        .wait_until_workspaces_are_initialized();
 
     server.initialization_result().unwrap();
 
@@ -369,9 +369,9 @@ b: Litera
 
     builder.open(&mut server);
 
-    server.collect_publish_diagnostic_notifications(2)?;
+    server.collect_publish_diagnostic_notifications(2);
 
-    let completions = literal_completions(&mut server, &second_cell, Position::new(1, 9))?;
+    let completions = literal_completions(&mut server, &second_cell, Position::new(1, 9));
 
     assert_json_snapshot!(completions);
 
@@ -385,8 +385,8 @@ fn auto_import_docstring() -> anyhow::Result<()> {
             SystemPath::new("src"),
             Some(ClientOptions::default().with_experimental_auto_import(true)),
         )?
-        .build()?
+        .build()
-        .wait_until_workspaces_are_initialized()?;
+        .wait_until_workspaces_are_initialized();
 
     server.initialization_result().unwrap();
 
@@ -405,9 +405,9 @@ b: Litera
 
     builder.open(&mut server);
 
-    server.collect_publish_diagnostic_notifications(2)?;
+    server.collect_publish_diagnostic_notifications(2);
 
-    let completions = literal_completions(&mut server, &second_cell, Position::new(1, 9))?;
+    let completions = literal_completions(&mut server, &second_cell, Position::new(1, 9));
 
     assert_json_snapshot!(completions);
 
@@ -417,7 +417,7 @@ b: Litera
 fn semantic_tokens_full_for_cell(
     server: &mut TestServer,
     cell_uri: &lsp_types::Url,
-) -> crate::Result<Option<lsp_types::SemanticTokensResult>> {
+) -> Option<lsp_types::SemanticTokensResult> {
     let cell1_tokens_req_id = server.send_request::<lsp_types::request::SemanticTokensFullRequest>(
         lsp_types::SemanticTokensParams {
             work_done_progress_params: lsp_types::WorkDoneProgressParams::default(),
@@ -502,7 +502,7 @@ fn literal_completions(
     server: &mut TestServer,
     cell: &lsp_types::Url,
     position: Position,
-) -> crate::Result<Vec<lsp_types::CompletionItem>> {
+) -> Vec<lsp_types::CompletionItem> {
     let completions_id =
         server.send_request::<lsp_types::request::Completion>(lsp_types::CompletionParams {
             text_document_position: lsp_types::TextDocumentPositionParams {
@@ -520,14 +520,14 @@ fn literal_completions(
     // There are a ton of imports we don't care about in here...
     // The import bit is that an edit is always restricted to the current cell. That means,
     // we can't add `Literal` to the `from typing import TYPE_CHECKING` import in cell 1
-    let completions = server.await_response::<lsp_types::request::Completion>(&completions_id)?;
+    let completions = server.await_response::<lsp_types::request::Completion>(&completions_id);
     let mut items = match completions {
         Some(CompletionResponse::Array(array)) => array,
         Some(CompletionResponse::List(lsp_types::CompletionList { items, .. })) => items,
-        None => return Ok(vec![]),
+        None => return vec![],
     };
 
     items.retain(|item| item.label.starts_with("Litera"));
 
-    Ok(items)
+    items
 }

crates/ty_server/tests/e2e/publish_diagnostics.rs

@@ -17,11 +17,11 @@ def foo() -> str:
         .with_workspace(workspace_root, None)?
         .with_file(foo, foo_content)?
         .enable_pull_diagnostics(false)
-        .build()?
+        .build()
-        .wait_until_workspaces_are_initialized()?;
+        .wait_until_workspaces_are_initialized();
 
     server.open_text_document(foo, &foo_content, 1);
-    let diagnostics = server.await_notification::<PublishDiagnostics>()?;
+    let diagnostics = server.await_notification::<PublishDiagnostics>();
 
     insta::assert_debug_snapshot!(diagnostics);
 

crates/ty_server/tests/e2e/pull_diagnostics.rs

@@ -1,3 +1,5 @@
+use std::time::Duration;
+
 use anyhow::Result;
 use insta::{assert_compact_debug_snapshot, assert_debug_snapshot};
 use lsp_server::RequestId;
@@ -9,7 +11,7 @@ use lsp_types::{
 use ruff_db::system::SystemPath;
 use ty_server::{ClientOptions, DiagnosticMode, PartialWorkspaceProgress};
 
-use crate::{TestServer, TestServerBuilder, TestServerError};
+use crate::{AwaitResponseError, TestServer, TestServerBuilder};
 
 #[test]
 fn on_did_open() -> Result<()> {
@@ -26,11 +28,11 @@ def foo() -> str:
         .with_workspace(workspace_root, None)?
         .with_file(foo, foo_content)?
         .enable_pull_diagnostics(true)
-        .build()?
+        .build()
-        .wait_until_workspaces_are_initialized()?;
+        .wait_until_workspaces_are_initialized();
 
     server.open_text_document(foo, &foo_content, 1);
-    let diagnostics = server.document_diagnostic_request(foo, None)?;
+    let diagnostics = server.document_diagnostic_request(foo, None);
 
     assert_debug_snapshot!(diagnostics);
 
@@ -52,13 +54,13 @@ def foo() -> str:
         .with_workspace(workspace_root, None)?
         .with_file(foo, foo_content)?
         .enable_pull_diagnostics(true)
-        .build()?
+        .build()
-        .wait_until_workspaces_are_initialized()?;
+        .wait_until_workspaces_are_initialized();
 
     server.open_text_document(foo, &foo_content, 1);
 
     // First request with no previous result ID
-    let first_response = server.document_diagnostic_request(foo, None)?;
+    let first_response = server.document_diagnostic_request(foo, None);
 
     // Extract result ID from first response
     let result_id = match &first_response {
@@ -74,7 +76,7 @@ def foo() -> str:
     };
 
     // Second request with the previous result ID - should return Unchanged
-    let second_response = server.document_diagnostic_request(foo, Some(result_id))?;
+    let second_response = server.document_diagnostic_request(foo, Some(result_id));
 
     // Verify it's an unchanged report
     match second_response {
@@ -108,13 +110,13 @@ def foo() -> str:
         .with_workspace(workspace_root, None)?
         .with_file(foo, foo_content_v1)?
         .enable_pull_diagnostics(true)
-        .build()?
+        .build()
-        .wait_until_workspaces_are_initialized()?;
+        .wait_until_workspaces_are_initialized();
 
     server.open_text_document(foo, &foo_content_v1, 1);
 
     // First request with no previous result ID
-    let first_response = server.document_diagnostic_request(foo, None)?;
+    let first_response = server.document_diagnostic_request(foo, None);
 
     // Extract result ID from first response
     let result_id = match &first_response {
@@ -141,7 +143,7 @@ def foo() -> str:
     );
 
     // Second request with the previous result ID - should return a new full report
-    let second_response = server.document_diagnostic_request(foo, Some(result_id))?;
+    let second_response = server.document_diagnostic_request(foo, Some(result_id));
 
     // Verify it's a full report (not unchanged)
     match second_response {
@@ -228,16 +230,14 @@ def foo() -> str:
         .with_file(file_d, file_d_content_v1)?
         .with_file(file_e, file_e_content_v1)?
         .enable_pull_diagnostics(true)
-        .build()?
+        .build()
-        .wait_until_workspaces_are_initialized()?;
+        .wait_until_workspaces_are_initialized();
 
     server.open_text_document(file_a, &file_a_content, 1);
 
     // First request with no previous result IDs
-    let mut first_response = server.workspace_diagnostic_request(
+    let mut first_response = server
-        Some(NumberOrString::String("progress-1".to_string())),
+        .workspace_diagnostic_request(Some(NumberOrString::String("progress-1".to_string())), None);
-        None,
-    )?;
     sort_workspace_diagnostic_response(&mut first_response);
 
     assert_debug_snapshot!("workspace_diagnostic_initial_state", first_response);
@@ -309,7 +309,7 @@ def foo() -> str:
     let mut second_response = server.workspace_diagnostic_request(
         Some(NumberOrString::String("progress-2".to_string())),
         Some(previous_result_ids),
-    )?;
+    );
     sort_workspace_diagnostic_response(&mut second_response);
 
     // Consume all progress notifications sent during the second workspace diagnostics
@@ -339,10 +339,10 @@ def foo() -> str:
             ClientOptions::default().with_diagnostic_mode(DiagnosticMode::Workspace),
         )
         .enable_pull_diagnostics(true)
-        .build()?
+        .build()
-        .wait_until_workspaces_are_initialized()?;
+        .wait_until_workspaces_are_initialized();
 
-    let first_response = server.workspace_diagnostic_request(None, None).unwrap();
+    let first_response = server.workspace_diagnostic_request(None, None);
 
     // Extract result IDs from the first response
     let mut previous_result_ids = extract_result_ids_from_response(&first_response);
@@ -366,7 +366,7 @@ def foo() -> str:
     // The server needs to match the previous result IDs by the path, not the URL.
     assert_workspace_diagnostics_suspends_for_long_polling(&mut server, &workspace_request_id);
 
-    let second_response = shutdown_and_await_workspace_diagnostic(server, &workspace_request_id)?;
+    let second_response = shutdown_and_await_workspace_diagnostic(server, &workspace_request_id);
 
     assert_compact_debug_snapshot!(second_response, @"Report(WorkspaceDiagnosticReport { items: [] })");
 
@@ -382,7 +382,7 @@ fn filter_result_id() -> insta::internals::SettingsBindDropGuard {
 
 fn consume_all_progress_notifications(server: &mut TestServer) -> Result<()> {
     // Always consume Begin
-    let begin_params = server.await_notification::<lsp_types::notification::Progress>()?;
+    let begin_params = server.await_notification::<lsp_types::notification::Progress>();
 
     // The params are already the ProgressParams type
     let lsp_types::ProgressParamsValue::WorkDone(lsp_types::WorkDoneProgress::Begin(_)) =
@@ -394,7 +394,7 @@ fn consume_all_progress_notifications(server: &mut TestServer) -> Result<()> {
     // Consume Report notifications - there may be multiple based on number of files
     // Keep consuming until we hit the End notification
     loop {
-        let params = server.await_notification::<lsp_types::notification::Progress>()?;
+        let params = server.await_notification::<lsp_types::notification::Progress>();
 
         if let lsp_types::ProgressParamsValue::WorkDone(lsp_types::WorkDoneProgress::End(_)) =
             params.value
@@ -442,8 +442,8 @@ def foo() -> str:
 
     let mut server = builder
         .enable_pull_diagnostics(true)
-        .build()?
+        .build()
-        .wait_until_workspaces_are_initialized()?;
+        .wait_until_workspaces_are_initialized();
 
     let partial_token = lsp_types::ProgressToken::String("streaming-diagnostics".to_string());
     let request_id = server.send_request::<WorkspaceDiagnosticRequest>(WorkspaceDiagnosticParams {
@@ -461,7 +461,7 @@ def foo() -> str:
 
     // First, read the response of the workspace diagnostic request.
     // Note: This response comes after the progress notifications but it simplifies the test to read it first.
-    let final_response = server.await_response::<WorkspaceDiagnosticRequest>(&request_id)?;
+    let final_response = server.await_response::<WorkspaceDiagnosticRequest>(&request_id);
 
     // Process the final report.
     // This should always be a partial report. However, the type definition in the LSP specification
@@ -478,7 +478,9 @@ def foo() -> str:
     received_results += response_items.len();
 
     // Collect any partial results sent via progress notifications
-    while let Ok(params) = server.await_notification::<PartialWorkspaceProgress>() {
+    while let Ok(params) =
+        server.try_await_notification::<PartialWorkspaceProgress>(Some(Duration::from_secs(1)))
+    {
         if params.token == partial_token {
             let streamed_items = match params.value {
                 // Ideally we'd assert that only the first response is a full report
@@ -531,15 +533,15 @@ fn workspace_diagnostic_streaming_with_caching() -> Result<()> {
 
     let mut server = builder
         .enable_pull_diagnostics(true)
-        .build()?
+        .build()
-        .wait_until_workspaces_are_initialized()?;
+        .wait_until_workspaces_are_initialized();
 
     server.open_text_document(SystemPath::new("src/error_0.py"), &error_content, 1);
     server.open_text_document(SystemPath::new("src/error_1.py"), &error_content, 1);
     server.open_text_document(SystemPath::new("src/error_2.py"), &error_content, 1);
 
     // First request to get result IDs (non-streaming for simplicity)
-    let first_response = server.workspace_diagnostic_request(None, None)?;
+    let first_response = server.workspace_diagnostic_request(None, None);
 
     let result_ids = extract_result_ids_from_response(&first_response);
 
@@ -590,7 +592,7 @@ fn workspace_diagnostic_streaming_with_caching() -> Result<()> {
             },
         });
 
-    let final_response2 = server.await_response::<WorkspaceDiagnosticRequest>(&request2_id)?;
+    let final_response2 = server.await_response::<WorkspaceDiagnosticRequest>(&request2_id);
 
     let mut all_items = Vec::new();
 
@@ -605,7 +607,9 @@ fn workspace_diagnostic_streaming_with_caching() -> Result<()> {
     all_items.extend(items);
 
     // Collect any partial results sent via progress notifications
-    while let Ok(params) = server.await_notification::<PartialWorkspaceProgress>() {
+    while let Ok(params) =
+        server.try_await_notification::<PartialWorkspaceProgress>(Some(Duration::from_secs(1)))
+    {
         if params.token == partial_token {
             let streamed_items = match params.value {
                 // Ideally we'd assert that only the first response is a full report
@@ -679,7 +683,7 @@ def hello() -> str:
     assert_workspace_diagnostics_suspends_for_long_polling(&mut server, &request_id);
 
     // The workspace diagnostic request should now respond with an empty report
-    let workspace_response = shutdown_and_await_workspace_diagnostic(server, &request_id)?;
+    let workspace_response = shutdown_and_await_workspace_diagnostic(server, &request_id);
 
     // Verify we got an empty report (default response during shutdown)
     assert_debug_snapshot!(
@@ -733,7 +737,7 @@ def hello() -> str:
     );
 
     // The workspace diagnostic request should now complete with the new diagnostic
-    let workspace_response = server.await_response::<WorkspaceDiagnosticRequest>(&request_id)?;
+    let workspace_response = server.await_response::<WorkspaceDiagnosticRequest>(&request_id);
 
     // Verify we got a report with one file containing the new diagnostic
     assert_debug_snapshot!(
@@ -775,7 +779,7 @@ def hello() -> str:
     server.cancel(&request_id);
 
     // The workspace diagnostic request should now respond with a cancellation response (Err).
-    let result = server.await_response::<WorkspaceDiagnosticRequest>(&request_id);
+    let result = server.try_await_response::<WorkspaceDiagnosticRequest>(&request_id, None);
     assert_debug_snapshot!(
         "workspace_diagnostic_long_polling_cancellation_result",
         result
@@ -833,7 +837,7 @@ def hello() -> str:
     );
 
     // First request should complete with diagnostics
-    let first_response = server.await_response::<WorkspaceDiagnosticRequest>(&request_id_1)?;
+    let first_response = server.await_response::<WorkspaceDiagnosticRequest>(&request_id_1);
 
     // Extract result IDs from the first response for the second request
     let previous_result_ids = extract_result_ids_from_response(&first_response);
@@ -866,7 +870,7 @@ def hello() -> str:
     );
 
     // Second request should complete with the fix (no diagnostics)
-    let second_response = server.await_response::<WorkspaceDiagnosticRequest>(&request_id_2)?;
+    let second_response = server.await_response::<WorkspaceDiagnosticRequest>(&request_id_2);
 
     // Snapshot both responses to verify the full cycle
     assert_debug_snapshot!(
@@ -887,15 +891,15 @@ fn create_workspace_server_with_file(
     file_path: &SystemPath,
     file_content: &str,
 ) -> Result<TestServer> {
-    TestServerBuilder::new()?
+    Ok(TestServerBuilder::new()?
         .with_workspace(workspace_root, None)?
         .with_file(file_path, file_content)?
         .with_initialization_options(
             ClientOptions::default().with_diagnostic_mode(DiagnosticMode::Workspace),
         )
         .enable_pull_diagnostics(true)
-        .build()?
+        .build()
-        .wait_until_workspaces_are_initialized()
+        .wait_until_workspaces_are_initialized())
 }
 
 /// Sends a workspace diagnostic request to the server.
@@ -917,7 +921,7 @@ fn send_workspace_diagnostic_request(server: &mut TestServer) -> lsp_server::Req
 fn shutdown_and_await_workspace_diagnostic(
     mut server: TestServer,
     request_id: &RequestId,
-) -> Result<WorkspaceDiagnosticReportResult> {
+) -> WorkspaceDiagnosticReportResult {
     // Send shutdown request - this should cause the suspended workspace diagnostic request to respond
     let shutdown_id = server.send_request::<lsp_types::request::Shutdown>(());
 
@@ -925,7 +929,7 @@ fn shutdown_and_await_workspace_diagnostic(
     let workspace_response = server.await_response::<WorkspaceDiagnosticRequest>(request_id);
 
     // Complete shutdown sequence
-    server.await_response::<lsp_types::request::Shutdown>(&shutdown_id)?;
+    server.await_response::<lsp_types::request::Shutdown>(&shutdown_id);
     server.send_notification::<lsp_types::notification::Exit>(());
 
     workspace_response
@@ -936,19 +940,19 @@ fn assert_workspace_diagnostics_suspends_for_long_polling(
     server: &mut TestServer,
     request_id: &lsp_server::RequestId,
 ) {
-    match server.await_response::<WorkspaceDiagnosticRequest>(request_id) {
+    match server
+        .try_await_response::<WorkspaceDiagnosticRequest>(request_id, Some(Duration::from_secs(2)))
+    {
         Ok(_) => {
             panic!("Expected workspace diagnostic request to suspend for long-polling.");
         }
-        Err(error) => {
+        Err(AwaitResponseError::Timeout) => {
-            if let Some(test_error) = error.downcast_ref::<TestServerError>() {
+            // Ok
-                assert!(
-                    matches!(test_error, TestServerError::RecvTimeoutError(_)),
-                    "Response should time out because the request is suspended for long-polling"
-                );
-            } else {
-                panic!("Unexpected error type: {error:?}");
         }
+        Err(err) => {
+            panic!(
+                "Response should time out because the request is suspended for long-polling but errored with: {err}"
+            )
         }
     }
 }

crates/ty_server/tests/e2e/snapshots/e2e__pull_diagnostics__workspace_diagnostic_long_polling_cancellation_result.snap

@@ -3,7 +3,7 @@ source: crates/ty_server/tests/e2e/pull_diagnostics.rs
 expression: result
 ---
 Err(
-    ResponseError(
+    RequestFailed(
         ResponseError {
             code: -32800,
             message: "request was cancelled by client",