Fix FFI close responsiveness and callback reclamation

java/sdk/src/main/java/com/github/copilot/ffi/JnaNativeBinding.java

@@ -49,9 +49,9 @@
  * <p>
  * The native runtime can still be inside an outbound callback when
  * {@link #connectionClose} returns {@code false}. Each JNA callback wrapper is
- * therefore retained for the lifetime of the JVM. After connection close
- * reports quiescence, its Java delegate is detached so the wrapper no longer
- * retains the complete host object graph.
+ * therefore retained until connection close reports quiescence. A successful
+ * close detaches its Java delegate and releases the wrapper for garbage
+ * collection.
  *
  * <h2>GraalVM Native Image</h2>
  * <p>
@@ -112,9 +112,9 @@ int copilot_runtime_connection_open(int serverId, OutboundCallback callback, Poi
     private static volatile CopilotRuntimeLibrary loadedLib;
 
     /**
-     * Process-lifetime roots for JNA callback trampolines. Native code can invoke a
-     * callback after connection and host teardown return, so entries are never
-     * removed in production.
+     * Roots for JNA callback trampolines until successful connection close confirms
+     * that native code can no longer invoke them. Host shutdown alone is not a
+     * callback-quiescence barrier.
      */
     private static final Set<OutboundCallback> RETAINED_CALLBACKS = ConcurrentHashMap.newKeySet();
 
@@ -141,7 +141,7 @@ int copilot_runtime_connection_open(int serverId, OutboundCallback callback, Poi
      * <p>
      * Registrations remain here through connection close because native callbacks
      * can still arrive while close reports non-quiescence. Successful connection
-     * close detaches their Java delegates; the wrappers themselves remain rooted by
+     * close detaches their Java delegates and removes their wrapper roots from
      * {@link #RETAINED_CALLBACKS}.
      */
     private final Map<Integer, CallbackRegistration> callbackRegistrations = new ConcurrentHashMap<>();
@@ -238,8 +238,10 @@ public int hostStart(byte[] argvJson, int argvJsonLen, byte[] envJson, int envJs
     public boolean hostShutdown(int serverId) {
         boolean shutdown = lib.copilot_runtime_host_shutdown(serverId) != 0;
         if (shutdown) {
-            callbackRegistrations.forEach((connectionId, registration) -> {
-                if (registration.serverId == serverId && callbackRegistrations.remove(connectionId, registration)) {
+            // Keep registrations so a later successful close can release their wrapper
+            // roots.
+            callbackRegistrations.values().forEach(registration -> {
+                if (registration.serverId == serverId) {
                     registration.detach();
                 }
             });
@@ -272,6 +274,7 @@ public boolean connectionClose(int connectionId) {
             CallbackRegistration registration = callbackRegistrations.remove(connectionId);
             if (registration != null) {
                 registration.detach();
+                RETAINED_CALLBACKS.remove(registration.wrapper);
             }
         }
         return closed;

java/sdk/src/test/java/com/github/copilot/ffi/JnaNativeBindingTest.java

@@ -7,10 +7,13 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assumptions.assumeTrue;
 
+import com.sun.jna.CallbackReference;
+import com.sun.jna.Function;
 import com.sun.jna.Pointer;
 
 import java.lang.ref.WeakReference;
@@ -22,6 +25,8 @@
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.io.TempDir;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
 
 /**
  * Unit tests for {@link JnaNativeBinding}.
@@ -52,6 +57,7 @@ private static class StubRuntimeLibrary implements JnaNativeBinding.CopilotRunti
         int connectionOpenReturn = 1;
         byte connectionWriteReturn = 1;
         byte connectionCloseReturn = 1;
+        RuntimeException connectionCloseFailure;
 
         byte[] lastArgvJson;
         int lastArgvJsonLen;
@@ -90,6 +96,9 @@ public byte copilot_runtime_connection_write(int connectionId, byte[] data, Size
         @Override
         public byte copilot_runtime_connection_close(int connectionId) {
             lastConnectionId = connectionId;
+            if (connectionCloseFailure != null) {
+                throw connectionCloseFailure;
+            }
             return connectionCloseReturn;
         }
     }
@@ -229,33 +238,60 @@ void activeCallbacksStartsAtZero() {
         assertEquals(0, binding.activeCallbacks.get(), "Active callback counter must start at zero");
     }
 
-    @Test
-    void callbackWrapperRemainsReachableAndIsDetachedAfterConnectionClose() throws InterruptedException {
+    @ParameterizedTest
+    @ValueSource(booleans = {false, true})
+    void callbackWrapperIsReleasedOnlyAfterSuccessfulConnectionClose(boolean shutdownBeforeClose)
+            throws InterruptedException {
         StubRuntimeLibrary stub = new StubRuntimeLibrary();
         stub.connectionOpenReturn = 99;
         JnaNativeBinding binding = new JnaNativeBinding(stub);
         AtomicInteger invocations = new AtomicInteger();
 
         WeakReference<OutboundCallback> callbackReference = openAndCloseConnection(binding, stub,
-                (userData, data, len) -> invocations.incrementAndGet());
+                (userData, data, len) -> invocations.incrementAndGet(), shutdownBeforeClose);
 
         awaitGarbageCollection(callbackReference);
 
-        OutboundCallback callback = callbackReference.get();
-        assertNotNull(callback, "Callback wrapper must remain strongly reachable after connection close");
-        callback.invoke(Pointer.NULL, Pointer.NULL, new SizeT(0));
-        assertEquals(0, invocations.get(), "Successful connection close must detach the Java delegate");
+        assertNull(callbackReference.get(), "Successful connection close must release the callback wrapper");
+        assertEquals(2, invocations.get(), "Failed closes must retain the delegate; successful close must detach it");
     }
 
     private static WeakReference<OutboundCallback> openAndCloseConnection(JnaNativeBinding binding,
-            StubRuntimeLibrary stub, OutboundCallback callback) {
+            StubRuntimeLibrary stub, OutboundCallback callback, boolean shutdownBeforeClose)
+            throws InterruptedException {
         int connectionId = binding.connectionOpen(1, callback, Pointer.NULL, null, 0, null, 0, null, 0);
         assertEquals(99, connectionId);
         assertNotNull(stub.lastCallback);
 
         WeakReference<OutboundCallback> callbackReference = new WeakReference<>(stub.lastCallback);
+        // Allocate a real JNA trampoline, whose weak references must not keep the
+        // wrapper alive once the native connection reports quiescence.
+        Function nativeCallback = Function.getFunction(CallbackReference.getFunctionPointer(stub.lastCallback));
         stub.lastCallback = null;
+
+        stub.connectionCloseReturn = 0;
+        assertFalse(binding.connectionClose(connectionId));
+        awaitGarbageCollection(callbackReference);
+        assertNotNull(callbackReference.get(), "Non-quiescent close must retain the callback wrapper");
+        nativeCallback.invokeVoid(new Object[]{Pointer.NULL, Pointer.NULL, new SizeT(0)});
+
+        stub.connectionCloseFailure = new IllegalStateException("close failed");
+        assertThrows(IllegalStateException.class, () -> binding.connectionClose(connectionId));
+        awaitGarbageCollection(callbackReference);
+        assertNotNull(callbackReference.get(), "Throwing close must retain the callback wrapper");
+        nativeCallback.invokeVoid(new Object[]{Pointer.NULL, Pointer.NULL, new SizeT(0)});
+        stub.connectionCloseFailure = null;
+
+        if (shutdownBeforeClose) {
+            assertTrue(binding.hostShutdown(1));
+            awaitGarbageCollection(callbackReference);
+            assertNotNull(callbackReference.get(), "Host shutdown alone must retain the callback wrapper");
+        }
+        OutboundCallback wrapper = callbackReference.get();
+        assertNotNull(wrapper);
+        stub.connectionCloseReturn = 1;
         assertTrue(binding.connectionClose(connectionId));
+        wrapper.invoke(Pointer.NULL, Pointer.NULL, new SizeT(0));
         return callbackReference;
     }
 

nodejs/src/client.ts

@@ -1200,7 +1200,7 @@ export class CopilotClient {
             const host = this.ffiHost;
             this.ffiHost = null;
             try {
-                host.dispose();
+                await host.dispose();
             } catch (error) {
                 errors.push(
                     new Error(
@@ -1314,12 +1314,13 @@ export class CopilotClient {
 
         // Tear down the in-process FFI host (if any).
         if (this.ffiHost) {
+            const host = this.ffiHost;
+            this.ffiHost = null;
             try {
-                this.ffiHost.dispose();
+                await host.dispose();
             } catch {
                 // Ignore errors during force stop
             }
-            this.ffiHost = null;
         }
 
         if (this.cliStartTimeout) {

nodejs/src/ffiRuntimeHost.ts

@@ -256,7 +256,7 @@ export class FfiRuntimeHost {
         } finally {
             this.starting = false;
             if (this.disposed) {
-                this.tryFinalizeCleanup();
+                void this.tryFinalizeCleanup();
             }
         }
     }
@@ -329,11 +329,11 @@ export class FfiRuntimeHost {
         }
         this.cleanupRetryTimer = setTimeout(() => {
             this.cleanupRetryTimer = undefined;
-            this.tryFinalizeCleanup();
+            void this.tryFinalizeCleanup();
         }, CLEANUP_RETRY_INTERVAL_MS);
     }
 
-    private tryFinalizeCleanup(): void {
+    private async tryFinalizeCleanup(): Promise<void> {
         if (this.cleanupInProgress) {
             this.scheduleCleanupRetry();
             return;
@@ -344,7 +344,20 @@ export class FfiRuntimeHost {
             if (this.connectionId) {
                 let closed = false;
                 try {
-                    closed = Boolean(this.lib.connectionClose(this.connectionId));
+                    // Close waits for outbound callbacks, which need the JS event loop
+                    // to run and return before native code can report quiescence.
+                    closed = await new Promise<boolean>((resolvePromise, rejectPromise) => {
+                        this.lib.connectionClose.async(
+                            this.connectionId,
+                            (error: Error | null, result: boolean) => {
+                                if (error) {
+                                    rejectPromise(error);
+                                } else {
+                                    resolvePromise(result);
+                                }
+                            }
+                        );
+                    });
                 } catch (error) {
                     console.error(
                         `Failed to close in-process FFI connection: ${error instanceof Error ? (error.stack ?? error.message) : String(error)}`
@@ -394,15 +407,15 @@ export class FfiRuntimeHost {
         }
     }
 
-    /** Closes the FFI connection, shuts down the native host, and releases resources. */
-    dispose(): void {
+    /** Awaits the initial cleanup attempt; a non-quiescent close is retried in the background. */
+    async dispose(): Promise<void> {
         if (this.disposed) {
             return;
         }
         this.disposed = true;
         this.receiveStream.end();
         if (!this.starting) {
-            this.tryFinalizeCleanup();
+            await this.tryFinalizeCleanup();
         }
     }
 }

nodejs/test/client.test.ts

@@ -4285,6 +4285,32 @@ describe("CopilotClient", () => {
     });
 
     describe("shutdown", () => {
+        it.each(["stop", "forceStop"] as const)(
+            "%s waits for the initial in-process cleanup attempt",
+            async (method) => {
+                const client = new CopilotClient({
+                    connection: RuntimeConnection.forInProcess(),
+                });
+                let finishCleanup!: () => void;
+                const cleanup = new Promise<void>((resolve) => {
+                    finishCleanup = resolve;
+                });
+                const dispose = vi.fn(() => cleanup);
+                (client as any).ffiHost = { dispose };
+
+                let stopped = false;
+                const shutdown = client[method]().then(() => {
+                    stopped = true;
+                });
+                await vi.waitFor(() => expect(dispose).toHaveBeenCalledTimes(1));
+                expect(stopped).toBe(false);
+
+                finishCleanup();
+                await shutdown;
+                expect(stopped).toBe(true);
+            }
+        );
+
         it("requests runtime shutdown when stopping an SDK-owned process", async () => {
             const client = new CopilotClient();
             const calls: string[] = [];

nodejs/test/e2e/inprocess_ffi.e2e.test.ts

@@ -2,6 +2,10 @@
  *  Copyright (c) Microsoft Corporation. All rights reserved.
  *--------------------------------------------------------------------------------------------*/
 
+import { mkdtemp, rm } from "node:fs/promises";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { setTimeout as delay } from "node:timers/promises";
 import { describe, expect, it } from "vitest";
 import { CopilotClient, RuntimeConnection } from "../../src/index.js";
 
@@ -22,4 +26,31 @@ describe("In-process FFI transport", () => {
 
         expect(await client.stop()).toHaveLength(0); // No errors on stop
     });
+
+    it("keeps the event loop responsive when force-stopping pending outbound traffic", async () => {
+        const baseDirectory = await mkdtemp(join(tmpdir(), "copilot-ffi-stop-"));
+        const client = new CopilotClient({
+            connection: RuntimeConnection.forInProcess(),
+            mode: "empty",
+            useLoggedInUser: false,
+            baseDirectory,
+        });
+        try {
+            await client.start();
+            const pending = Array.from({ length: 200 }, () => client.ping("x").catch(() => {}));
+            await delay(1);
+
+            const start = performance.now();
+            const timer = delay(10).then(() => performance.now() - start);
+            await client.forceStop();
+            const timerDelay = await timer;
+            await Promise.all(pending);
+
+            // Native close waits up to five seconds if it blocks callback delivery.
+            expect(timerDelay).toBeLessThan(2000);
+        } finally {
+            await client.forceStop();
+            await rm(baseDirectory, { recursive: true, force: true });
+        }
+    });
 });