diff --git a/java/sdk/src/main/java/com/github/copilot/ffi/JnaNativeBinding.java b/java/sdk/src/main/java/com/github/copilot/ffi/JnaNativeBinding.java
index bf6c912f5f..001e8d0760 100644
--- a/java/sdk/src/main/java/com/github/copilot/ffi/JnaNativeBinding.java
+++ b/java/sdk/src/main/java/com/github/copilot/ffi/JnaNativeBinding.java
@@ -49,9 +49,9 @@
*
* 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.
*
*
GraalVM Native Image
*
@@ -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 RETAINED_CALLBACKS = ConcurrentHashMap.newKeySet();
@@ -141,7 +141,7 @@ int copilot_runtime_connection_open(int serverId, OutboundCallback callback, Poi
*
* 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 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;
diff --git a/java/sdk/src/test/java/com/github/copilot/ffi/JnaNativeBindingTest.java b/java/sdk/src/test/java/com/github/copilot/ffi/JnaNativeBindingTest.java
index f0ec519d47..2c6a0111d7 100644
--- a/java/sdk/src/test/java/com/github/copilot/ffi/JnaNativeBindingTest.java
+++ b/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 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 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 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;
}
diff --git a/nodejs/src/client.ts b/nodejs/src/client.ts
index 035f0f5e60..6e4b4fb5b4 100644
--- a/nodejs/src/client.ts
+++ b/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) {
diff --git a/nodejs/src/ffiRuntimeHost.ts b/nodejs/src/ffiRuntimeHost.ts
index ee71c44717..f2db8eb163 100644
--- a/nodejs/src/ffiRuntimeHost.ts
+++ b/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 {
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((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 {
if (this.disposed) {
return;
}
this.disposed = true;
this.receiveStream.end();
if (!this.starting) {
- this.tryFinalizeCleanup();
+ await this.tryFinalizeCleanup();
}
}
}
diff --git a/nodejs/test/client.test.ts b/nodejs/test/client.test.ts
index 119ed40a10..24adb5cb6e 100644
--- a/nodejs/test/client.test.ts
+++ b/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((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[] = [];
diff --git a/nodejs/test/e2e/inprocess_ffi.e2e.test.ts b/nodejs/test/e2e/inprocess_ffi.e2e.test.ts
index e3b5f75ee4..e047440118 100644
--- a/nodejs/test/e2e/inprocess_ffi.e2e.test.ts
+++ b/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 });
+ }
+ });
});
diff --git a/nodejs/test/ffiRuntimeHost.test.ts b/nodejs/test/ffiRuntimeHost.test.ts
index d538b45385..e0fb3d71af 100644
--- a/nodejs/test/ffiRuntimeHost.test.ts
+++ b/nodejs/test/ffiRuntimeHost.test.ts
@@ -20,7 +20,11 @@ const ffi = vi.hoisted(() => {
const hostShutdown = vi.fn(() => true);
const connectionOpen = vi.fn(() => 21);
const connectionWrite = vi.fn(() => true);
- const connectionClose = vi.fn<() => boolean>();
+ const connectionClose = Object.assign(vi.fn<() => boolean>(), {
+ async: vi.fn<
+ (connectionId: number, callback: (error: Error | null, result: boolean) => void) => void
+ >(),
+ });
const register = vi.fn(
(callback: (userData: unknown, bytesPtr: unknown, bytesLen: number) => void) => {
registeredCallback = callback;
@@ -73,6 +77,9 @@ describe("FfiRuntimeHost callback cleanup", () => {
beforeEach(() => {
vi.useFakeTimers();
ffi.connectionClose.mockReset();
+ ffi.connectionClose.async
+ .mockReset()
+ .mockImplementation((_id, callback) => callback(null, true));
ffi.connectionOpen.mockClear();
ffi.hostShutdown.mockClear();
ffi.hostStart.mockClear();
@@ -82,18 +89,19 @@ describe("FfiRuntimeHost callback cleanup", () => {
});
afterEach(() => {
+ expect(ffi.connectionClose).not.toHaveBeenCalled();
vi.clearAllTimers();
vi.useRealTimers();
});
- it("retains the callback and keepalive until a retry closes the connection", async () => {
- ffi.connectionClose.mockReturnValueOnce(false).mockReturnValueOnce(true);
+ it("finishes disposal after a false close while retaining resources for the detached retry", async () => {
+ ffi.connectionClose.async.mockImplementationOnce((_id, callback) => callback(null, false));
const host = FfiRuntimeHost.create("runtime.node", undefined, undefined, []);
await host.start();
- host.dispose();
+ await host.dispose();
- expect(ffi.connectionClose).toHaveBeenCalledTimes(1);
+ expect(ffi.connectionClose.async).toHaveBeenCalledTimes(1);
expect(ffi.unregister).not.toHaveBeenCalled();
expect(ffi.hostShutdown).not.toHaveBeenCalled();
expect((host as any).outboundCallback).toBe(ffi.callbackToken);
@@ -101,25 +109,60 @@ describe("FfiRuntimeHost callback cleanup", () => {
await vi.advanceTimersByTimeAsync(100);
- expect(ffi.connectionClose).toHaveBeenCalledTimes(2);
+ expect(ffi.connectionClose.async).toHaveBeenCalledTimes(2);
expect(ffi.unregister).toHaveBeenCalledTimes(1);
expect(ffi.hostShutdown).toHaveBeenCalledTimes(1);
expect((host as any).outboundCallback).toBeUndefined();
expect((host as any).keepAliveTimer).toBeUndefined();
expect(vi.getTimerCount()).toBe(0);
- host.dispose();
+ await host.dispose();
await vi.advanceTimersByTimeAsync(100);
- expect(ffi.connectionClose).toHaveBeenCalledTimes(2);
+ expect(ffi.connectionClose.async).toHaveBeenCalledTimes(2);
expect(ffi.unregister).toHaveBeenCalledTimes(1);
expect(ffi.hostShutdown).toHaveBeenCalledTimes(1);
});
+ it("waits for successful initial cleanup without overlapping close calls", async () => {
+ let finishClose!: (error: Error | null, result: boolean) => void;
+ ffi.connectionClose.async.mockImplementationOnce((_id, callback) => {
+ finishClose = callback;
+ });
+ const host = FfiRuntimeHost.create("runtime.node", undefined, undefined, []);
+ await host.start();
+
+ let disposed = false;
+ const cleanup = host.dispose().then(() => {
+ disposed = true;
+ });
+ await host.dispose();
+ await vi.advanceTimersByTimeAsync(1000);
+
+ expect(disposed).toBe(false);
+ expect(ffi.connectionClose.async).toHaveBeenCalledTimes(1);
+ expect(ffi.unregister).not.toHaveBeenCalled();
+ expect(ffi.hostShutdown).not.toHaveBeenCalled();
+ expect((host as any).outboundCallback).toBe(ffi.callbackToken);
+ expect((host as any).keepAliveTimer).toBeDefined();
+
+ finishClose(null, true);
+ await cleanup;
+
+ expect(disposed).toBe(true);
+ expect(ffi.unregister).toHaveBeenCalledTimes(1);
+ expect(ffi.hostShutdown).toHaveBeenCalledTimes(1);
+ expect(vi.getTimerCount()).toBe(0);
+ });
+
it("defers reclamation when dispose is called from the outbound callback", async () => {
- ffi.connectionClose.mockReturnValueOnce(false).mockReturnValueOnce(true);
+ ffi.connectionClose.async.mockImplementationOnce((_id, callback) => {
+ setImmediate(() => callback(null, true));
+ });
const host = FfiRuntimeHost.create("runtime.node", undefined, undefined, []);
await host.start();
- vi.spyOn(host as any, "feedInbound").mockImplementation(() => host.dispose());
+ host.receiveStream.once("data", () => {
+ void host.dispose();
+ });
ffi.getRegisteredCallback()?.(null, {}, 1);
@@ -133,17 +176,44 @@ describe("FfiRuntimeHost callback cleanup", () => {
expect(vi.getTimerCount()).toBe(0);
});
+ it.each(["callback", "throw"])(
+ "quarantines the callback on a close %s error",
+ async (failure) => {
+ const closeError = new Error("close failed");
+ ffi.connectionClose.async.mockImplementationOnce((_id, callback) => {
+ if (failure === "throw") {
+ throw closeError;
+ }
+ callback(closeError, false);
+ });
+ const error = vi.spyOn(console, "error").mockImplementation(() => {});
+ const host = FfiRuntimeHost.create("runtime.node", undefined, undefined, []);
+ await host.start();
+
+ await host.dispose();
+ await vi.advanceTimersByTimeAsync(500);
+
+ expect(ffi.connectionClose.async).toHaveBeenCalledTimes(1);
+ expect(ffi.unregister).not.toHaveBeenCalled();
+ expect(ffi.hostShutdown).not.toHaveBeenCalled();
+ expect((host as any).outboundCallback).toBe(ffi.callbackToken);
+ expect((FfiRuntimeHost as any).quarantinedHosts.has(host)).toBe(true);
+ expect(error).toHaveBeenCalledTimes(1);
+ expect(vi.getTimerCount()).toBe(0);
+ error.mockRestore();
+ }
+ );
+
it("does not retry a terminal host shutdown failure", async () => {
- ffi.connectionClose.mockReturnValue(true);
ffi.hostShutdown.mockReturnValueOnce(false);
const error = vi.spyOn(console, "error").mockImplementation(() => {});
const host = FfiRuntimeHost.create("runtime.node", undefined, undefined, []);
await host.start();
- host.dispose();
+ await host.dispose();
await vi.advanceTimersByTimeAsync(500);
- expect(ffi.connectionClose).toHaveBeenCalledTimes(1);
+ expect(ffi.connectionClose.async).toHaveBeenCalledTimes(1);
expect(ffi.unregister).toHaveBeenCalledTimes(1);
expect(ffi.hostShutdown).toHaveBeenCalledTimes(1);
expect(vi.getTimerCount()).toBe(0);
@@ -151,7 +221,6 @@ describe("FfiRuntimeHost callback cleanup", () => {
});
it("retains the callback token when Koffi unregistration fails", async () => {
- ffi.connectionClose.mockReturnValue(true);
ffi.unregister.mockImplementationOnce(() => {
throw new Error("unregister failed");
});
@@ -159,10 +228,10 @@ describe("FfiRuntimeHost callback cleanup", () => {
const host = FfiRuntimeHost.create("runtime.node", undefined, undefined, []);
await host.start();
- host.dispose();
+ await host.dispose();
await vi.advanceTimersByTimeAsync(500);
- expect(ffi.connectionClose).toHaveBeenCalledTimes(1);
+ expect(ffi.connectionClose.async).toHaveBeenCalledTimes(1);
expect(ffi.unregister).toHaveBeenCalledTimes(1);
expect(ffi.hostShutdown).toHaveBeenCalledTimes(1);
expect((host as any).outboundCallback).toBe(ffi.callbackToken);
@@ -170,7 +239,7 @@ describe("FfiRuntimeHost callback cleanup", () => {
expect((host as any).keepAliveTimer).toBeUndefined();
expect(vi.getTimerCount()).toBe(0);
- host.dispose();
+ await host.dispose();
expect(ffi.unregister).toHaveBeenCalledTimes(1);
error.mockRestore();
});
diff --git a/rust/src/ffi.rs b/rust/src/ffi.rs
index 02fce3f030..450f50ae1c 100644
--- a/rust/src/ffi.rs
+++ b/rust/src/ffi.rs
@@ -710,7 +710,7 @@ mod tests {
shared.close();
- assert_eq!(TEST_CLOSE_CALLS.load(Ordering::SeqCst), 1);
+ assert!(TEST_CLOSE_CALLS.load(Ordering::SeqCst) >= 1);
assert_eq!(TEST_SHUTDOWN_CALLS.load(Ordering::SeqCst), 0);
assert_eq!(shared.connection_id.load(Ordering::SeqCst), 0);
assert!(shared.callback_state.load(Ordering::SeqCst).is_null());