Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 40 additions & 2 deletions dotnet/src/Client.cs
Original file line number Diff line number Diff line change
Expand Up @@ -463,6 +463,11 @@ async Task<Connection> StartCoreAsync(CancellationToken ct)
"CopilotClient.StartAsync protocol verification complete. Elapsed={Elapsed}",
startTimestamp);

if (_options.ExtensionLaunchProvider is not null)
{
await connection.Server.RegisterExtensionLaunchProviderAsync(ct);
Comment thread
stephentoub marked this conversation as resolved.
}

if (_builtinPluginDirectories.Length > 0)
{
var request = new BuiltinPluginDirectoriesRequest(_builtinPluginDirectories);
Expand Down Expand Up @@ -2041,8 +2046,7 @@ await Rpc.SessionFs.SetProviderAsync(

/// <summary>
/// Builds the client-global RPC handler bag at construction time. Registers
/// the LLM inference provider adapter and/or the GitHub telemetry adapter
/// depending on which options are configured. The GitHub token dispatcher is
/// the configured connection-level adapters. The GitHub token dispatcher is
/// always registered because providers are configured per session.
/// </summary>
private ClientGlobalApiHandlers? BuildClientGlobalApis()
Expand All @@ -2051,6 +2055,7 @@ await Rpc.SessionFs.SetProviderAsync(
var onGitHubTelemetry = _options.OnGitHubTelemetry;
return new ClientGlobalApiHandlers
{
ExtensionLaunchProvider = _options.ExtensionLaunchProvider,
LlmInference = handler is null ? null : new LlmInferenceAdapter(handler, () => _serverRpc),
GitHubTelemetry = onGitHubTelemetry is null ? null : new GitHubTelemetryAdapter(onGitHubTelemetry, _logger),
GitHubToken = new GitHubTokenAdapter(this),
Expand Down Expand Up @@ -2698,6 +2703,10 @@ private async Task<Connection> ConnectToServerAsync(Process? cliProcess, string?
{
ClientGlobalApiRegistration.RegisterClientGlobalApiHandlers(rpc, _clientGlobalApis);
}
if (cliProcess is not null)
{
RegisterRpcProcessExit(cliProcess, rpc);
}
rpc.StartListening();
_ = CancelExternalToolsWhenConnectionClosesAsync(rpc);
LoggingHelpers.LogTiming(_logger, LogLevel.Debug, null,
Expand Down Expand Up @@ -2729,6 +2738,35 @@ private async Task<Connection> ConnectToServerAsync(Process? cliProcess, string?
}
}

private void RegisterRpcProcessExit(Process cliProcess, JsonRpc rpc)
{
try
{
cliProcess.EnableRaisingEvents = true;
cliProcess.Exited += (_, _) => DisposeRpcAfterProcessExit(rpc);
if (cliProcess.HasExited)
{
DisposeRpcAfterProcessExit(rpc);
}
}
catch (Exception ex) when (ex is InvalidOperationException or ObjectDisposedException)
{
_logger.LogDebug(ex, "Unable to monitor the Copilot CLI process for transport closure");
}
}

private void DisposeRpcAfterProcessExit(JsonRpc rpc)
{
try
{
rpc.Dispose(new ConnectionLostException());
}
catch (Exception ex) when (IsRecoverableConnectionCleanupFailure(ex))
{
_logger.LogDebug(ex, "Failed to dispose JSON-RPC connection after Copilot CLI process exit");
}
}

private static bool IsRecoverableConnectionCleanupFailure(Exception exception)
=> exception is not OutOfMemoryException
and not StackOverflowException
Expand Down
59 changes: 36 additions & 23 deletions dotnet/src/JsonRpc.cs
Original file line number Diff line number Diff line change
Expand Up @@ -39,10 +39,11 @@
private readonly ConcurrentDictionary<long, PendingRequest> _pendingRequests = new();
private readonly ConcurrentDictionary<string, MethodRegistration> _methods = new();
private readonly TaskCompletionSource _completionSource = new(TaskCreationOptions.RunContinuationsAsynchronously);
private readonly SemaphoreSlim _writeLock = new(1, 1);

Check notice

Code scanning / CodeQL

Missed 'using' opportunity Note

This variable is manually
disposed
in a
finally block
- consider a C# using statement as a preferable resource management technique.
private readonly CancellationTokenSource _disposeCts = new();
private long _nextId;
private bool _disposed;
private int _disposeStarted;
private Exception? _terminalError;

/// <summary>
/// Initializes a new <see cref="JsonRpc"/>.
Expand Down Expand Up @@ -96,6 +97,11 @@
CancellationTokenRegistration cancelRegistration = default;
try
{
if (Volatile.Read(ref _terminalError) is { } terminalError)
{
throw terminalError;
}

if (cancellationToken.CanBeCanceled)
{
cancelRegistration = cancellationToken.Register(static state =>
Expand Down Expand Up @@ -136,6 +142,11 @@
LogInvokeTiming(LogLevel.Debug, ex, method, id, "Canceled", timingTimestamp);
throw;
}
catch (ObjectDisposedException ex) when (Volatile.Read(ref _terminalError) is ConnectionLostException)
{
LogInvokeTiming(LogLevel.Warning, ex, method, id, "Failed", timingTimestamp);
throw new ConnectionLostException();
}
catch (Exception ex)
{
LogInvokeTiming(LogLevel.Warning, ex, method, id, "Failed", timingTimestamp);
Expand Down Expand Up @@ -183,27 +194,25 @@
}

/// <inheritdoc />
public void Dispose()
public void Dispose() => Dispose(new ObjectDisposedException(nameof(JsonRpc)));

internal void Dispose(Exception reason)
{
if (_disposed)
if (Interlocked.Exchange(ref _disposeStarted, 1) != 0)
{
return;
}

_disposed = true;
_disposeCts.Cancel();

// Fail all pending requests
foreach (var kvp in _pendingRequests)
FailPendingRequests(reason);
try
{
if (_pendingRequests.TryRemove(kvp.Key, out var pending))
{
pending.TrySetException(new ObjectDisposedException(nameof(JsonRpc)));
}
_disposeCts.Cancel();
}
finally
{
_completionSource.TrySetResult();
_writeLock.Dispose();
}

_completionSource.TrySetResult();
_writeLock.Dispose();
}

private async Task SendMessageAsync<T>(T message, JsonTypeInfo<T> typeInfo, CancellationToken cancellationToken)
Expand Down Expand Up @@ -338,17 +347,21 @@
}
finally
{
// Fail all pending requests
foreach (var kvp in _pendingRequests)
FailPendingRequests(new ConnectionLostException());
_completionSource.TrySetResult();
}
}

private void FailPendingRequests(Exception reason)
{
var terminalError = Interlocked.CompareExchange(ref _terminalError, reason, null) ?? reason;
foreach (var kvp in _pendingRequests)
{
if (_pendingRequests.TryRemove(kvp.Key, out var pending))
{
if (_pendingRequests.TryRemove(kvp.Key, out var pending))
{
pending.TrySetException(new ConnectionLostException());
}
pending.TrySetException(terminalError);
}

_completionSource.TrySetResult();
}

Check notice

Code scanning / CodeQL

Missed opportunity to use Where Note

This foreach loop
implicitly filters its target sequence
- consider filtering the sequence explicitly using '.Where(...)'.
}

/// <summary>
Expand Down
9 changes: 9 additions & 0 deletions dotnet/src/Types.cs
Original file line number Diff line number Diff line change
Expand Up @@ -319,6 +319,7 @@ private CopilotClientOptions(CopilotClientOptions? other)
OnListModels = other.OnListModels;
SessionFs = other.SessionFs;
RequestHandler = other.RequestHandler;
ExtensionLaunchProvider = other.ExtensionLaunchProvider;
OnGitHubTelemetry = other.OnGitHubTelemetry;
SessionIdleTimeoutSeconds = other.SessionIdleTimeoutSeconds;
EnableRemoteSessions = other.EnableRemoteSessions;
Expand Down Expand Up @@ -433,6 +434,14 @@ private CopilotClientOptions(CopilotClientOptions? other)
[Experimental(Diagnostics.Experimental)]
public CopilotRequestHandler? RequestHandler { get; set; }

/// <summary>
/// Connection-level extension launch profile provider.
/// When set, the SDK registers the provider during <c>StartAsync()</c>
/// before any session can be created.
/// </summary>
[Experimental(Diagnostics.Experimental)]
public IExtensionLaunchProviderHandler? ExtensionLaunchProvider { get; set; }
Comment thread
stephentoub marked this conversation as resolved.
Comment thread
stephentoub marked this conversation as resolved.

/// <summary>
/// Experimental. Receives GitHub telemetry events the runtime forwards to this
/// connection; setting a handler opts created/resumed sessions into forwarding.
Expand Down
62 changes: 0 additions & 62 deletions dotnet/test/E2E/ExternalToolCancellationE2ETests.cs

This file was deleted.

Loading
Loading