Harden session detach cleanup

Serialize and finalize disconnect cleanup consistently across SDKs, unregister manually detached Go and Java sessions, propagate Java detach failures, and teach fake runtimes the detach response contract.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

Copilot-Session: 69723dd2-e732-43c6-9f9a-a16c2de3a628

Author: jmoseley

dotnet/src/Session.cs

@@ -1939,18 +1939,17 @@ public async ValueTask DisposeAsync()
         finally
         {
             RemoveFromClient();
+            _eventHandlers = ImmutableInterlocked.InterlockedExchange(ref _eventHandlers, ImmutableArray<EventSubscription>.Empty);
+            _toolHandlers.Clear();
+            _commandHandlers.Clear();
+
+            _permissionHandler = null;
+            _userInputHandler = null;
+            _elicitationHandler = null;
+            _exitPlanModeHandler = null;
+            _autoModeSwitchHandler = null;
             GC.SuppressFinalize(this);
         }
-
-        _eventHandlers = ImmutableInterlocked.InterlockedExchange(ref _eventHandlers, ImmutableArray<EventSubscription>.Empty);
-        _toolHandlers.Clear();
-        _commandHandlers.Clear();
-
-        _permissionHandler = null;
-        _userInputHandler = null;
-        _elicitationHandler = null;
-        _exitPlanModeHandler = null;
-        _autoModeSwitchHandler = null;
     }
 
     [LoggerMessage(Level = LogLevel.Error, Message = "Unhandled exception in broadcast event handler")]

dotnet/test/E2E/ClientOptionsE2ETests.cs

@@ -1068,6 +1068,10 @@ function handleMessage(message) {
             writeResponse(message.id, { messageId: "fake-message" });
             return;
           }
+          if (message.method === "session.detach") {
+            writeResponse(message.id, { success: true });
+            return;
+          }
 
           writeResponse(message.id, {});
         }

go/client.go

@@ -924,6 +924,13 @@ func (c *Client) CreateSession(ctx context.Context, config *SessionConfig) (*Ses
 			"",
 			hasManagedSettings(config.EnableManagedSettings, config.ManagedSettings),
 		)
+		s.onDisconnected = func() {
+			c.sessionsMux.Lock()
+			defer c.sessionsMux.Unlock()
+			if c.sessions[sessionID] == s {
+				delete(c.sessions, sessionID)
+			}
+		}
 
 		s.registerTools(config.Tools)
 		s.registerPermissionHandler(config.OnPermissionRequest)
@@ -1258,6 +1265,13 @@ func (c *Client) ResumeSessionWithOptions(ctx context.Context, sessionID string,
 		"",
 		hasManagedSettings(config.EnableManagedSettings, config.ManagedSettings),
 	)
+	session.onDisconnected = func() {
+		c.sessionsMux.Lock()
+		defer c.sessionsMux.Unlock()
+		if c.sessions[sessionID] == session {
+			delete(c.sessions, sessionID)
+		}
+	}
 
 	session.registerTools(config.Tools)
 	session.registerPermissionHandler(config.OnPermissionRequest)

go/client_test.go

@@ -1939,6 +1939,36 @@ func TestClient_MCPAuthInterestRegistration(t *testing.T) {
 	})
 }
 
+func TestSessionDisconnectUnregistersBeforeClientStop(t *testing.T) {
+	client, requests, cleanup := newInMemoryClient(t)
+	defer cleanup()
+
+	session, err := client.CreateSession(t.Context(), &SessionConfig{
+		OnPermissionRequest: PermissionHandler.ApproveAll,
+	})
+	if err != nil {
+		t.Fatalf("CreateSession failed: %v", err)
+	}
+	requests.clear()
+
+	if err := session.Disconnect(); err != nil {
+		t.Fatalf("Disconnect failed: %v", err)
+	}
+	if err := client.Stop(); err != nil {
+		t.Fatalf("Stop failed: %v", err)
+	}
+
+	detachCount := 0
+	for _, request := range requests.snapshot() {
+		if request.Method == "session.detach" {
+			detachCount++
+		}
+	}
+	if detachCount != 1 {
+		t.Fatalf("expected exactly one session.detach request, got %d", detachCount)
+	}
+}
+
 type recordedRequest struct {
 	Method string
 	Params map[string]any

go/session.go

@@ -95,8 +95,11 @@ type Session struct {
 
 	// eventCh serializes user event handler dispatch. dispatchEvent enqueues;
 	// a single goroutine (processEvents) dequeues and invokes handlers in FIFO order.
-	eventCh   chan SessionEvent
-	closeOnce sync.Once // guards eventCh close so Disconnect is safe to call more than once
+	eventCh        chan SessionEvent
+	closeOnce      sync.Once // guards eventCh close so Disconnect is safe to call more than once
+	disconnectMu   sync.Mutex
+	disconnected   bool
+	onDisconnected func()
 
 	// RPC provides typed session-scoped RPC methods.
 	RPC *rpc.SessionRPC
@@ -1716,6 +1719,12 @@ func (s *Session) GetEvents(ctx context.Context) ([]SessionEvent, error) {
 //	    log.Printf("Failed to disconnect session: %v", err)
 //	}
 func (s *Session) Disconnect() error {
+	s.disconnectMu.Lock()
+	defer s.disconnectMu.Unlock()
+	if s.disconnected {
+		return nil
+	}
+
 	result, err := s.client.Request(context.Background(), "session.detach", sessionDetachRequest{SessionID: s.SessionID})
 	if err != nil {
 		return fmt.Errorf("failed to disconnect session: %w", err)
@@ -1731,6 +1740,7 @@ func (s *Session) Disconnect() error {
 		return fmt.Errorf("failed to disconnect session: %s", response.Error)
 	}
 
+	s.disconnected = true
 	s.closeOnce.Do(func() { close(s.eventCh) })
 
 	// Clear handlers
@@ -1754,6 +1764,9 @@ func (s *Session) Disconnect() error {
 	s.elicitationHandler = nil
 	s.elicitationMu.Unlock()
 
+	if s.onDisconnected != nil {
+		s.onDisconnected()
+	}
 	return nil
 }
 

java/src/main/java/com/github/copilot/CopilotClient.java

@@ -381,26 +381,32 @@ public CompletableFuture<Void> stop() {
 
         for (CopilotSession session : new ArrayList<>(sessions.values())) {
             Runnable closeTask = () -> {
-                try {
-                    session.close();
-                } catch (Exception e) {
-                    LOG.log(Level.WARNING, "Error closing session " + session.getSessionId(), e);
-                }
+                session.close();
             };
             CompletableFuture<Void> future;
             try {
                 future = CompletableFuture.runAsync(closeTask, executor);
             } catch (RejectedExecutionException e) {
                 LOG.log(Level.WARNING, "Executor rejected session close task; closing inline", e);
-                closeTask.run();
-                future = CompletableFuture.completedFuture(null);
+                try {
+                    closeTask.run();
+                    future = CompletableFuture.completedFuture(null);
+                } catch (RuntimeException closeError) {
+                    future = CompletableFuture.failedFuture(closeError);
+                }
             }
             closeFutures.add(future);
         }
         sessions.clear();
 
         return CompletableFuture.allOf(closeFutures.toArray(new CompletableFuture[0]))
-                .thenCompose(v -> cleanupConnection(true));
+                .handle((ignored, closeError) -> closeError)
+                .thenCompose(closeError -> cleanupConnection(true).thenApply(ignored -> {
+                    if (closeError != null) {
+                        throw new CompletionException(closeError);
+                    }
+                    return null;
+                }));
     }
 
     /**
@@ -571,6 +577,7 @@ public CompletableFuture<CopilotSession> createSession(SessionConfig config) {
                 long setupNanos = System.nanoTime();
                 var s = new CopilotSession(sid, connection.rpc);
                 s.setExecutor(executor);
+                s.setOnClosed(() -> sessions.remove(sid, s));
                 SessionRequestBuilder.configureSession(s, config);
                 if (extracted.transformCallbacks() != null) {
                     s.registerTransformCallbacks(extracted.transformCallbacks());
@@ -743,6 +750,7 @@ public CompletableFuture<CopilotSession> resumeSession(String sessionId, ResumeS
             long setupNanos = System.nanoTime();
             var session = new CopilotSession(sessionId, connection.rpc);
             session.setExecutor(executor);
+            session.setOnClosed(() -> sessions.remove(sessionId, session));
             SessionRequestBuilder.configureSession(session, config);
             sessions.put(sessionId, session);
             LoggingHelpers.logTiming(LOG, Level.FINE,

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

@@ -201,6 +201,8 @@ public final class CopilotSession implements AutoCloseable {
 
     /** Tracks whether this session instance has been terminated via close(). */
     private volatile boolean isTerminated = false;
+    private volatile Runnable onClosed = () -> {
+    };
 
     /**
      * Creates a new session with the given ID and RPC client.
@@ -252,6 +254,10 @@ void setExecutor(Executor executor) {
         this.executor = executor;
     }
 
+    void setOnClosed(Runnable onClosed) {
+        this.onClosed = onClosed;
+    }
+
     /**
      * Gets the unique identifier for this session.
      *
@@ -2302,19 +2308,20 @@ public void close() {
 
         timeoutScheduler.shutdownNow();
 
+        RuntimeException detachFailure = null;
         try {
             SessionDetachResponse response = rpc
                     .invoke("session.detach", Map.of("sessionId", sessionId), SessionDetachResponse.class)
                     .get(5, TimeUnit.SECONDS);
             if (response == null || !response.success()) {
-                LOG.log(Level.FINE, "Failed to detach session {0}: {1}",
-                        new Object[] {
-                                sessionId,
-                                response != null && response.error() != null ? response.error() : "unknown error"
-                        });
+                String detail = response != null && response.error() != null ? response.error() : "unknown error";
+                detachFailure = new IllegalStateException("Failed to detach session " + sessionId + ": " + detail);
             }
         } catch (Exception e) {
-            LOG.log(Level.FINE, "Error detaching session", e);
+            if (e instanceof InterruptedException) {
+                Thread.currentThread().interrupt();
+            }
+            detachFailure = new IllegalStateException("Failed to detach session " + sessionId, e);
         }
 
         eventHandlers.clear();
@@ -2326,13 +2333,17 @@ public void close() {
         exitPlanModeHandler.set(null);
         autoModeSwitchHandler.set(null);
         hooksHandler.set(null);
+        onClosed.run();
+
+        if (detachFailure != null) {
+            throw detachFailure;
+        }
     }
 
     // ===== Internal response types for agent API =====
 
     @JsonIgnoreProperties(ignoreUnknown = true)
-    private record SessionDetachResponse(@JsonProperty("success") boolean success,
-            @JsonProperty("error") String error) {
+    record SessionDetachResponse(@JsonProperty("success") boolean success, @JsonProperty("error") String error) {
     }
 
     @JsonIgnoreProperties(ignoreUnknown = true)

java/src/test/java/com/github/copilot/ClientOptionsE2ETest.java

@@ -265,6 +265,8 @@ function resultFor(message) {
                   return { sessionId: message.params?.sessionId ?? 'fake-session', openCanvases: [] };
                 case 'session.options.update':
                   return { success: true };
+                case 'session.detach':
+                  return { success: true };
                 default:
                   return {};
               }

java/src/test/java/com/github/copilot/TimeoutEdgeCaseTest.java

@@ -5,6 +5,7 @@
 package com.github.copilot;
 
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import java.io.ByteArrayOutputStream;
@@ -81,7 +82,7 @@ void testTimeoutDoesNotFireAfterSessionClose() throws Exception {
 
                 // close() blocks up to 5s on session.detach RPC. The 2s timeout
                 // fires during that window with the current per-call scheduler.
-                session.close();
+                assertThrows(IllegalStateException.class, session::close);
 
                 assertFalse(result.isDone(), "Future should not be completed by a timeout after session is closed. "
                         + "The per-call ScheduledExecutorService leaked a TimeoutException.");
@@ -126,6 +127,7 @@ void testSendAndWaitReusesTimeoutThread() throws Exception {
 
                 result1.cancel(true);
                 result2.cancel(true);
+                assertThrows(IllegalStateException.class, session::close);
             }
         } finally {
             rpc.close();

java/src/test/java/com/github/copilot/ZeroTimeoutContractTest.java

@@ -21,6 +21,21 @@
  */
 public class ZeroTimeoutContractTest {
 
+    @SuppressWarnings("unchecked")
+    @Test
+    void closeShouldPropagateDetachFailureAndRemainTerminal() {
+        var mockRpc = mock(JsonRpcClient.class);
+        when(mockRpc.invoke(eq("session.detach"), any(), any())).thenReturn(
+                CompletableFuture.completedFuture(new CopilotSession.SessionDetachResponse(false, "cleanup failed")));
+        var session = new CopilotSession("detach-failure-test", mockRpc);
+
+        var error = assertThrows(IllegalStateException.class, session::close);
+
+        assertTrue(error.getMessage().contains("cleanup failed"));
+        assertThrows(IllegalStateException.class, () -> session.send("test"));
+        assertDoesNotThrow(session::close);
+    }
+
     @SuppressWarnings("unchecked")
     @Test
     void sendAndWaitWithZeroTimeoutShouldNotTimeOut() throws Exception {
@@ -33,7 +48,7 @@ void sendAndWaitWithZeroTimeoutShouldNotTimeOut() throws Exception {
             Object method = invocation.getArgument(0);
             if ("session.detach".equals(method)) {
                 // Make session.close() non-blocking by completing detach immediately
-                return CompletableFuture.completedFuture(null);
+                return CompletableFuture.completedFuture(new CopilotSession.SessionDetachResponse(true, null));
             }
             // For other calls (e.g., message send), return an incomplete future so the
             // sendAndWait result does not complete due to a mock response.