Fix scheduler memory leak and close() race condition in CopilotSession

pom.xml

  Add mockito-core 5.17.0 as a test dependency.

src/main/java/com/github/copilot/sdk/CopilotSession.java

  Replace Executors.newSingleThreadScheduledExecutor with explicit
  ScheduledThreadPoolExecutor and enable removeOnCancelPolicy(true) so
  cancelled timeout tasks are purged from the work queue immediately.

  Wrap timeoutScheduler.schedule() in a try-catch for
  RejectedExecutionException. On rejection (close() race), the event
  subscription is cleaned up and the returned future completes
  exceptionally instead of throwing uncaught.

src/test/java/com/github/copilot/sdk/SchedulerShutdownRaceTest.java (new)

  TDD test that reproduces the scheduler shutdown race. Uses Mockito to
  stub JsonRpcClient.invoke(), then shuts down the scheduler without
  setting isTerminated, and asserts sendAndWait() returns a failed future
  rather than throwing RejectedExecutionException.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Ed Burns <edburns@microsoft.com>

Author: edburns

pom.xml

@@ -86,6 +86,12 @@
             <version>5.14.1</version>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.mockito</groupId>
+            <artifactId>mockito-core</artifactId>
+            <version>5.17.0</version>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 
     <build>

src/main/java/com/github/copilot/sdk/CopilotSession.java

@@ -14,7 +14,9 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Executors;
+import java.util.concurrent.RejectedExecutionException;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.ScheduledThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
@@ -417,14 +419,26 @@ public CompletableFuture<AssistantMessageEvent> sendAndWait(MessageOptions optio
             return null;
         });
 
+        var result = new CompletableFuture<AssistantMessageEvent>();
+
         // Schedule timeout on the shared session-level scheduler
-        var timeoutTask = timeoutScheduler.schedule(() -> {
-            if (!future.isDone()) {
-                future.completeExceptionally(new TimeoutException("sendAndWait timed out after " + timeoutMs + "ms"));
+        ScheduledFuture<?> timeoutTask;
+        try {
+            timeoutTask = timeoutScheduler.schedule(() -> {
+                if (!future.isDone()) {
+                    future.completeExceptionally(
+                            new TimeoutException("sendAndWait timed out after " + timeoutMs + "ms"));
+                }
+            }, timeoutMs, TimeUnit.MILLISECONDS);
+        } catch (RejectedExecutionException e) {
+            try {
+                subscription.close();
+            } catch (IOException closeEx) {
+                e.addSuppressed(closeEx);
             }
-        }, timeoutMs, TimeUnit.MILLISECONDS);
-
-        var result = new CompletableFuture<AssistantMessageEvent>();
+            result.completeExceptionally(e);
+            return result;
+        }
 
         // When inner future completes, run cleanup and propagate to result
         future.whenComplete((r, ex) -> {

src/test/java/com/github/copilot/sdk/SchedulerShutdownRaceTest.java

@@ -0,0 +1,65 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *--------------------------------------------------------------------------------------------*/
+
+package com.github.copilot.sdk;
+
+import static org.junit.jupiter.api.Assertions.*;
+import static org.mockito.ArgumentMatchers.*;
+import static org.mockito.Mockito.*;
+
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.RejectedExecutionException;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+
+import org.junit.jupiter.api.Test;
+
+import com.github.copilot.sdk.json.MessageOptions;
+
+/**
+ * Reproduces the race between {@code sendAndWait()} and {@code close()}.
+ * <p>
+ * If {@code close()} shuts down the timeout scheduler after
+ * {@code ensureNotTerminated()} passes but before
+ * {@code timeoutScheduler.schedule()} executes, the schedule call throws
+ * {@link RejectedExecutionException}. Without a fix the exception propagates
+ * uncaught, leaking the event subscription and leaving the returned future
+ * incomplete.
+ */
+public class SchedulerShutdownRaceTest {
+
+    @SuppressWarnings("unchecked")
+    @Test
+    void sendAndWaitShouldReturnFailedFutureWhenSchedulerIsShutDown() throws Exception {
+        // Build a session via reflection (package-private constructor)
+        var ctor = CopilotSession.class.getDeclaredConstructor(
+                String.class, JsonRpcClient.class, String.class);
+        ctor.setAccessible(true);
+
+        // Mock JsonRpcClient so send() returns a pending future instead of NPE
+        var mockRpc = mock(JsonRpcClient.class);
+        when(mockRpc.invoke(any(), any(), any()))
+                .thenReturn(new CompletableFuture<>());
+
+        var session = ctor.newInstance("race-test", mockRpc, null);
+
+        // Shut down the scheduler without setting isTerminated,
+        // simulating the race window between ensureNotTerminated() and schedule()
+        var schedulerField = CopilotSession.class.getDeclaredField("timeoutScheduler");
+        schedulerField.setAccessible(true);
+        var scheduler = (ScheduledExecutorService) schedulerField.get(session);
+        scheduler.shutdownNow();
+
+        // With the fix: sendAndWait returns a future that completes exceptionally.
+        // Without the fix: sendAndWait throws RejectedExecutionException directly.
+        CompletableFuture<?> result = session.sendAndWait(
+                new MessageOptions().setPrompt("test"), 5000);
+
+        assertNotNull(result, "sendAndWait should return a future, not throw");
+        var ex = assertThrows(ExecutionException.class,
+                () -> result.get(1, TimeUnit.SECONDS));
+        assertInstanceOf(RejectedExecutionException.class, ex.getCause());
+    }
+}