diff --git a/dotnet/src/Client.cs b/dotnet/src/Client.cs index cbc7d5fa2d..3f9176bfc3 100644 --- a/dotnet/src/Client.cs +++ b/dotnet/src/Client.cs @@ -365,7 +365,56 @@ private static bool IsFullyQualifiedPath(string path) /// public Task StartAsync(CancellationToken cancellationToken = default) { - return _connectionTask ??= StartCoreAsync(cancellationToken); + var connectionTask = Volatile.Read(ref _connectionTask); + if (connectionTask is not null) + { + return connectionTask; + } + + var completion = new TaskCompletionSource( + TaskCreationOptions.RunContinuationsAsynchronously); + // Publish before startup begins because StartCoreAsync executes synchronously + // until its first incomplete await and can re-enter user-provided callbacks. + connectionTask = Interlocked.CompareExchange( + ref _connectionTask, + completion.Task, + null); + if (connectionTask is not null) + { + return connectionTask; + } + + _ = CompleteStartAsync(completion, cancellationToken); + return completion.Task; + + async Task CompleteStartAsync( + TaskCompletionSource startCompletion, + CancellationToken ct) + { + try + { + var connection = await StartCoreAsync(ct).ConfigureAwait(false); + startCompletion.TrySetResult(connection); + } + catch (OperationCanceledException) + { + // Clear before waking waiters so continuations can immediately retry. + _ = Interlocked.CompareExchange( + ref _connectionTask, + null, + startCompletion.Task); + startCompletion.TrySetCanceled(ct); + } + catch (Exception ex) + { + // Clear before waking waiters so continuations can immediately retry. + _ = Interlocked.CompareExchange( + ref _connectionTask, + null, + startCompletion.Task); + startCompletion.TrySetException(ex); + } + } async Task StartCoreAsync(CancellationToken ct) { @@ -375,6 +424,7 @@ async Task StartCoreAsync(CancellationToken ct) Connection? connection = null; Process? cliProcess = null; ProcessStderrPump? stderrPump = null; + FfiRuntimeHost? ffiHost = null; try { @@ -427,7 +477,7 @@ async Task StartCoreAsync(CancellationToken ct) ?? throw new InvalidOperationException( $"In-process FFI runtime library not found at '{searchedRuntime}'.") : ResolveRuntimePathForExplicitCli(explicitCliPath); - var ffiHost = FfiRuntimeHost.Create( + ffiHost = FfiRuntimeHost.Create( ffiRuntimePath, explicitCliPath, ffiEnvironment, @@ -503,6 +553,12 @@ await InvokeRpcAsync( { await CleanupCliProcessAsync(cliProcess, stderrPump, errors: null, _logger); } + else if (ffiHost is not null) + { + try { ffiHost.Dispose(); } + catch (Exception cleanupError) { AddCleanupError(null, cleanupError, _logger); } + _ffiHost = null; + } if (ex is IOException && cliProcess is not null @@ -627,14 +683,12 @@ private static void ThrowErrors(List? errors) private async Task CleanupConnectionAsync(List? errors, bool gracefulRuntimeShutdown) { - var connectionTask = _connectionTask; + var connectionTask = Interlocked.Exchange(ref _connectionTask, null); if (connectionTask is null) { return; } - _connectionTask = null; - Connection ctx; try { @@ -2741,7 +2795,7 @@ private async Task CancelExternalToolsWhenConnectionClosesAsync(JsonRpc rpc) _logger.LogDebug(exception, "JSON-RPC connection completed with an error"); } - var connectionTask = _connectionTask; + var connectionTask = Volatile.Read(ref _connectionTask); if (connectionTask is null || connectionTask.Status != System.Threading.Tasks.TaskStatus.RanToCompletion || !ReferenceEquals(connectionTask.Result.Rpc, rpc)) diff --git a/dotnet/test/Unit/ClientSessionLifetimeTests.cs b/dotnet/test/Unit/ClientSessionLifetimeTests.cs index 173569788c..75b33f0454 100644 --- a/dotnet/test/Unit/ClientSessionLifetimeTests.cs +++ b/dotnet/test/Unit/ClientSessionLifetimeTests.cs @@ -13,6 +13,7 @@ using System.Text.Json; using GitHub.Copilot.Rpc; using Microsoft.Extensions.AI; +using Microsoft.Extensions.Logging; using Xunit; namespace GitHub.Copilot.Test.Unit; @@ -21,6 +22,66 @@ public sealed class ClientSessionLifetimeTests { private sealed record RpcRequestRecord(string Method, JsonElement Params); + [Fact] + public async Task StartAsync_Concurrent_Callers_Share_One_Startup() + { + await using var server = await FakeCopilotServer.StartAsync(); + var logger = new BlockingStartLogger(); + await using var client = new CopilotClient(new CopilotClientOptions + { + Connection = RuntimeConnection.ForUri(server.Url), + Logger = logger + }); + using var cancellation = new CancellationTokenSource(); + + var firstInvocation = Task.Run(() => client.StartAsync(cancellation.Token)); + await logger.FirstStartEntered.WaitAsync(TimeSpan.FromSeconds(5)); + + Task secondStart; + int startCount; + try + { + // Cancel before releasing the blocked first attempt so both code paths + // terminate without opening a connection if startup is duplicated. + cancellation.Cancel(); + secondStart = client.StartAsync(cancellation.Token); + startCount = logger.StartCount; + } + finally + { + logger.ReleaseFirstStart(); + } + + var firstStart = await firstInvocation; + await Assert.ThrowsAnyAsync(() => firstStart); + await Assert.ThrowsAnyAsync(() => secondStart); + + Assert.Equal(1, startCount); + Assert.Same(firstStart, secondStart); + + await client.StartAsync(); + Assert.Equal(2, logger.StartCount); + } + + [Fact] + public async Task StartAsync_Retries_After_Failed_Startup() + { + await using var server = await FakeCopilotServer.StartAsync(); + var logger = new FailFirstStartLogger(); + await using var client = new CopilotClient(new CopilotClientOptions + { + Connection = RuntimeConnection.ForUri(server.Url), + Logger = logger + }); + + var error = await Assert.ThrowsAsync(() => client.StartAsync()); + Assert.Equal("first startup failed", error.Message); + + await client.StartAsync(); + + Assert.Equal(2, logger.StartCount); + } + [Theory] [InlineData("static")] [InlineData("")] @@ -2183,6 +2244,70 @@ private static Process StartExitedProcess() return process; } + private sealed class BlockingStartLogger : ILogger + { + private readonly TaskCompletionSource _firstStartEntered = + new(TaskCreationOptions.RunContinuationsAsynchronously); + private readonly TaskCompletionSource _allowFirstStart = + new(TaskCreationOptions.RunContinuationsAsynchronously); + private int _startCount; + + public Task FirstStartEntered => _firstStartEntered.Task; + + public int StartCount => Volatile.Read(ref _startCount); + + public IDisposable? BeginScope(TState state) where TState : notnull => null; + + public bool IsEnabled(LogLevel logLevel) => logLevel == LogLevel.Debug; + + public void Log( + LogLevel logLevel, + EventId eventId, + TState state, + Exception? exception, + Func formatter) + { + if (formatter(state, exception) != "Starting Copilot client") + { + return; + } + + var startCount = Interlocked.Increment(ref _startCount); + if (startCount == 1) + { + _firstStartEntered.TrySetResult(); + _allowFirstStart.Task.GetAwaiter().GetResult(); + } + } + + public void ReleaseFirstStart() => _allowFirstStart.TrySetResult(); + } + + private sealed class FailFirstStartLogger : ILogger + { + private int _startCount; + + public int StartCount => Volatile.Read(ref _startCount); + + public IDisposable? BeginScope(TState state) where TState : notnull => null; + + public bool IsEnabled(LogLevel logLevel) => logLevel == LogLevel.Debug; + + public void Log( + LogLevel logLevel, + EventId eventId, + TState state, + Exception? exception, + Func formatter) + { + if (formatter(state, exception) == "Starting Copilot client" + && Interlocked.Increment(ref _startCount) == 1) + { + throw new InvalidOperationException("first startup failed"); + } + } + } + private sealed class FakeCopilotServer : IAsyncDisposable { private readonly TcpListener _listener; diff --git a/nodejs/src/client.ts b/nodejs/src/client.ts index eb92cf0bed..a25a74121a 100644 --- a/nodejs/src/client.ts +++ b/nodejs/src/client.ts @@ -448,6 +448,8 @@ export class CopilotClient { private runtimePort: number | null = null; private actualHost: string = "localhost"; private state: "disconnected" | "connecting" | "connected" | "error" = "disconnected"; + /** Shared in-flight start; concurrent callers await it instead of spawning another CLI. */ + private startPromise: Promise | null = null; private sessions: Map = new Map(); private stderrBuffer: string = ""; // Captures CLI stderr for error messages /** Resolved connection mode chosen in the constructor. */ @@ -935,6 +937,20 @@ export class CopilotClient { return; } + // Concurrent callers share one in-progress start instead of each spawning a CLI. + if (this.startPromise) { + return this.startPromise; + } + + this.startPromise = this.doStart(); + try { + await this.startPromise; + } finally { + this.startPromise = null; + } + } + + private async doStart(): Promise { this.forceStopping = false; this.connectionClosed = false; this.processTransportError = null; diff --git a/nodejs/test/client.test.ts b/nodejs/test/client.test.ts index 3db96ea47b..119ed40a10 100644 --- a/nodejs/test/client.test.ts +++ b/nodejs/test/client.test.ts @@ -62,6 +62,64 @@ describe("approveAll", () => { }); describe("CopilotClient", () => { + it("start() is single-flight: concurrent callers share one startup", async () => { + const client = new CopilotClient({ autoStart: false }); + onTestFinished(() => client.forceStop()); + + // Stub the underlying startup (doStart) that the single-flight guard + // dedupes. Transport-independent: this is the same regardless of the + // stdio vs in-process connection path. The delay makes all three + // start() calls overlap; on success it marks the client connected like + // the real doStart does. + const doStart = vi.fn().mockImplementation( + () => + new Promise((resolve) => + setTimeout(() => { + (client as any).state = "connected"; + resolve(); + }, 50) + ) + ); + (client as any).doStart = doStart; + + // Before the fix, each concurrent caller ran startup (and spawned its own + // CLI, orphaning all but the last). With single-flight they share one. + await Promise.all([client.start(), client.start(), client.start()]); + + expect(doStart).toHaveBeenCalledTimes(1); + expect((client as any).state).toBe("connected"); + + // Once connected, a further start() is a no-op (no extra startup). + await client.start(); + expect(doStart).toHaveBeenCalledTimes(1); + }); + + it("start() retries after a failed attempt (single-flight guard is cleared)", async () => { + const client = new CopilotClient({ autoStart: false }); + onTestFinished(() => client.forceStop()); + + // Stub the underlying startup: fail once, then succeed. Transport- + // independent (does not depend on the stdio vs in-process path). + const doStart = vi + .fn() + .mockImplementationOnce(async () => { + (client as any).state = "error"; + throw new Error("boom"); + }) + .mockImplementationOnce(async () => { + (client as any).state = "connected"; + }); + (client as any).doStart = doStart; + + await expect(client.start()).rejects.toThrow(/boom/); + expect((client as any).state).toBe("error"); + + // The guard must have cleared so a later start() can retry. + await client.start(); + expect(doStart).toHaveBeenCalledTimes(2); + expect((client as any).state).toBe("connected"); + }); + it.each([ { source: "connection path", diff --git a/python/copilot/client.py b/python/copilot/client.py index dd531a2be9..4af6da09ea 100644 --- a/python/copilot/client.py +++ b/python/copilot/client.py @@ -1753,6 +1753,7 @@ def __init__( self._cli_process: subprocess.Popen | None = None self._client: JsonRpcClient | None = None self._state: _ConnectionState = "disconnected" + self._start_lock = asyncio.Lock() self._sessions: dict[str, CopilotSession] = {} self._sessions_lock = threading.Lock() self._github_token_providers: dict[str, _GitHubTokenProviderRegistration] = {} @@ -1938,12 +1939,15 @@ async def start(self) -> None: >>> await client.start() >>> # Now ready to create sessions """ - if self._state == "connected": - return + async with self._start_lock: + if self._state == "connected": + return + + await self._start_once() + async def _start_once(self) -> None: start_time = time.perf_counter() self._state = "connecting" - try: # Only start CLI server process if not connecting to external server if not self._is_external_server: @@ -1969,14 +1973,10 @@ async def start(self) -> None: if self._options.builtin_plugin_directories: assert self._client is not None - try: - await self._client.request( - "plugins.builtin.set", - {"paths": list(self._options.builtin_plugin_directories)}, - ) - except Exception: - await self.force_stop() - raise + await self._client.request( + "plugins.builtin.set", + {"paths": list(self._options.builtin_plugin_directories)}, + ) if self._session_fs_config: session_fs_start = time.perf_counter() @@ -2000,6 +2000,7 @@ async def start(self) -> None: ) except ProcessExitedError as e: # Process exited with error - reraise as RuntimeError with stderr + await self._cleanup_failed_start() self._state = "error" log_timing( logger, @@ -2009,16 +2010,13 @@ async def start(self) -> None: exc_info=True, ) raise RuntimeError(str(e)) from None - except Exception as e: + except asyncio.CancelledError: + await self._cleanup_failed_start() self._state = "error" - log_timing( - logger, - logging.WARNING, - "CopilotClient.start failed", - start_time, - exc_info=True, - ) + raise + except Exception as e: # Check if process exited and capture any remaining stderr + startup_error: RuntimeError | None = None process = self._cli_process if self._cli_process is not None else self._process if process and hasattr(process, "poll"): if isinstance(e, BrokenPipeError) and process.poll() is None: @@ -2028,9 +2026,24 @@ async def start(self) -> None: pass return_code = process.poll() if return_code is not None and self._client: - raise RuntimeError(self._client._get_process_exit_error()) from e + startup_error = RuntimeError(self._client._get_process_exit_error()) + + await self._cleanup_failed_start() + self._state = "error" + log_timing( + logger, + logging.WARNING, + "CopilotClient.start failed", + start_time, + exc_info=True, + ) + if startup_error is not None: + raise startup_error from e raise + async def _cleanup_failed_start(self) -> None: + await self.force_stop() + async def stop(self) -> None: """ Stop the CLI server and close all active sessions. diff --git a/python/e2e/test_client_e2e.py b/python/e2e/test_client_e2e.py index 1e8ea82e55..a6d41af9f7 100644 --- a/python/e2e/test_client_e2e.py +++ b/python/e2e/test_client_e2e.py @@ -206,15 +206,18 @@ async def test_should_report_error_with_stderr_when_cli_fails_to_start(self): f"Expected error to contain 'nonexistent', got: {error_message}" ) - # Verify subsequent calls also fail (don't hang) - with pytest.raises(Exception) as exc_info2: - session = await client.create_session( - on_permission_request=PermissionHandler.approve_all - ) - await session.send("test") - # Error message varies by platform (EINVAL on Windows, EPIPE on Linux) - error_msg = str(exc_info2.value).lower() - assert "invalid" in error_msg or "pipe" in error_msg or "closed" in error_msg + # A subsequent lazy start retries and reports the same invalid CLI + # configuration rather than using the failed process transport. + with pytest.raises(RuntimeError) as retry_exc_info: + await client.create_session(on_permission_request=PermissionHandler.approve_all) + + retry_error_message = str(retry_exc_info.value) + assert "stderr" in retry_error_message, ( + f"Expected retry error to contain 'stderr', got: {retry_error_message}" + ) + assert "nonexistent" in retry_error_message, ( + f"Expected retry error to contain 'nonexistent', got: {retry_error_message}" + ) finally: await client.force_stop() diff --git a/python/test_client.py b/python/test_client.py index 2e3868ef1c..84e9ed2133 100644 --- a/python/test_client.py +++ b/python/test_client.py @@ -127,6 +127,109 @@ def test_relative_path_is_rejected(self): ) +class TestClientStart: + @staticmethod + def _create_client() -> CopilotClient: + client = CopilotClient(connection=RuntimeConnection.for_uri("localhost:1234")) + client._verify_protocol_version = AsyncMock() + return client + + @pytest.mark.asyncio + async def test_concurrent_callers_share_one_startup(self): + client = self._create_client() + connect_started = asyncio.Event() + allow_connect = asyncio.Event() + + async def connect(): + connect_started.set() + await allow_connect.wait() + + client._connect_to_server = AsyncMock(side_effect=connect) + + first_start = asyncio.create_task(client.start()) + await connect_started.wait() + second_start = asyncio.create_task(client.start()) + await asyncio.sleep(0) + allow_connect.set() + + try: + await asyncio.gather(first_start, second_start) + assert client._state == "connected" + finally: + await client.force_stop() + + client._connect_to_server.assert_awaited_once() + client._verify_protocol_version.assert_awaited_once() + + @pytest.mark.asyncio + async def test_failed_start_cleans_up_before_retry(self): + client = self._create_client() + failed_transport = Mock() + failed_transport.poll.return_value = None + attempts = 0 + + async def connect(): + nonlocal attempts + attempts += 1 + if attempts == 1: + client._process = failed_transport + raise RuntimeError("first startup failed") + + client._connect_to_server = AsyncMock(side_effect=connect) + + with pytest.raises(RuntimeError, match="first startup failed"): + await client.start() + + assert client._process is None + assert client._state == "error" + failed_transport.terminate.assert_called_once() + + try: + await client.start() + assert client._state == "connected" + finally: + await client.force_stop() + + assert attempts == 2 + assert client._state == "disconnected" + + @pytest.mark.asyncio + async def test_cancelled_start_cleans_up_before_retry(self): + client = self._create_client() + cancelled_transport = Mock() + connect_started = asyncio.Event() + attempts = 0 + + async def connect(): + nonlocal attempts + attempts += 1 + if attempts == 1: + client._process = cancelled_transport + connect_started.set() + await asyncio.Future() + + client._connect_to_server = AsyncMock(side_effect=connect) + + first_start = asyncio.create_task(client.start()) + await connect_started.wait() + first_start.cancel() + with pytest.raises(asyncio.CancelledError): + await first_start + + assert client._process is None + assert client._state == "error" + cancelled_transport.terminate.assert_called_once() + + try: + await client.start() + assert client._state == "connected" + finally: + await client.force_stop() + + assert attempts == 2 + assert client._state == "disconnected" + + class TestClientShutdown: @pytest.mark.asyncio async def test_stop_requests_runtime_shutdown_for_owned_process(self):