Add from https://github.com/github/copilot-sdk-java/pull/233

java/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<String, CopilotSession> getSessionsMap(CopilotClient client) throws Exception {
+        Field f = CopilotClient.class.getDeclaredField("sessions");
+        f.setAccessible(true);
+        return (Map<String, CopilotSession>) 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<String, CopilotSession> 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<String, CopilotSession> 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<String, CopilotSession> 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();
+        }
+    }
+}