Speed up Rust E2E tests with shared clients

rust/src/lib.rs

@@ -2196,6 +2196,60 @@ impl Client {
         Ok(())
     }
 
+    /// Start this client's notification and request router on the current runtime.
+    /// This is test-harness plumbing, not part of the supported SDK API.
+    #[cfg(feature = "test-support")]
+    #[doc(hidden)]
+    pub fn start_router_for_test(&self) {
+        self.inner.router.ensure_started(
+            &self.inner.notification_tx,
+            &self.inner.request_rx,
+            self.inner.llm_inference.get().cloned(),
+            self.inner.on_github_telemetry.clone(),
+        );
+    }
+
+    #[cfg(feature = "test-support")]
+    #[doc(hidden)]
+    /// Disconnect and delete every session owned by this test client's isolated
+    /// runtime. This is test-harness plumbing, not part of the supported SDK API.
+    pub async fn cleanup_sessions_for_test(&self) -> Result<()> {
+        let mut first_error = None;
+
+        for session_id in self.inner.router.session_ids() {
+            if let Err(error) = self
+                .call(
+                    "session.destroy",
+                    Some(serde_json::json!({ "sessionId": session_id })),
+                )
+                .await
+                && first_error.is_none()
+            {
+                first_error = Some(error);
+            }
+            self.inner.router.unregister(&session_id);
+        }
+
+        match self.list_sessions(None).await {
+            Ok(sessions) => {
+                for session in sessions {
+                    if let Err(error) = self.delete_session(&session.session_id).await
+                        && first_error.is_none()
+                    {
+                        first_error = Some(error);
+                    }
+                }
+            }
+            Err(error) if first_error.is_none() => first_error = Some(error),
+            Err(_) => {}
+        }
+
+        match first_error {
+            Some(error) => Err(error),
+            None => Ok(()),
+        }
+    }
+
     /// Return the ID of the most recently updated session, if any.
     ///
     /// Useful for resuming the last conversation when the session ID was

rust/tests/e2e/abort.rs

@@ -10,77 +10,80 @@ use tokio::sync::{Mutex, mpsc, oneshot};
 
 use super::support::{
     DEFAULT_TEST_TOKEN, assistant_message_content, recv_with_timeout, wait_for_event,
-    with_e2e_context,
 };
 
 #[tokio::test]
 async fn should_abort_during_active_streaming() {
-    with_e2e_context("abort", "should_abort_during_active_streaming", |ctx| {
-        Box::pin(async move {
-            ctx.set_default_copilot_user();
-            let client = ctx.start_client().await;
-            let session = client
-                .create_session(ctx.approve_all_session_config().with_streaming(true))
-                .await
-                .expect("create session");
-            let events = session.subscribe();
+    super::support::with_dedicated_e2e_context(
+        "abort",
+        "should_abort_during_active_streaming",
+        |ctx| {
+            Box::pin(async move {
+                ctx.set_default_copilot_user();
+                let client = ctx.start_client().await;
+                let session = client
+                    .create_session(ctx.approve_all_session_config().with_streaming(true))
+                    .await
+                    .expect("create session");
+                let events = session.subscribe();
 
-            session
+                session
                 .send(
                     "Write a very long essay about the history of computing, covering every decade \
                          from the 1940s to the 2020s in great detail.",
                 )
                 .await
                 .expect("send long streaming turn");
 
-            let delta = wait_for_event(events, "assistant.message_delta", |event| {
-                event.parsed_type() == SessionEventType::AssistantMessageDelta
+                let delta = wait_for_event(events, "assistant.message_delta", |event| {
+                    event.parsed_type() == SessionEventType::AssistantMessageDelta
+                })
+                .await;
+                assert!(
+                    !delta
+                        .typed_data::<AssistantMessageDeltaData>()
+                        .expect("assistant.message_delta data")
+                        .delta_content
+                        .is_empty()
+                );
+
+                session.abort().await.expect("abort session");
+
+                // Session should be usable after abort. Wait for the specific recovery
+                // message rather than racing against a late idle from the aborted turn.
+                let recovery_events = session.subscribe();
+                session
+                    .send("Say 'abort_recovery_ok'.")
+                    .await
+                    .expect("send recovery");
+                let recovery = wait_for_event(
+                    recovery_events,
+                    "assistant.message containing abort_recovery_ok",
+                    |event| {
+                        event.parsed_type() == SessionEventType::AssistantMessage
+                            && assistant_message_content(event)
+                                .to_lowercase()
+                                .contains("abort_recovery_ok")
+                    },
+                )
+                .await;
+                assert!(
+                    assistant_message_content(&recovery)
+                        .to_lowercase()
+                        .contains("abort_recovery_ok")
+                );
+
+                session.disconnect().await.expect("disconnect session");
+                client.stop().await.expect("stop client");
             })
-            .await;
-            assert!(
-                !delta
-                    .typed_data::<AssistantMessageDeltaData>()
-                    .expect("assistant.message_delta data")
-                    .delta_content
-                    .is_empty()
-            );
-
-            session.abort().await.expect("abort session");
-
-            // Session should be usable after abort. Wait for the specific recovery
-            // message rather than racing against a late idle from the aborted turn.
-            let recovery_events = session.subscribe();
-            session
-                .send("Say 'abort_recovery_ok'.")
-                .await
-                .expect("send recovery");
-            let recovery = wait_for_event(
-                recovery_events,
-                "assistant.message containing abort_recovery_ok",
-                |event| {
-                    event.parsed_type() == SessionEventType::AssistantMessage
-                        && assistant_message_content(event)
-                            .to_lowercase()
-                            .contains("abort_recovery_ok")
-                },
-            )
-            .await;
-            assert!(
-                assistant_message_content(&recovery)
-                    .to_lowercase()
-                    .contains("abort_recovery_ok")
-            );
-
-            session.disconnect().await.expect("disconnect session");
-            client.stop().await.expect("stop client");
-        })
-    })
+        },
+    )
     .await;
 }
 
 #[tokio::test]
 async fn should_abort_during_active_tool_execution() {
-    with_e2e_context(
+    super::support::with_dedicated_e2e_context(
         "abort",
         "should_abort_during_active_tool_execution",
         |ctx| {

rust/tests/e2e/ask_user.rs

@@ -12,13 +12,11 @@ use github_copilot_sdk::{
 use serde_json::json;
 use tokio::sync::{Notify, mpsc};
 
-use super::support::{
-    DEFAULT_TEST_TOKEN, assistant_message_content, recv_with_timeout, with_e2e_context,
-};
+use super::support::{DEFAULT_TEST_TOKEN, assistant_message_content, recv_with_timeout};
 
 #[tokio::test]
 async fn should_invoke_user_input_handler_when_model_uses_ask_user_tool() {
-    with_e2e_context(
+    super::support::with_shared_e2e_context(&E2E,
         "ask_user",
         "should_invoke_user_input_handler_when_model_uses_ask_user_tool",
         |ctx| {
@@ -62,7 +60,7 @@ async fn should_invoke_user_input_handler_when_model_uses_ask_user_tool() {
 
 #[tokio::test]
 async fn should_receive_choices_in_user_input_request() {
-    with_e2e_context(
+    super::support::with_shared_e2e_context(&E2E,
         "ask_user",
         "should_receive_choices_in_user_input_request",
         |ctx| {
@@ -107,7 +105,7 @@ async fn should_receive_choices_in_user_input_request() {
 
 #[tokio::test]
 async fn should_handle_freeform_user_input_response() {
-    with_e2e_context(
+    super::support::with_shared_e2e_context(&E2E,
         "ask_user",
         "should_handle_freeform_user_input_response",
         |ctx| {
@@ -164,7 +162,8 @@ async fn should_handle_freeform_user_input_response() {
 /// the handler observes the sibling tool while its own request is still pending.
 #[tokio::test]
 async fn ask_user_does_not_block_sibling_tool_call_in_same_turn() {
-    with_e2e_context(
+    super::support::with_shared_e2e_context(
+        &E2E,
         "ask_user",
         "ask_user_does_not_block_sibling_tool_call_in_same_turn",
         |ctx| {
@@ -346,3 +345,5 @@ impl ToolHandler for SetMarkerTool {
         Ok(ToolResult::Text(format!("MARKER_{}", value.to_uppercase())))
     }
 }
+static E2E: super::support::SharedE2eGroup =
+    super::support::SharedE2eGroup::standard("ask_user", 4);

rust/tests/e2e/builtin_tools.rs

@@ -2,7 +2,7 @@ use std::time::Duration;
 
 use github_copilot_sdk::MessageOptions;
 
-use super::support::{assistant_message_content, with_e2e_context};
+use super::support::assistant_message_content;
 
 /// Built-in tool tests spawn a real CLI subprocess and execute actual shell /
 /// file tools. Under concurrent Windows CI load (e2e runs 4-wide on a 4-vCPU
@@ -16,7 +16,8 @@ fn message(prompt: &str) -> MessageOptions {
 
 #[tokio::test]
 async fn should_capture_exit_code_in_output() {
-    with_e2e_context(
+    super::support::with_shared_e2e_context(
+        &E2E,
         "builtin_tools",
         "should_capture_exit_code_in_output",
         |ctx| {
@@ -49,7 +50,7 @@ async fn should_capture_exit_code_in_output() {
 
 #[tokio::test]
 async fn should_capture_stderr_output() {
-    with_e2e_context("builtin_tools", "should_capture_stderr_output", |ctx| {
+    super::support::with_shared_e2e_context(&E2E, "builtin_tools", "should_capture_stderr_output", |ctx| {
         Box::pin(async move {
             if cfg!(windows) {
                 return;
@@ -77,7 +78,7 @@ async fn should_capture_stderr_output() {
 
 #[tokio::test]
 async fn should_read_file_with_line_range() {
-    with_e2e_context("builtin_tools", "should_read_file_with_line_range", |ctx| {
+    super::support::with_shared_e2e_context(&E2E, "builtin_tools", "should_read_file_with_line_range", |ctx| {
         Box::pin(async move {
             ctx.set_default_copilot_user();
             std::fs::write(ctx.work_dir().join("lines.txt"), "line1\nline2\nline3\nline4\nline5\n")
@@ -106,7 +107,7 @@ async fn should_read_file_with_line_range() {
 
 #[tokio::test]
 async fn should_handle_nonexistent_file_gracefully() {
-    with_e2e_context(
+    super::support::with_shared_e2e_context(&E2E,
         "builtin_tools",
         "should_handle_nonexistent_file_gracefully",
         |ctx| {
@@ -144,7 +145,7 @@ async fn should_handle_nonexistent_file_gracefully() {
 
 #[tokio::test]
 async fn should_edit_a_file_successfully() {
-    with_e2e_context("builtin_tools", "should_edit_a_file_successfully", |ctx| {
+    super::support::with_shared_e2e_context(&E2E, "builtin_tools", "should_edit_a_file_successfully", |ctx| {
         Box::pin(async move {
             ctx.set_default_copilot_user();
             std::fs::write(ctx.work_dir().join("edit_me.txt"), "Hello World\nGoodbye World\n")
@@ -171,7 +172,7 @@ async fn should_edit_a_file_successfully() {
 
 #[tokio::test]
 async fn should_create_a_new_file() {
-    with_e2e_context("builtin_tools", "should_create_a_new_file", |ctx| {
+    super::support::with_shared_e2e_context(&E2E, "builtin_tools", "should_create_a_new_file", |ctx| {
         Box::pin(async move {
             ctx.set_default_copilot_user();
             let client = ctx.start_client().await;
@@ -196,7 +197,7 @@ async fn should_create_a_new_file() {
 
 #[tokio::test]
 async fn should_search_for_patterns_in_files() {
-    with_e2e_context(
+    super::support::with_shared_e2e_context(&E2E,
         "builtin_tools",
         "should_search_for_patterns_in_files",
         |ctx| {
@@ -229,7 +230,7 @@ async fn should_search_for_patterns_in_files() {
 
 #[tokio::test]
 async fn should_find_files_by_pattern() {
-    with_e2e_context("builtin_tools", "should_find_files_by_pattern", |ctx| {
+    super::support::with_shared_e2e_context(&E2E, "builtin_tools", "should_find_files_by_pattern", |ctx| {
         Box::pin(async move {
             ctx.set_default_copilot_user();
             let src = ctx.work_dir().join("src");
@@ -256,3 +257,5 @@ async fn should_find_files_by_pattern() {
     })
     .await;
 }
+static E2E: super::support::SharedE2eGroup =
+    super::support::SharedE2eGroup::standard("builtin_tools", 8);