diff --git a/dotnet/src/Client.cs b/dotnet/src/Client.cs
index 86178ba74..7b9161adb 100644
--- a/dotnet/src/Client.cs
+++ b/dotnet/src/Client.cs
@@ -663,7 +663,17 @@ or IOException
if (ctx.FfiHost is { } ffiHost)
{
- try { ffiHost.Dispose(); }
+ try
+ {
+ if (gracefulRuntimeShutdown)
+ {
+ ffiHost.Dispose();
+ }
+ else
+ {
+ ffiHost.ForceDispose();
+ }
+ }
catch (Exception ex) { AddCleanupError(errors, ex, _logger); }
_ffiHost = null;
}
diff --git a/dotnet/src/FfiRuntimeHost.cs b/dotnet/src/FfiRuntimeHost.cs
index a838b9fd1..be7470671 100644
--- a/dotnet/src/FfiRuntimeHost.cs
+++ b/dotnet/src/FfiRuntimeHost.cs
@@ -39,6 +39,7 @@ internal sealed partial class FfiRuntimeHost : IDisposable
{
/// Logical name the native interop layer binds the cdylib to.
private const string LibraryName = "copilot_runtime";
+ private const uint ConnectionDrainTimeoutMilliseconds = 30_000;
private readonly ILogger _logger;
private readonly string _cliEntrypoint;
@@ -53,6 +54,10 @@ internal sealed partial class FfiRuntimeHost : IDisposable
private uint _connectionId;
private bool _disposed;
+ // Roots this host while native code can invoke its outbound callback. It is
+ // released only after connection_close_and_wait confirms callbacks drained.
+ private GCHandle _selfHandle;
+
private FfiRuntimeHost(string libraryPath, string cliEntrypoint, IReadOnlyDictionary? environment, IReadOnlyList args, ILogger logger)
{
_libraryPath = libraryPath;
@@ -224,7 +229,13 @@ private void FeedInbound(IntPtr bytesPtr, UIntPtr bytesLen)
_receiveStream.Feed(buffer);
}
- public void Dispose()
+ public void Dispose() =>
+ Dispose(ConnectionDrainTimeoutMilliseconds, throwOnDrainFailure: true);
+
+ internal void ForceDispose() =>
+ Dispose(timeoutMilliseconds: 0, throwOnDrainFailure: false);
+
+ private void Dispose(uint timeoutMilliseconds, bool throwOnDrainFailure)
{
if (_disposed)
{
@@ -232,17 +243,31 @@ public void Dispose()
}
_disposed = true;
+ Exception? connectionCloseError = null;
+ var callbackDrained = _connectionId == 0;
try
{
if (_connectionId != 0)
{
- NativeConnectionClose(_connectionId);
+ callbackDrained = NativeConnectionCloseAndWait(
+ _connectionId,
+ timeoutMilliseconds);
+ if (!callbackDrained)
+ {
+ connectionCloseError = new TimeoutException(
+ $"FfiRuntimeHost timed out after {timeoutMilliseconds} ms "
+ + "waiting for runtime callbacks to drain; callback state was retained.");
+ }
_connectionId = 0;
}
}
catch (Exception ex)
{
- _logger.LogDebug(ex, "FfiRuntimeHost: connection_close failed");
+ connectionCloseError = new InvalidOperationException(
+ "FfiRuntimeHost failed to close and drain the runtime connection; "
+ + "callback state was retained.",
+ ex);
+ _connectionId = 0;
}
try
@@ -259,7 +284,21 @@ public void Dispose()
}
_receiveStream.Complete();
- DisposeNativeCallback();
+ if (callbackDrained)
+ {
+ DisposeNativeCallback();
+ }
+
+ if (connectionCloseError is not null)
+ {
+ if (throwOnDrainFailure)
+ {
+ throw connectionCloseError;
+ }
+ _logger.LogWarning(
+ connectionCloseError,
+ "FfiRuntimeHost force-closed before runtime callbacks drained; callback state was retained");
+ }
}
/// Length as the native pointer-sized unsigned integer the ABI expects.
@@ -272,10 +311,6 @@ public void Dispose()
private static bool s_resolverRegistered;
private static string? s_resolvedLibraryPath;
- // A normal (non-pinned) handle to this instance, passed to the native side as
- // the callback's user_data so the static outbound callback can route back here.
- private GCHandle _selfHandle;
-
///
/// Registers (once) a process-wide
/// that maps to the absolute runtime.node path so the
@@ -332,7 +367,8 @@ private uint NativeOpenConnection(uint serverId)
private static bool NativeConnectionWrite(uint connectionId, ReadOnlySpan frame) => ConnectionWrite(connectionId, frame, Len(frame.Length));
- private static bool NativeConnectionClose(uint connectionId) => ConnectionClose(connectionId);
+ private static bool NativeConnectionCloseAndWait(uint connectionId, uint timeoutMilliseconds) =>
+ ConnectionCloseAndWait(connectionId, timeoutMilliseconds);
private void DisposeNativeCallback()
{
@@ -366,7 +402,7 @@ private static partial uint HostStart(
[return: MarshalAs(UnmanagedType.U1)]
private static partial bool HostShutdown(uint serverId);
- [LibraryImport(LibraryName, EntryPoint = "copilot_runtime_connection_open")]
+ [LibraryImport(LibraryName, EntryPoint = "copilot_runtime_connection_open_tracked")]
[UnmanagedCallConv(CallConvs = new[] { typeof(CallConvCdecl) })]
private static unsafe partial uint ConnectionOpen(
uint serverId,
@@ -381,10 +417,10 @@ private static unsafe partial uint ConnectionOpen(
[return: MarshalAs(UnmanagedType.U1)]
private static partial bool ConnectionWrite(uint connectionId, ReadOnlySpan bytes, nuint bytesLen);
- [LibraryImport(LibraryName, EntryPoint = "copilot_runtime_connection_close")]
+ [LibraryImport(LibraryName, EntryPoint = "copilot_runtime_connection_close_and_wait")]
[UnmanagedCallConv(CallConvs = new[] { typeof(CallConvCdecl) })]
[return: MarshalAs(UnmanagedType.U1)]
- private static partial bool ConnectionClose(uint connectionId);
+ private static partial bool ConnectionCloseAndWait(uint connectionId, uint timeoutMilliseconds);
#else
// ---- Legacy interop: delegate-based P/Invoke for netstandard2.0 ----
// netstandard2.0 has neither LibraryImport, NativeLibrary, nor UnmanagedCallersOnly,
@@ -416,7 +452,7 @@ private delegate uint ConnectionOpenDelegate(
[UnmanagedFunctionPointer(CallingConvention.Cdecl)]
[return: MarshalAs(UnmanagedType.U1)]
- private delegate bool ConnectionCloseDelegate(uint connectionId);
+ private delegate bool ConnectionCloseAndWaitDelegate(uint connectionId, uint timeoutMilliseconds);
[UnmanagedFunctionPointer(CallingConvention.Cdecl)]
private delegate void OutboundCallbackDelegate(IntPtr userData, IntPtr bytesPtr, UIntPtr bytesLen);
@@ -428,7 +464,7 @@ private delegate uint ConnectionOpenDelegate(
private static HostShutdownDelegate? s_hostShutdown;
private static ConnectionOpenDelegate? s_connectionOpen;
private static ConnectionWriteDelegate? s_connectionWrite;
- private static ConnectionCloseDelegate? s_connectionClose;
+ private static ConnectionCloseAndWaitDelegate? s_connectionCloseAndWait;
// Held for the connection's lifetime so the marshaled function pointer handed to the
// native side is not collected while Rust may still invoke it.
@@ -457,9 +493,11 @@ private static void PrepareNativeLibrary(string libraryPath)
s_hostStart = Bind(handle, "copilot_runtime_host_start");
s_hostShutdown = Bind(handle, "copilot_runtime_host_shutdown");
- s_connectionOpen = Bind(handle, "copilot_runtime_connection_open");
+ s_connectionOpen = Bind(handle, "copilot_runtime_connection_open_tracked");
s_connectionWrite = Bind(handle, "copilot_runtime_connection_write");
- s_connectionClose = Bind(handle, "copilot_runtime_connection_close");
+ s_connectionCloseAndWait = Bind(
+ handle,
+ "copilot_runtime_connection_close_and_wait");
s_loaded = true;
s_loadedPath = libraryPath;
}
@@ -480,11 +518,12 @@ private static uint NativeHostStart(byte[] argvJson, byte[]? env) =>
private uint NativeOpenConnection(uint serverId)
{
+ _selfHandle = GCHandle.Alloc(this);
_outboundDelegate = OnOutbound;
return s_connectionOpen!(
serverId,
_outboundDelegate,
- IntPtr.Zero,
+ GCHandle.ToIntPtr(_selfHandle),
null, UIntPtr.Zero,
null, UIntPtr.Zero,
null, UIntPtr.Zero);
@@ -500,17 +539,28 @@ private static unsafe bool NativeConnectionWrite(uint connectionId, ReadOnlySpan
}
}
- private static bool NativeConnectionClose(uint connectionId) => s_connectionClose!(connectionId);
+ private static bool NativeConnectionCloseAndWait(uint connectionId, uint timeoutMilliseconds) =>
+ s_connectionCloseAndWait!(connectionId, timeoutMilliseconds);
- private void DisposeNativeCallback() => _outboundDelegate = null;
+ private void DisposeNativeCallback()
+ {
+ _outboundDelegate = null;
+ if (_selfHandle.IsAllocated)
+ {
+ _selfHandle.Free();
+ }
+ }
- private void OnOutbound(IntPtr userData, IntPtr bytesPtr, UIntPtr bytesLen)
+ private static void OnOutbound(IntPtr userData, IntPtr bytesPtr, UIntPtr bytesLen)
{
- if (bytesPtr == IntPtr.Zero || bytesLen == UIntPtr.Zero)
+ if (userData == IntPtr.Zero || bytesPtr == IntPtr.Zero || bytesLen == UIntPtr.Zero)
{
return;
}
- FeedInbound(bytesPtr, bytesLen);
+ if (GCHandle.FromIntPtr(userData).Target is FfiRuntimeHost self)
+ {
+ self.FeedInbound(bytesPtr, bytesLen);
+ }
}
///
diff --git a/dotnet/test/E2E/ClientE2ETests.cs b/dotnet/test/E2E/ClientE2ETests.cs
index b6bdfd90f..401dd24d4 100644
--- a/dotnet/test/E2E/ClientE2ETests.cs
+++ b/dotnet/test/E2E/ClientE2ETests.cs
@@ -64,6 +64,18 @@ public async Task Should_Start_And_Connect_Over_InProcess_Ffi()
}
}
+ [Fact]
+ public async Task Should_Force_Stop_Over_InProcess_Ffi()
+ {
+ using var client = new CopilotClient(new CopilotClientOptions
+ {
+ Connection = RuntimeConnection.ForInProcess(),
+ });
+
+ await client.StartAsync();
+ await client.ForceStopAsync();
+ }
+
[Theory]
[InlineData(true)] // stdio transport
[InlineData(false)] // TCP transport