From 00c194ebb62c4f83f6628bd1fc08ff98127b14f7 Mon Sep 17 00:00:00 2001 From: Ed Burns Date: Wed, 27 May 2026 16:22:10 -0700 Subject: [PATCH 1/3] Fix session registry leak when server re-keys session ID When the server returns a different sessionId and a subsequent step (e.g. updateSessionOptionsForMode) fails, the exceptionally blocks only removed the original sessionId from the sessions map, leaving the re-keyed entry behind as a stale/closed session. Fix: remove both the original and the active (possibly re-keyed) session ID in createSession/resumeSession exceptionally blocks, and unregister the session from the registry in updateSessionOptionsForMode before closing it. Addresses review comments on copilot-sdk PR #1473. --- src/main/java/com/github/copilot/CopilotClient.java | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/src/main/java/com/github/copilot/CopilotClient.java b/src/main/java/com/github/copilot/CopilotClient.java index 0379b97e6..662b66c7b 100644 --- a/src/main/java/com/github/copilot/CopilotClient.java +++ b/src/main/java/com/github/copilot/CopilotClient.java @@ -515,6 +515,11 @@ public CompletableFuture createSession(SessionConfig config) { }); }).exceptionally(ex -> { sessions.remove(sessionId); + // Also remove the re-keyed entry if the server returned a different ID + String activeId = session.getSessionId(); + if (!sessionId.equals(activeId)) { + sessions.remove(activeId); + } LoggingHelpers.logTiming(LOG, Level.WARNING, ex, "CopilotClient.createSession failed. Elapsed={Elapsed}, SessionId=" + sessionId, totalNanos); @@ -622,6 +627,11 @@ public CompletableFuture resumeSession(String sessionId, ResumeS }); }).exceptionally(ex -> { sessions.remove(sessionId); + // Also remove the re-keyed entry if the server returned a different ID + String activeId = session.getSessionId(); + if (!sessionId.equals(activeId)) { + sessions.remove(activeId); + } LoggingHelpers.logTiming(LOG, Level.WARNING, ex, "CopilotClient.resumeSession failed. Elapsed={Elapsed}, SessionId=" + sessionId, totalNanos); @@ -739,6 +749,7 @@ CompletableFuture updateSessionOptionsForMode(CopilotSession session, Bool // Best-effort disconnect so we don't leak it (in empty mode it would // otherwise stay alive with permissive defaults). LOG.log(Level.WARNING, "session.options.update failed for session " + session.getSessionId(), ex); + sessions.remove(session.getSessionId()); try { session.close(); } catch (Exception closeEx) { From 4002f806cb4500a1c0a51903559183715f5cff1d Mon Sep 17 00:00:00 2001 From: Ed Burns Date: Wed, 27 May 2026 16:23:34 -0700 Subject: [PATCH 2/3] Add more to comparison --- scripts/compare-standalone-to-monorepo.sh | 42 +++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/scripts/compare-standalone-to-monorepo.sh b/scripts/compare-standalone-to-monorepo.sh index a7c6a9ce7..ad329fdb1 100755 --- a/scripts/compare-standalone-to-monorepo.sh +++ b/scripts/compare-standalone-to-monorepo.sh @@ -118,6 +118,35 @@ while IFS= read -r relpath; do fi done < "$TMPFILE_MONO" +# ── Compare scripts/codegen/java.ts ─────────────────────────────────── +JAVATS_STATUS="" +STANDALONE_JAVATS="${STANDALONE}/scripts/codegen/java.ts" +MONO_JAVATS="${MONO_JAVA}/scripts/codegen/java.ts" + +if [ -f "$STANDALONE_JAVATS" ] && [ -f "$MONO_JAVATS" ]; then + if diff -q "$STANDALONE_JAVATS" "$MONO_JAVATS" >/dev/null 2>&1; then + JAVATS_STATUS="identical" + SAME_COUNT=$((SAME_COUNT + 1)) + else + JAVATS_STATUS="differ" + DIFFER_COUNT=$((DIFFER_COUNT + 1)) + DIFFER_LIST="${DIFFER_LIST}scripts/codegen/java.ts +" + fi +elif [ -f "$STANDALONE_JAVATS" ] && [ ! -f "$MONO_JAVATS" ]; then + JAVATS_STATUS="only-standalone" + MISSING_FROM_MONO_COUNT=$((MISSING_FROM_MONO_COUNT + 1)) + MISSING_FROM_MONO_LIST="${MISSING_FROM_MONO_LIST}scripts/codegen/java.ts +" +elif [ ! -f "$STANDALONE_JAVATS" ] && [ -f "$MONO_JAVATS" ]; then + JAVATS_STATUS="only-monorepo" + MISSING_FROM_STANDALONE_COUNT=$((MISSING_FROM_STANDALONE_COUNT + 1)) + MISSING_FROM_STANDALONE_LIST="${MISSING_FROM_STANDALONE_LIST}scripts/codegen/java.ts +" +else + JAVATS_STATUS="missing-both" +fi + # ── Compare .lastmerge ──────────────────────────────────────────────── LASTMERGE_STATUS="" STANDALONE_LASTMERGE="${STANDALONE}/.lastmerge" @@ -172,6 +201,19 @@ elif [ "$LASTMERGE_STATUS" = "only-monorepo" ]; then else echo ".lastmerge: not found in either location" fi + +# scripts/codegen/java.ts status +if [ "$JAVATS_STATUS" = "identical" ]; then + echo "scripts/codegen/java.ts: identical" +elif [ "$JAVATS_STATUS" = "differ" ]; then + echo "scripts/codegen/java.ts: DIFFERS" +elif [ "$JAVATS_STATUS" = "only-standalone" ]; then + echo "scripts/codegen/java.ts: only in standalone" +elif [ "$JAVATS_STATUS" = "only-monorepo" ]; then + echo "scripts/codegen/java.ts: only in monorepo" +else + echo "scripts/codegen/java.ts: not found in either location" +fi echo "" if [ "$DIFFER_COUNT" -gt 0 ]; then From d82e8b0786c7c705ae330cda6da9ef5ad799ea15 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 27 May 2026 23:41:35 +0000 Subject: [PATCH 3/3] Add CreateSessionReKeyEntryTest to cover session-map re-key cleanup paths Co-authored-by: edburns <75821+edburns@users.noreply.github.com> --- .../copilot/CreateSessionReKeyEntryTest.java | 262 ++++++++++++++++++ 1 file changed, 262 insertions(+) create mode 100644 src/test/java/com/github/copilot/CreateSessionReKeyEntryTest.java diff --git a/src/test/java/com/github/copilot/CreateSessionReKeyEntryTest.java b/src/test/java/com/github/copilot/CreateSessionReKeyEntryTest.java new file mode 100644 index 000000000..94f4edbdf --- /dev/null +++ b/src/test/java/com/github/copilot/CreateSessionReKeyEntryTest.java @@ -0,0 +1,262 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +package com.github.copilot; + +import static org.junit.jupiter.api.Assertions.*; + +import java.io.OutputStream; +import java.lang.reflect.Field; +import java.net.ServerSocket; +import java.net.Socket; +import java.nio.charset.StandardCharsets; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; + +import org.junit.jupiter.api.Test; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import com.github.copilot.rpc.CopilotClientOptions; +import com.github.copilot.rpc.PermissionHandler; +import com.github.copilot.rpc.SessionConfig; + +/** + * Tests for the session-map re-key cleanup paths in CopilotClient when the + * server returns a different session ID than the client-supplied one. + */ +class CreateSessionReKeyEntryTest { + + private static final ObjectMapper MAPPER = JsonRpcClient.getObjectMapper(); + + /** + * A connected socket pair where the server replies to "session.create" with a + * configurable sessionId and then replies to "session.options.update" with + * success or failure. + */ + private static final class ReKeyServer implements AutoCloseable { + + final Socket clientSocket; + final Socket serverSocket; + final JsonRpcClient rpcClient; + private volatile boolean running = true; + private final Thread replyThread; + + /** The sessionId to return in the session.create response. */ + private final String returnedSessionId; + /** If true, the session.options.update call will fail. */ + private final boolean failOptionsUpdate; + + ReKeyServer(String returnedSessionId, boolean failOptionsUpdate) throws Exception { + this.returnedSessionId = returnedSessionId; + this.failOptionsUpdate = failOptionsUpdate; + + try (var ss = new ServerSocket(0)) { + clientSocket = new Socket("localhost", ss.getLocalPort()); + serverSocket = ss.accept(); + } + serverSocket.setSoTimeout(5000); + rpcClient = JsonRpcClient.fromSocket(clientSocket); + + replyThread = new Thread(() -> { + try { + var in = serverSocket.getInputStream(); + var out = serverSocket.getOutputStream(); + while (running) { + // Read Content-Length header + var header = new StringBuilder(); + int b; + while ((b = in.read()) != -1) { + if (b == '\n' && header.toString().endsWith("\r")) { + break; + } + header.append((char) b); + } + if (b == -1) + break; + // Skip blank line + in.read(); // '\r' + in.read(); // '\n' + + String hdr = header.toString().trim(); + int colon = hdr.indexOf(':'); + int len = Integer.parseInt(hdr.substring(colon + 1).trim()); + byte[] body = in.readNBytes(len); + JsonNode msg = MAPPER.readTree(body); + + String method = msg.get("method").asText(); + long id = msg.get("id").asLong(); + + if ("session.create".equals(method)) { + // Return a response with the (possibly different) session ID + ObjectNode result = MAPPER.createObjectNode(); + result.put("sessionId", returnedSessionId); + String response = MAPPER.writeValueAsString(MAPPER.createObjectNode().put("jsonrpc", "2.0") + .put("id", id).set("result", result)); + sendRpcMessage(out, response); + } else if ("session.options.update".equals(method)) { + if (failOptionsUpdate) { + // Send an error response + ObjectNode error = MAPPER.createObjectNode(); + error.put("code", -32000); + error.put("message", "simulated options update failure"); + String response = MAPPER.writeValueAsString(MAPPER.createObjectNode() + .put("jsonrpc", "2.0").put("id", id).set("error", error)); + sendRpcMessage(out, response); + } else { + // Send a success response + String response = MAPPER.writeValueAsString( + MAPPER.createObjectNode().put("jsonrpc", "2.0").put("id", id).set("result", + MAPPER.createObjectNode().put("success", true))); + sendRpcMessage(out, response); + } + } else { + // Generic success for anything else + String response = MAPPER.writeValueAsString(MAPPER.createObjectNode().put("jsonrpc", "2.0") + .put("id", id).set("result", MAPPER.createObjectNode().put("success", true))); + sendRpcMessage(out, response); + } + } + } catch (Exception e) { + if (running) { + // Ignore expected exceptions on shutdown + } + } + }); + replyThread.setDaemon(true); + replyThread.start(); + } + + private static void sendRpcMessage(OutputStream out, String json) throws Exception { + byte[] bytes = json.getBytes(StandardCharsets.UTF_8); + String header = "Content-Length: " + bytes.length + "\r\n\r\n"; + out.write(header.getBytes(StandardCharsets.UTF_8)); + out.write(bytes); + out.flush(); + } + + @Override + public void close() throws Exception { + running = false; + rpcClient.close(); + clientSocket.close(); + serverSocket.close(); + replyThread.join(3000); + } + } + + @SuppressWarnings("unchecked") + private static Map getSessionsMap(CopilotClient client) throws Exception { + Field f = CopilotClient.class.getDeclaredField("sessions"); + f.setAccessible(true); + return (Map) f.get(client); + } + + private static void injectConnection(CopilotClient client, JsonRpcClient rpc) throws Exception { + // Build a Connection record via the private record constructor + Class connClass = null; + for (Class c : CopilotClient.class.getDeclaredClasses()) { + if (c.getSimpleName().equals("Connection")) { + connClass = c; + break; + } + } + assertNotNull(connClass, "Could not find Connection inner class"); + + var ctor = connClass.getDeclaredConstructors()[0]; + ctor.setAccessible(true); + // Connection(JsonRpcClient rpc, Process process, ServerRpc serverRpc) + Object connection = ctor.newInstance(rpc, null, null); + + Field f = CopilotClient.class.getDeclaredField("connectionFuture"); + f.setAccessible(true); + f.set(client, CompletableFuture.completedFuture(connection)); + } + + @Test + void createSessionReKeyEntry_successfulReKey_removesOldKeyAndAddsNewKey() throws Exception { + String clientSessionId = "client-supplied-id"; + String serverSessionId = "server-returned-id"; + + try (var server = new ReKeyServer(serverSessionId, false)) { + var client = new CopilotClient(new CopilotClientOptions().setAutoStart(false)); + injectConnection(client, server.rpcClient); + + var config = new SessionConfig().setSessionId(clientSessionId) + .setOnPermissionRequest(PermissionHandler.APPROVE_ALL); + + CopilotSession session = client.createSession(config).get(); + + Map sessions = getSessionsMap(client); + + // The old client-supplied key should be removed + assertNull(sessions.get(clientSessionId), + "Old client-supplied sessionId should be removed from sessions map after re-key"); + // The new server-returned key should be present + assertSame(session, sessions.get(serverSessionId), + "Server-returned sessionId should be the key in sessions map"); + // The session object should report the server-returned ID + assertEquals(serverSessionId, session.getSessionId(), + "Session should report the server-returned sessionId"); + + client.close(); + } + } + + @Test + void createSessionReKeyEntry_failureAfterReKey_removesBothKeys() throws Exception { + String clientSessionId = "client-supplied-id"; + String serverSessionId = "server-returned-id"; + + try (var server = new ReKeyServer(serverSessionId, true)) { + var client = new CopilotClient(new CopilotClientOptions().setAutoStart(false)); + injectConnection(client, server.rpcClient); + + // Set skipCustomInstructions so that session.options.update is actually invoked + var config = new SessionConfig().setSessionId(clientSessionId) + .setOnPermissionRequest(PermissionHandler.APPROVE_ALL).setSkipCustomInstructions(true); + + // The session.options.update will fail, triggering the exceptionally handler + ExecutionException ex = assertThrows(ExecutionException.class, () -> client.createSession(config).get()); + assertNotNull(ex.getCause()); + + Map sessions = getSessionsMap(client); + + // Both the original and re-keyed entries should be cleaned up + assertNull(sessions.get(clientSessionId), + "Original client-supplied sessionId should be removed on failure"); + assertNull(sessions.get(serverSessionId), + "Re-keyed server-returned sessionId should be removed on failure"); + assertTrue(sessions.isEmpty(), "Sessions map should be empty after failed create with re-key"); + + client.close(); + } + } + + @Test + void createSessionReKeyEntry_noReKey_sameIdKept() throws Exception { + String sessionId = "same-id-for-both"; + + try (var server = new ReKeyServer(sessionId, false)) { + var client = new CopilotClient(new CopilotClientOptions().setAutoStart(false)); + injectConnection(client, server.rpcClient); + + var config = new SessionConfig().setSessionId(sessionId) + .setOnPermissionRequest(PermissionHandler.APPROVE_ALL); + + CopilotSession session = client.createSession(config).get(); + + Map sessions = getSessionsMap(client); + + // When IDs match, the session stays under the original key + assertSame(session, sessions.get(sessionId), + "Session should remain under original key when server returns same ID"); + assertEquals(1, sessions.size(), "Should have exactly one entry in sessions map"); + + client.close(); + } + } +}