diff --git a/dotnet/src/Client.cs b/dotnet/src/Client.cs index 2c830e0d5d..0d7b937b88 100644 --- a/dotnet/src/Client.cs +++ b/dotnet/src/Client.cs @@ -557,6 +557,7 @@ await InvokeRpcAsync( public async Task StopAsync() { List errors = []; + CancelPendingExternalTools(); foreach (var session in _sessions.Values.ToArray()) { @@ -602,10 +603,7 @@ public async Task StopAsync() /// public async Task ForceStopAsync() { - foreach (var session in _sessions.Values) - { - session.CancelPendingExternalTools(); - } + CancelPendingExternalTools(); _sessions.Clear(); ClearGitHubTokenProviders(); @@ -2788,6 +2786,15 @@ private async Task CancelExternalToolsWhenConnectionClosesAsync(JsonRpc rpc) { return; } + CancelPendingExternalTools(); + } + + private void CancelPendingExternalTools() + { + if (_clientGlobalApis?.LlmInference is LlmInferenceAdapter llmInferenceAdapter) + { + llmInferenceAdapter.CancelPending(); + } foreach (var session in _sessions.Values) { session.CancelPendingExternalTools(); diff --git a/dotnet/src/CopilotRequestHandler.cs b/dotnet/src/CopilotRequestHandler.cs index 514d77da6f..cafc08ea94 100644 --- a/dotnet/src/CopilotRequestHandler.cs +++ b/dotnet/src/CopilotRequestHandler.cs @@ -8,6 +8,7 @@ using System.Diagnostics.CodeAnalysis; using System.Net.WebSockets; using System.Runtime.CompilerServices; +using System.Runtime.ExceptionServices; using System.Text; using System.Threading.Channels; @@ -462,6 +463,7 @@ private static Uri ToWebSocketUri(string url) [Experimental(Diagnostics.Experimental)] public class CopilotRequestHandler { + private const int HttpResponseReadAheadSize = 32 * 1024; private static readonly HttpClient s_sharedHttpClient = new(); private readonly HttpClient _httpClient; @@ -554,23 +556,45 @@ private static async Task BuildHttpRequestAsync(LlmInference private static async Task StreamResponseAsync(HttpResponseMessage response, LlmInferenceExchange exchange) { - await exchange.StartResponseAsync( - (int)response.StatusCode, - response.ReasonPhrase, - HeadersToMultiMap(response)).ConfigureAwait(false); - var ct = exchange.Context.CancellationToken; - using var stream = await response.Content.ReadAsStreamAsync(ct).ConfigureAwait(false); - var buffer = new byte[16 * 1024]; + await AwaitRpcAsync( + exchange.StartResponseAsync( + (int)response.StatusCode, + response.ReasonPhrase, + HeadersToMultiMap(response)), + ct).ConfigureAwait(false); + + var stream = await response.Content.ReadAsStreamAsync(ct).ConfigureAwait(false); + await using var reader = new BoundedHttpResponseReader(stream, HttpResponseReadAheadSize, ct); + var buffer = new byte[HttpResponseReadAheadSize]; int read; - while ((read = await stream.ReadAsync(buffer.AsMemory(), ct).ConfigureAwait(false)) > 0) + while ((read = await reader.ReadChunkAsync(buffer, ct).ConfigureAwait(false)) > 0) { - await exchange.WriteResponseAsync(new ReadOnlyMemory(buffer, 0, read)).ConfigureAwait(false); + await AwaitRpcAsync( + exchange.WriteResponseAsync(new ReadOnlyMemory(buffer, 0, read)), + ct).ConfigureAwait(false); } await exchange.EndResponseAsync().ConfigureAwait(false); } + private static async Task AwaitRpcAsync(Task rpcTask, CancellationToken cancellationToken) + { + try + { + await rpcTask.WaitAsync(Timeout.InfiniteTimeSpan, cancellationToken).ConfigureAwait(false); + } + catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) + { + _ = rpcTask.ContinueWith( + static completed => _ = completed.Exception, + CancellationToken.None, + TaskContinuationOptions.ExecuteSynchronously | TaskContinuationOptions.OnlyOnFaulted, + TaskScheduler.Default); + throw; + } + } + private async Task HandleWebSocketAsync(LlmInferenceExchange exchange) { var ctx = exchange.Context; @@ -657,6 +681,191 @@ private static Dictionary> HeadersToMultiMap(HttpR return result; } + + /// + /// Reads an HTTP response into a fixed-size ring while the current response + /// chunk waits for its runtime acknowledgement. The producer reserves ring + /// space before every read, so committed bytes plus an in-flight read never + /// exceed the configured read-ahead bound. + /// + private sealed class BoundedHttpResponseReader : IAsyncDisposable + { + private readonly Stream _stream; + private readonly byte[] _buffer; + private readonly CancellationTokenSource _disposeCts; + private readonly object _gate = new(); + private TaskCompletionSource _changed = CreateSignal(); + private readonly Task _pump; + + private long _head; + private long _committedTail; + private long _reservedTail; + private bool _completed; + private ExceptionDispatchInfo? _error; + + internal BoundedHttpResponseReader(Stream stream, int capacity, CancellationToken cancellationToken) + { + _stream = stream; + _buffer = new byte[capacity]; + _disposeCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken); + _pump = PumpAsync(_disposeCts.Token); + } + + internal async Task ReadChunkAsync(byte[] destination, CancellationToken cancellationToken) + { + while (true) + { + Task waitTask; + ExceptionDispatchInfo? error; + lock (_gate) + { + var count = checked((int)(_committedTail - _head)); + if (count > 0) + { + var headIndex = (int)(_head % _buffer.Length); + var firstCount = Math.Min(count, _buffer.Length - headIndex); + Buffer.BlockCopy(_buffer, headIndex, destination, 0, firstCount); + if (firstCount < count) + { + Buffer.BlockCopy(_buffer, 0, destination, firstCount, count - firstCount); + } + + _head += count; + PulseLocked(); + return count; + } + + error = _error; + if (error is null) + { + if (_completed) + { + cancellationToken.ThrowIfCancellationRequested(); + return 0; + } + + waitTask = _changed.Task; + } + else + { + waitTask = Task.CompletedTask; + } + } + + if (error is not null) + { + error.Throw(); + } + + await waitTask.WaitAsync(Timeout.InfiniteTimeSpan, cancellationToken).ConfigureAwait(false); + } + } + + public async ValueTask DisposeAsync() + { + _disposeCts.Cancel(); + _stream.Dispose(); + try + { + await _pump.ConfigureAwait(false); + } + catch (OperationCanceledException) when (_disposeCts.IsCancellationRequested) + { + // Cancellation is the expected result of disposing an active pump. + } + finally + { + _disposeCts.Dispose(); + } + } + + private async Task PumpAsync(CancellationToken cancellationToken) + { + while (true) + { + Task? waitTask = null; + long reservationStart = 0; + int reservationLength = 0; + int reservationIndex = 0; + + lock (_gate) + { + var used = checked((int)(_reservedTail - _head)); + if (used == _buffer.Length) + { + waitTask = _changed.Task; + } + else + { + reservationStart = _reservedTail; + reservationIndex = (int)(reservationStart % _buffer.Length); + reservationLength = Math.Min(_buffer.Length - used, _buffer.Length - reservationIndex); + _reservedTail += reservationLength; + } + } + + if (waitTask is not null) + { + await waitTask.WaitAsync(Timeout.InfiniteTimeSpan, cancellationToken).ConfigureAwait(false); + continue; + } + + int read; + try + { + read = await _stream.ReadAsync( + _buffer.AsMemory(reservationIndex, reservationLength), + cancellationToken).ConfigureAwait(false); + } + catch (Exception ex) + { + lock (_gate) + { + _reservedTail = reservationStart; + if (cancellationToken.IsCancellationRequested) + { + _completed = true; + } + else + { + _error = ExceptionDispatchInfo.Capture(ex); + } + + PulseLocked(); + } + + return; + } + + lock (_gate) + { + _committedTail = reservationStart + read; + _reservedTail = _committedTail; + if (read == 0) + { + _completed = true; + } + + PulseLocked(); + } + + if (read == 0) + { + return; + } + } + } + + private void PulseLocked() + { + var changed = _changed; + _changed = CreateSignal(); + changed.TrySetResult(true); + } + + private static TaskCompletionSource CreateSignal() => + new(TaskCreationOptions.RunContinuationsAsynchronously); + } } /// @@ -917,6 +1126,17 @@ public Task HttpRequestChunkAsync(LlmInferen return Task.FromResult(new LlmInferenceHttpRequestChunkResult()); } + internal void CancelPending() + { + foreach (var (requestId, exchange) in _pending) + { + if (_pending.TryRemove(requestId, out _)) + { + exchange.PushCancel("RPC connection closed"); + } + } + } + private async Task RunAsync(LlmInferenceExchange exchange) { try diff --git a/dotnet/test/E2E/ScenarioTestingEventSubscriptionsE2ETests.cs b/dotnet/test/E2E/ScenarioTestingEventSubscriptionsE2ETests.cs index c6e4649ac3..933d3b782f 100644 --- a/dotnet/test/E2E/ScenarioTestingEventSubscriptionsE2ETests.cs +++ b/dotnet/test/E2E/ScenarioTestingEventSubscriptionsE2ETests.cs @@ -141,7 +141,13 @@ await TestHelper.WaitForConditionAsync( }); var newEvents = new List(); - using var newSubscription = secondSession.On(newEvents.Add); + using var newSubscription = secondSession.On(evt => + { + lock (newEvents) + { + newEvents.Add(evt); + } + }); var newInfo = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); using var infoSubscription = secondSession.On(evt => { @@ -154,7 +160,10 @@ await TestHelper.WaitForConditionAsync( await newInfo.Task.WaitAsync(EventTimeout); Assert.Equal(countAfterClose, Volatile.Read(ref oldEventCount)); - Assert.Contains(newEvents, evt => evt is SessionInfoEvent info && info.Data.Message == "SCENARIO_EVENT_SOURCE_TWO"); + lock (newEvents) + { + Assert.Contains(newEvents, evt => evt is SessionInfoEvent info && info.Data.Message == "SCENARIO_EVENT_SOURCE_TWO"); + } } private static bool IsEventChannelClosed(CopilotSession session) diff --git a/dotnet/test/Unit/CopilotRequestHandlerProtocolTests.cs b/dotnet/test/Unit/CopilotRequestHandlerProtocolTests.cs new file mode 100644 index 0000000000..ddeaf7e341 --- /dev/null +++ b/dotnet/test/Unit/CopilotRequestHandlerProtocolTests.cs @@ -0,0 +1,651 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +#if NET8_0_OR_GREATER +using System.Net; +using System.Net.Sockets; +using System.Text; +using System.Text.Json; +using System.Threading.Channels; +using Xunit; + +namespace GitHub.Copilot.Test.Unit; + +public sealed class CopilotRequestHandlerProtocolTests +{ + private static readonly TimeSpan s_timeout = TimeSpan.FromSeconds(5); + + [Fact] + public async Task HttpResponse_ReadsAheadAndCoalescesUnderOneOutstandingDataRpc() + { + var source = new ControlledResponseStream(); + await using var peer = await ProtocolPeer.StartAsync(new StreamRequestHandler(source)); + + await source.PushAsync("first"u8.ToArray()); + await peer.BeginRequestAsync(); + var head = await peer.NextMethodAsync(); + Assert.Equal("llmInference.httpResponseStart", Method(head)); + await peer.AcknowledgeAsync(head); + + var first = await peer.NextMethodAsync(); + Assert.Equal("first"u8.ToArray(), Data(first)); + + var expectedReadAhead = new byte[32 * 1024]; + for (var i = 0; i < 32; i++) + { + var fragment = Enumerable.Repeat((byte)i, 1024).ToArray(); + fragment.CopyTo(expectedReadAhead, i * fragment.Length); + await source.PushAsync(fragment); + } + + await source.WaitForBytesAsync(5 + (32 * 1024)); + await source.PushAsync("last"u8.ToArray()); + await Task.Delay(TimeSpan.FromMilliseconds(100)); + Assert.Equal(5 + (32 * 1024), source.BytesRead); + await peer.AssertNoMethodAsync(); + + await peer.AcknowledgeAsync(first); + var combined = await peer.NextMethodAsync(); + Assert.Equal(expectedReadAhead, Data(combined)); + + await source.CompleteAsync(); + await peer.AcknowledgeAsync(combined); + var last = await peer.NextMethodAsync(); + Assert.Equal("last"u8.ToArray(), Data(last)); + await peer.AcknowledgeAsync(last); + + var end = await peer.NextMethodAsync(); + Assert.True(end.GetProperty("params").GetProperty("end").GetBoolean()); + Assert.False(end.GetProperty("params").TryGetProperty("error", out _)); + await peer.AcknowledgeAsync(end); + await source.Disposed.WaitAsync(s_timeout); + } + + [Fact] + public async Task HttpResponse_UpstreamErrorFollowsBufferedBytesAndOutstandingAck() + { + var source = new ControlledResponseStream(); + await using var peer = await ProtocolPeer.StartAsync(new StreamRequestHandler(source)); + + await source.PushAsync("first"u8.ToArray()); + await peer.BeginRequestAsync(); + await peer.AcknowledgeAsync(await peer.NextMethodAsync()); + var first = await peer.NextMethodAsync(); + + await source.PushAsync("partial"u8.ToArray()); + await source.FailAsync(new IOException("upstream failed")); + await source.WaitForReadsAsync(3); + await peer.AssertNoMethodAsync(); + + await peer.AcknowledgeAsync(first); + var partial = await peer.NextMethodAsync(); + Assert.Equal("partial"u8.ToArray(), Data(partial)); + await peer.AssertNoMethodAsync(); + + await peer.AcknowledgeAsync(partial); + var error = await peer.NextMethodAsync(); + Assert.True(error.GetProperty("params").GetProperty("end").GetBoolean()); + Assert.Equal("upstream failed", error.GetProperty("params").GetProperty("error").GetProperty("message").GetString()); + await peer.AcknowledgeAsync(error); + await source.Disposed.WaitAsync(s_timeout); + } + + [Fact] + public async Task HttpResponse_RuntimeCancellationDisposesSourceWithoutWaitingForDataAck() + { + var source = new ControlledResponseStream(); + await using var peer = await ProtocolPeer.StartAsync(new StreamRequestHandler(source)); + + await source.PushAsync("first"u8.ToArray()); + await peer.BeginRequestAsync(); + await peer.AcknowledgeAsync(await peer.NextMethodAsync()); + var first = await peer.NextMethodAsync(); + Assert.Equal("first"u8.ToArray(), Data(first)); + + await peer.CancelRequestAsync(); + var error = await peer.NextMethodAsync(); + Assert.True(error.GetProperty("params").GetProperty("end").GetBoolean()); + Assert.Equal("cancelled", error.GetProperty("params").GetProperty("error").GetProperty("code").GetString()); + await source.Disposed.WaitAsync(s_timeout); + await peer.AcknowledgeAsync(error); + } + + [Fact] + public async Task HttpResponse_RuntimeCancellationWhileSourceIsIdleReportsCancellation() + { + var source = new ControlledResponseStream(); + await using var peer = await ProtocolPeer.StartAsync(new StreamRequestHandler(source)); + + await peer.BeginRequestAsync(); + await peer.AcknowledgeAsync(await peer.NextMethodAsync()); + await source.ReadStarted.WaitAsync(s_timeout); + + await peer.CancelRequestAsync(); + var error = await peer.NextMethodAsync(); + Assert.True(error.GetProperty("params").GetProperty("end").GetBoolean()); + Assert.Equal("cancelled", error.GetProperty("params").GetProperty("error").GetProperty("code").GetString()); + await source.Disposed.WaitAsync(s_timeout); + await peer.AcknowledgeAsync(error); + } + + [Fact] + public async Task HttpResponse_RejectedDataRpcDisposesSourceAndReportsError() + { + var source = new ControlledResponseStream(); + await using var peer = await ProtocolPeer.StartAsync(new StreamRequestHandler(source)); + + await source.PushAsync("first"u8.ToArray()); + await peer.BeginRequestAsync(); + await peer.AcknowledgeAsync(await peer.NextMethodAsync()); + var first = await peer.NextMethodAsync(); + + await peer.RejectAsync(first, "write rejected"); + await source.Disposed.WaitAsync(s_timeout); + var error = await peer.NextMethodAsync(); + Assert.Contains( + "write rejected", + error.GetProperty("params").GetProperty("error").GetProperty("message").GetString(), + StringComparison.Ordinal); + await peer.AcknowledgeAsync(error); + } + + [Fact] + public async Task HttpResponse_ConnectionLossDisposesSourceWithOutstandingDataRpc() + { + var source = new ControlledResponseStream(); + await using var peer = await ProtocolPeer.StartAsync(new StreamRequestHandler(source)); + + await source.PushAsync("first"u8.ToArray()); + await peer.BeginRequestAsync(); + await peer.AcknowledgeAsync(await peer.NextMethodAsync()); + Assert.Equal("first"u8.ToArray(), Data(await peer.NextMethodAsync())); + + peer.CloseConnection(); + await source.Disposed.WaitAsync(s_timeout); + } + + [Fact] + public async Task HttpResponse_ConnectionLossDisposesIdleSource() + { + var source = new ControlledResponseStream(); + await using var peer = await ProtocolPeer.StartAsync(new StreamRequestHandler(source)); + + await peer.BeginRequestAsync(); + await peer.AcknowledgeAsync(await peer.NextMethodAsync()); + await source.ReadStarted.WaitAsync(s_timeout); + + peer.CloseConnection(); + await source.Disposed.WaitAsync(s_timeout); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task HttpResponse_ClientShutdownDisposesIdleSource(bool force) + { + var source = new ControlledResponseStream(); + await using var peer = await ProtocolPeer.StartAsync(new StreamRequestHandler(source)); + + await peer.BeginRequestAsync(); + await peer.AcknowledgeAsync(await peer.NextMethodAsync()); + await source.ReadStarted.WaitAsync(s_timeout); + + await peer.StopClientAsync(force); + await source.Disposed.WaitAsync(s_timeout); + } + + private static string? Method(JsonElement message) => + message.GetProperty("method").GetString(); + + private static byte[] Data(JsonElement message) + { + Assert.Equal("llmInference.httpResponseChunk", Method(message)); + var parameters = message.GetProperty("params"); + Assert.False(parameters.GetProperty("end").GetBoolean()); + Assert.True(parameters.GetProperty("binary").GetBoolean()); + return Convert.FromBase64String(parameters.GetProperty("data").GetString()!); + } + + private sealed class StreamRequestHandler(Stream source) : CopilotRequestHandler + { + protected override Task SendRequestAsync(HttpRequestMessage request, CopilotRequestContext ctx) => + Task.FromResult(CreateResponse(source)); + + private static HttpResponseMessage CreateResponse(Stream source) => + new(HttpStatusCode.OK) + { + Content = new StreamContent(source), + }; + } + + private sealed class ControlledResponseStream : Stream + { + private readonly Channel _items = Channel.CreateBounded( + new BoundedChannelOptions(1) + { + SingleReader = true, + SingleWriter = true, + FullMode = BoundedChannelFullMode.Wait, + }); + private readonly TaskCompletionSource _disposed = new(TaskCreationOptions.RunContinuationsAsynchronously); + private readonly TaskCompletionSource _readStarted = new(TaskCreationOptions.RunContinuationsAsynchronously); + private int _readCount; + private int _bytesRead; + private int _isDisposed; + private byte[]? _pendingData; + private int _pendingOffset; + + internal int ReadCount => Volatile.Read(ref _readCount); + + internal int BytesRead => Volatile.Read(ref _bytesRead); + + internal Task Disposed => _disposed.Task; + + internal Task ReadStarted => _readStarted.Task; + + public override bool CanRead => true; + + public override bool CanSeek => false; + + public override bool CanWrite => false; + + public override long Length => throw new NotSupportedException(); + + public override long Position + { + get => throw new NotSupportedException(); + set => throw new NotSupportedException(); + } + + internal ValueTask PushAsync(byte[] data) => + _items.Writer.WriteAsync(new ReadItem(data, Error: null, End: false)); + + internal ValueTask FailAsync(Exception error) => + _items.Writer.WriteAsync(new ReadItem(Data: null, Error: error, End: false)); + + internal ValueTask CompleteAsync() => + _items.Writer.WriteAsync(new ReadItem(Data: null, Error: null, End: true)); + + internal async Task WaitForReadsAsync(int expected) + { + using var cts = new CancellationTokenSource(s_timeout); + while (ReadCount < expected) + { + await Task.Delay(TimeSpan.FromMilliseconds(10), cts.Token); + } + } + + internal async Task WaitForBytesAsync(int expected) + { + using var cts = new CancellationTokenSource(s_timeout); + while (BytesRead < expected) + { + await Task.Delay(TimeSpan.FromMilliseconds(10), cts.Token); + } + } + + public override int Read(byte[] buffer, int offset, int count) => + throw new NotSupportedException(); + + public override Task ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken) => + ReadCoreAsync(buffer.AsMemory(offset, count), cancellationToken).AsTask(); + + public override ValueTask ReadAsync(Memory buffer, CancellationToken cancellationToken = default) => + ReadCoreAsync(buffer, cancellationToken); + + public override void Flush() + { + } + + public override long Seek(long offset, SeekOrigin origin) => + throw new NotSupportedException(); + + public override void SetLength(long value) => + throw new NotSupportedException(); + + public override void Write(byte[] buffer, int offset, int count) => + throw new NotSupportedException(); + + protected override void Dispose(bool disposing) + { + if (disposing && Interlocked.Exchange(ref _isDisposed, 1) == 0) + { + _items.Writer.TryComplete(); + _disposed.TrySetResult(); + } + + base.Dispose(disposing); + } + + private async ValueTask ReadCoreAsync(Memory buffer, CancellationToken cancellationToken) + { + if (_pendingData is not null) + { + return CopyPendingData(buffer); + } + + _readStarted.TrySetResult(); + var item = await _items.Reader.ReadAsync(cancellationToken); + Interlocked.Increment(ref _readCount); + if (item.Error is not null) + { + throw item.Error; + } + + if (item.End) + { + return 0; + } + + Assert.NotNull(item.Data); + _pendingData = item.Data; + _pendingOffset = 0; + return CopyPendingData(buffer); + } + + private int CopyPendingData(Memory buffer) + { + var available = _pendingData!.Length - _pendingOffset; + var count = Math.Min(available, buffer.Length); + _pendingData.AsMemory(_pendingOffset, count).CopyTo(buffer); + _pendingOffset += count; + if (_pendingOffset == _pendingData.Length) + { + _pendingData = null; + _pendingOffset = 0; + } + + Interlocked.Add(ref _bytesRead, count); + return count; + } + + private sealed record ReadItem(byte[]? Data, Exception? Error, bool End); + } + + private sealed class ProtocolPeer : IAsyncDisposable + { + private readonly CopilotClient _client; + private readonly TcpClient _tcpClient; + private readonly NetworkStream _stream; + private readonly CancellationTokenSource _disposeCts = new(); + private readonly Channel _messages = Channel.CreateUnbounded(); + private readonly Task _readLoop; + private int _nextRequestId = 100_000; + + private ProtocolPeer(CopilotClient client, TcpClient tcpClient) + { + _client = client; + _tcpClient = tcpClient; + _stream = tcpClient.GetStream(); + _readLoop = ReadLoopAsync(); + } + + internal static async Task StartAsync(CopilotRequestHandler handler) + { + using var listener = new TcpListener(IPAddress.Loopback, 0); + listener.Start(); + var endpoint = (IPEndPoint)listener.LocalEndpoint; + var client = new CopilotClient(new CopilotClientOptions + { + Connection = RuntimeConnection.ForUri($"http://127.0.0.1:{endpoint.Port}"), + RequestHandler = handler, + }); + var startTask = client.StartAsync(); + var tcpClient = await listener.AcceptTcpClientAsync().WaitAsync(s_timeout); + var stream = tcpClient.GetStream(); + + var connect = await ReadFrameAsync(stream, CancellationToken.None).WaitAsync(s_timeout); + Assert.Equal("connect", Method(connect)); + await WriteResultAsync( + stream, + connect, + writer => + { + writer.WriteStartObject(); + writer.WriteBoolean("ok", true); + writer.WriteNumber("protocolVersion", 3); + writer.WriteString("version", "test"); + writer.WriteEndObject(); + }, + CancellationToken.None); + + var setProvider = await ReadFrameAsync(stream, CancellationToken.None).WaitAsync(s_timeout); + Assert.Equal("llmInference.setProvider", Method(setProvider)); + await WriteResultAsync( + stream, + setProvider, + writer => + { + writer.WriteStartObject(); + writer.WriteBoolean("success", true); + writer.WriteEndObject(); + }, + CancellationToken.None); + + await startTask.WaitAsync(s_timeout); + return new ProtocolPeer(client, tcpClient); + } + + internal async Task BeginRequestAsync() + { + await WriteRequestAsync( + "llmInference.httpRequestStart", + writer => + { + writer.WriteString("requestId", "test"); + writer.WriteString("method", "GET"); + writer.WriteString("url", "http://unused.test"); + writer.WriteStartObject("headers"); + writer.WriteEndObject(); + }); + await WriteRequestAsync( + "llmInference.httpRequestChunk", + writer => + { + writer.WriteString("requestId", "test"); + writer.WriteString("data", string.Empty); + writer.WriteBoolean("end", true); + }); + } + + internal Task CancelRequestAsync() => + WriteRequestAsync( + "llmInference.httpRequestChunk", + writer => + { + writer.WriteString("requestId", "test"); + writer.WriteString("data", string.Empty); + writer.WriteBoolean("cancel", true); + }); + + internal async Task NextMethodAsync(CancellationToken cancellationToken = default) + { + using var timeout = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken); + timeout.CancelAfter(s_timeout); + while (true) + { + var message = await _messages.Reader.ReadAsync(timeout.Token); + if (message.TryGetProperty("method", out _)) + { + return message; + } + } + } + + internal async Task AssertNoMethodAsync() + { + using var cts = new CancellationTokenSource(TimeSpan.FromMilliseconds(100)); + await Assert.ThrowsAnyAsync( + () => NextMethodAsync(cts.Token)); + } + + internal Task AcknowledgeAsync(JsonElement request) => + WriteResultAsync( + _stream, + request, + writer => + { + writer.WriteStartObject(); + writer.WriteBoolean("accepted", true); + writer.WriteEndObject(); + }, + _disposeCts.Token); + + internal Task RejectAsync(JsonElement request, string message) => + WriteFrameAsync( + _stream, + writer => + { + writer.WriteStartObject(); + writer.WriteString("jsonrpc", "2.0"); + writer.WritePropertyName("id"); + request.GetProperty("id").WriteTo(writer); + writer.WriteStartObject("error"); + writer.WriteNumber("code", -32603); + writer.WriteString("message", message); + writer.WriteEndObject(); + writer.WriteEndObject(); + }, + _disposeCts.Token); + + internal void CloseConnection() => _tcpClient.Dispose(); + + internal Task StopClientAsync(bool force) => + force ? _client.ForceStopAsync() : _client.StopAsync(); + + public async ValueTask DisposeAsync() + { + _disposeCts.Cancel(); + _tcpClient.Dispose(); + try + { + await _readLoop; + } + catch (Exception ex) when (ex is OperationCanceledException or IOException or ObjectDisposedException) + { + // These exceptions are expected when the protocol peer is torn down. + } + + await _client.ForceStopAsync(); + await _client.DisposeAsync(); + _disposeCts.Dispose(); + } + + private async Task WriteRequestAsync(string method, Action writeParams) + { + var id = Interlocked.Increment(ref _nextRequestId); + await WriteFrameAsync( + _stream, + writer => + { + writer.WriteStartObject(); + writer.WriteString("jsonrpc", "2.0"); + writer.WriteNumber("id", id); + writer.WriteString("method", method); + writer.WriteStartObject("params"); + writeParams(writer); + writer.WriteEndObject(); + writer.WriteEndObject(); + }, + _disposeCts.Token); + } + + private async Task ReadLoopAsync() + { + try + { + while (!_disposeCts.IsCancellationRequested) + { + await _messages.Writer.WriteAsync( + await ReadFrameAsync(_stream, _disposeCts.Token), + _disposeCts.Token); + } + } + catch (Exception ex) when (ex is OperationCanceledException or IOException or ObjectDisposedException) + { + // Connection closure is the expected way to stop the read loop. + } + finally + { + _messages.Writer.TryComplete(); + } + } + + private static async Task ReadFrameAsync(Stream stream, CancellationToken cancellationToken) + { + using var header = new MemoryStream(); + var current = new byte[1]; + while (true) + { + var read = await stream.ReadAsync(current, cancellationToken); + if (read == 0) + { + throw new EndOfStreamException(); + } + + header.WriteByte(current[0]); + if (header.Length >= 4) + { + var bytes = header.GetBuffer(); + var length = (int)header.Length; + if (bytes[length - 4] == '\r' + && bytes[length - 3] == '\n' + && bytes[length - 2] == '\r' + && bytes[length - 1] == '\n') + { + break; + } + } + } + + var headerText = Encoding.ASCII.GetString(header.GetBuffer(), 0, (int)header.Length); + var contentLengthLine = headerText + .Split("\r\n", StringSplitOptions.RemoveEmptyEntries) + .Single(line => line.StartsWith("Content-Length:", StringComparison.OrdinalIgnoreCase)); + var contentLength = int.Parse(contentLengthLine["Content-Length:".Length..].Trim()); + var body = new byte[contentLength]; + await stream.ReadExactlyAsync(body, cancellationToken); + using var document = JsonDocument.Parse(body); + return document.RootElement.Clone(); + } + + private static Task WriteResultAsync( + Stream stream, + JsonElement request, + Action writeResult, + CancellationToken cancellationToken) => + WriteFrameAsync( + stream, + writer => + { + writer.WriteStartObject(); + writer.WriteString("jsonrpc", "2.0"); + writer.WritePropertyName("id"); + request.GetProperty("id").WriteTo(writer); + writer.WritePropertyName("result"); + writeResult(writer); + writer.WriteEndObject(); + }, + cancellationToken); + + private static async Task WriteFrameAsync( + Stream stream, + Action writeMessage, + CancellationToken cancellationToken) + { + using var bodyStream = new MemoryStream(); + using (var writer = new Utf8JsonWriter(bodyStream)) + { + writeMessage(writer); + } + + var body = bodyStream.ToArray(); + var header = Encoding.ASCII.GetBytes($"Content-Length: {body.Length}\r\n\r\n"); + await stream.WriteAsync(header, cancellationToken); + await stream.WriteAsync(body, cancellationToken); + await stream.FlushAsync(cancellationToken); + } + } +} +#endif diff --git a/dotnet/test/Unit/JsonRpcTests.cs b/dotnet/test/Unit/JsonRpcTests.cs index f8bf44d6b1..2eeefa19ad 100644 --- a/dotnet/test/Unit/JsonRpcTests.cs +++ b/dotnet/test/Unit/JsonRpcTests.cs @@ -404,8 +404,6 @@ public JsonRpcReflection(Stream sendStream, Stream receiveStream) public void StartListening() => JsonRpcType.GetMethod(nameof(StartListening))!.Invoke(_instance, null); - public Task Completion => (Task)JsonRpcType.GetProperty(nameof(Completion))!.GetValue(_instance)!; - public void SetLocalRpcMethod(string methodName, Delegate handler, bool singleObjectParam = false) => JsonRpcType.GetMethod("SetLocalRpcMethod")!.Invoke(_instance, [methodName, handler, singleObjectParam]); diff --git a/go/client.go b/go/client.go index 4623872a53..5399662aa0 100644 --- a/go/client.go +++ b/go/client.go @@ -156,6 +156,8 @@ type Client struct { sessionsMux sync.Mutex gitHubTokenProviders map[string]GitHubTokenProvider gitHubTokenProvidersMux sync.RWMutex + requestAdapter *copilotRequestAdapter + requestAdapterMux sync.Mutex sessionOperations map[string]*sessionOperation sessionOperationsMux sync.Mutex isExternalServer bool @@ -601,6 +603,7 @@ func (c *Client) Stop() error { c.sessions = make(map[string]*Session) c.sessionsMux.Unlock() c.clearGitHubTokenProviders() + c.closeCopilotRequestAdapter() c.startStopMux.Lock() defer c.startStopMux.Unlock() @@ -724,6 +727,7 @@ func (c *Client) ForceStop() { session.cancelPendingExternalTools() } c.clearGitHubTokenProviders() + c.closeCopilotRequestAdapter() c.startStopMux.Lock() defer c.startStopMux.Unlock() @@ -2506,9 +2510,17 @@ func (c *Client) setupNotificationHandler() { if c.options.RequestHandler != nil { llmInference := c.RPC.LlmInference - handlers.LlmInference = newCopilotRequestAdapter(c.options.RequestHandler, func() *rpc.ServerLlmInferenceAPI { + adapter := newCopilotRequestAdapter(c.options.RequestHandler, func() *rpc.ServerLlmInferenceAPI { return llmInference }) + c.requestAdapterMux.Lock() + previous := c.requestAdapter + c.requestAdapter = adapter + c.requestAdapterMux.Unlock() + if previous != nil { + previous.close() + } + handlers.LlmInference = adapter } if c.options.OnGitHubTelemetry != nil { handlers.GitHubTelemetry = &gitHubTelemetryAdapter{callback: c.options.OnGitHubTelemetry} @@ -2546,6 +2558,7 @@ func (c *Client) clearGitHubTokenProviders() { } func (c *Client) handleConnectionClose() { + c.closeCopilotRequestAdapter() c.clearGitHubTokenProviders() c.sessionsMux.Lock() sessions := make([]*Session, 0, len(c.sessions)) @@ -2565,6 +2578,15 @@ func (c *Client) handleConnectionClose() { }() } +func (c *Client) closeCopilotRequestAdapter() { + c.requestAdapterMux.Lock() + adapter := c.requestAdapter + c.requestAdapterMux.Unlock() + if adapter != nil { + adapter.close() + } +} + func (c *Client) lockSessionOperation(sessionID string) func() { c.sessionOperationsMux.Lock() if c.sessionOperations == nil { diff --git a/go/copilot_request_handler.go b/go/copilot_request_handler.go index ba8bb9b919..cc5254f5e3 100644 --- a/go/copilot_request_handler.go +++ b/go/copilot_request_handler.go @@ -8,6 +8,7 @@ import ( "bytes" "context" "encoding/base64" + "errors" "fmt" "io" "net/http" @@ -182,7 +183,7 @@ func (h *CopilotRequestHandler) handleHTTP(rctx *CopilotRequestContext, sink *re return err } defer resp.Body.Close() - return streamResponseToSink(resp, sink) + return streamResponseToSink(rctx.Context, resp, sink) } func buildHTTPRequest(rctx *CopilotRequestContext) (*http.Request, error) { @@ -218,28 +219,26 @@ func drainBody(ch <-chan CopilotWebSocketMessage) []byte { return buf.Bytes() } -func streamResponseToSink(resp *http.Response, sink *responseSink) error { +func streamResponseToSink(ctx context.Context, resp *http.Response, sink *responseSink) error { if err := sink.start(resp.StatusCode, statusText(resp), cloneHeader(resp.Header)); err != nil { return err } - buf := make([]byte, 32*1024) + reader := newHTTPResponseReader(resp.Body) + defer reader.Close() + + var chunk []byte for { - n, readErr := resp.Body.Read(buf) - if n > 0 { - // writeText copies eagerly via string(...), so the reused read - // buffer can be passed directly without an extra per-chunk alloc. - if err := sink.writeText(buf[:n]); err != nil { - return err - } + ok, err := reader.nextChunk(ctx, &chunk) + if err != nil { + return err } - if readErr == io.EOF { - break + if !ok { + return sink.end() } - if readErr != nil { - return sink.sinkError(readErr.Error(), "") + if err := sink.writeBinary(chunk); err != nil { + return err } } - return sink.end() } func statusText(resp *http.Response) string { @@ -542,7 +541,7 @@ type pendingExchange struct { mu sync.Mutex queue *frameQueue ctx context.Context - cancel context.CancelFunc + cancel context.CancelCauseFunc started bool finished bool } @@ -553,13 +552,20 @@ type copilotRequestAdapter struct { mu sync.Mutex pending map[string]*pendingExchange + closed bool + + connectionCtx context.Context + connectionCancel context.CancelCauseFunc } -func newCopilotRequestAdapter(handler *CopilotRequestHandler, getRPC func() *rpc.ServerLlmInferenceAPI) rpc.LlmInferenceHandler { +func newCopilotRequestAdapter(handler *CopilotRequestHandler, getRPC func() *rpc.ServerLlmInferenceAPI) *copilotRequestAdapter { + connectionCtx, connectionCancel := context.WithCancelCause(context.Background()) return &copilotRequestAdapter{ - handler: handler, - getRPC: getRPC, - pending: make(map[string]*pendingExchange), + handler: handler, + getRPC: getRPC, + pending: make(map[string]*pendingExchange), + connectionCtx: connectionCtx, + connectionCancel: connectionCancel, } } @@ -576,8 +582,12 @@ func (a *copilotRequestAdapter) getOrCreateExchange(requestID string) *pendingEx if exchange, ok := a.pending[requestID]; ok { return exchange } - ctx, cancel := context.WithCancel(context.Background()) + ctx, cancel := context.WithCancelCause(a.connectionCtx) exchange := &pendingExchange{queue: newFrameQueue(), ctx: ctx, cancel: cancel} + if a.closed { + exchange.queue.close() + return exchange + } a.pending[requestID] = exchange return exchange } @@ -645,7 +655,7 @@ func (a *copilotRequestAdapter) HttpRequestChunk(params *rpc.LlmInferenceHTTPReq func (a *copilotRequestAdapter) routeChunk(exchange *pendingExchange, params *rpc.LlmInferenceHTTPRequestChunkRequest) { if params.Cancel != nil && *params.Cancel { - exchange.cancel() + exchange.cancel(errRuntimeRequestCancelled) exchange.queue.close() return } @@ -661,15 +671,31 @@ func (a *copilotRequestAdapter) routeChunk(exchange *pendingExchange, params *rp } func (a *copilotRequestAdapter) runHandler(rctx *CopilotRequestContext, sink *responseSink, exchange *pendingExchange) { + defer exchange.cancel(nil) + err := a.handler.handle(rctx, sink) if err != nil { - if exchange.ctx.Err() != nil { + cause := context.Cause(exchange.ctx) + if errors.Is(cause, errRuntimeRequestCancelled) { a.finishCancelled(sink, exchange) return } + if cause != nil { + a.removePending(sink.requestID) + return + } a.failViaSink(sink, exchange, err.Error()) return } + cause := context.Cause(exchange.ctx) + if errors.Is(cause, errRuntimeRequestCancelled) { + a.finishCancelled(sink, exchange) + return + } + if cause != nil { + a.removePending(sink.requestID) + return + } exchange.mu.Lock() finished := exchange.finished exchange.mu.Unlock() @@ -701,7 +727,7 @@ func (a *copilotRequestAdapter) finishCancelled(sink *responseSink, exchange *pe return } if !started { - _ = sink.start(499, "", http.Header{}) + _ = sink.startWithContext(sink.adapter.connectionCtx, 499, "", http.Header{}) } _ = sink.sinkError("Request cancelled by runtime", "cancelled") } @@ -712,6 +738,31 @@ func (a *copilotRequestAdapter) removePending(requestID string) { a.mu.Unlock() } +var ( + errRuntimeRequestCancelled = errors.New("request cancelled by runtime") + errRPCConnectionClosed = errors.New("RPC connection closed") +) + +func (a *copilotRequestAdapter) close() { + a.mu.Lock() + if a.closed { + a.mu.Unlock() + return + } + a.closed = true + exchanges := make([]*pendingExchange, 0, len(a.pending)) + for _, exchange := range a.pending { + exchanges = append(exchanges, exchange) + } + clear(a.pending) + a.mu.Unlock() + + a.connectionCancel(errRPCConnectionClosed) + for _, exchange := range exchanges { + exchange.queue.close() + } +} + func stringOrEmpty(value *string) string { if value == nil { return "" @@ -742,6 +793,10 @@ func (s *responseSink) rpcAPI() (*rpc.ServerLlmInferenceAPI, error) { } func (s *responseSink) start(status int, statusTxt string, headers http.Header) error { + return s.startWithContext(s.exchange.ctx, status, statusTxt, headers) +} + +func (s *responseSink) startWithContext(ctx context.Context, status int, statusTxt string, headers http.Header) error { s.exchange.mu.Lock() if s.exchange.started { s.exchange.mu.Unlock() @@ -766,12 +821,15 @@ func (s *responseSink) start(status int, statusTxt string, headers http.Header) if h == nil { h = map[string][]string{} } - _, err = api.HttpResponseStart(context.Background(), &rpc.LlmInferenceHTTPResponseStartRequest{ + result, err := api.HttpResponseStart(ctx, &rpc.LlmInferenceHTTPResponseStartRequest{ RequestID: s.requestID, Status: int64(status), StatusText: st, Headers: h, }) + if err == nil && !result.Accepted { + return fmt.Errorf("llmInference.httpResponseStart was rejected") + } return err } @@ -808,7 +866,10 @@ func (s *responseSink) writeRaw(data string, binary bool) error { b := true chunk.Binary = &b } - _, err = api.HttpResponseChunk(context.Background(), chunk) + result, err := api.HttpResponseChunk(s.exchange.ctx, chunk) + if err == nil && !result.Accepted { + return fmt.Errorf("llmInference.httpResponseChunk was rejected") + } return err } @@ -826,11 +887,14 @@ func (s *responseSink) end() error { return err } end := true - _, err = api.HttpResponseChunk(context.Background(), &rpc.LlmInferenceHTTPResponseChunkRequest{ + result, err := api.HttpResponseChunk(s.exchange.ctx, &rpc.LlmInferenceHTTPResponseChunkRequest{ RequestID: s.requestID, Data: "", End: &end, }) + if err == nil && !result.Accepted { + return fmt.Errorf("llmInference.httpResponseChunk was rejected") + } return err } @@ -853,11 +917,14 @@ func (s *responseSink) sinkError(message string, code string) error { c := code chunkErr.Code = &c } - _, err = api.HttpResponseChunk(context.Background(), &rpc.LlmInferenceHTTPResponseChunkRequest{ + result, err := api.HttpResponseChunk(s.adapter.connectionCtx, &rpc.LlmInferenceHTTPResponseChunkRequest{ RequestID: s.requestID, Data: "", End: &end, Error: chunkErr, }) + if err == nil && !result.Accepted { + return fmt.Errorf("llmInference.httpResponseChunk was rejected") + } return err } diff --git a/go/http_response_forwarding_test.go b/go/http_response_forwarding_test.go new file mode 100644 index 0000000000..18f8ae82cc --- /dev/null +++ b/go/http_response_forwarding_test.go @@ -0,0 +1,643 @@ +package copilot + +import ( + "bufio" + "encoding/base64" + "encoding/json" + "errors" + "fmt" + "io" + "net" + "net/http" + "strconv" + "strings" + "sync" + "testing" + "time" + + "github.com/github/copilot-sdk/go/internal/jsonrpc2" + "github.com/github/copilot-sdk/go/rpc" +) + +const responseProtocolTimeout = 5 * time.Second + +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) { + return f(req) +} + +type trackedResponseBody struct { + reader *io.PipeReader + closed chan struct{} + closeOnce sync.Once +} + +type zeroProgressResponseBody struct { + closed chan struct{} + closeOnce sync.Once +} + +func (b *zeroProgressResponseBody) Read([]byte) (int, error) { + return 0, nil +} + +func (b *zeroProgressResponseBody) Close() error { + b.closeOnce.Do(func() { + close(b.closed) + }) + return nil +} + +func newTrackedResponseBody() (*trackedResponseBody, *io.PipeWriter) { + reader, writer := io.Pipe() + return &trackedResponseBody{reader: reader, closed: make(chan struct{})}, writer +} + +func (b *trackedResponseBody) Read(p []byte) (int, error) { + return b.reader.Read(p) +} + +func (b *trackedResponseBody) Close() error { + var err error + b.closeOnce.Do(func() { + close(b.closed) + err = b.reader.Close() + }) + return err +} + +type responseProtocolMessage struct { + JSONRPC string `json:"jsonrpc"` + ID json.RawMessage `json:"id,omitempty"` + Method string `json:"method,omitempty"` + Params json.RawMessage `json:"params,omitempty"` +} + +type responseProtocolPeer struct { + t *testing.T + conn net.Conn + reader *bufio.Reader + client *jsonrpc2.Client + adapter *copilotRequestAdapter + stop sync.Once +} + +func newResponseProtocolPeer(t *testing.T, body io.ReadCloser) *responseProtocolPeer { + return newResponseProtocolPeerWithTransport(t, roundTripFunc(func(*http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusOK, + Status: "200 OK", + Header: http.Header{"Content-Type": {"text/event-stream"}}, + Body: body, + }, nil + })) +} + +func newResponseProtocolPeerWithTransport(t *testing.T, transport http.RoundTripper) *responseProtocolPeer { + t.Helper() + sdkConn, runtimeConn := net.Pipe() + client := jsonrpc2.NewClient(sdkConn, sdkConn) + serverRPC := rpc.NewServerRPC(client) + handler := &CopilotRequestHandler{ + Transport: transport, + } + adapter := newCopilotRequestAdapter(handler, func() *rpc.ServerLlmInferenceAPI { + return serverRPC.LlmInference + }) + rpc.RegisterClientGlobalAPIHandlers(client, &rpc.ClientGlobalAPIHandlers{LlmInference: adapter}) + client.SetOnClose(adapter.close) + client.Start() + + peer := &responseProtocolPeer{ + t: t, + conn: runtimeConn, + reader: bufio.NewReader(runtimeConn), + client: client, + adapter: adapter, + } + t.Cleanup(peer.close) + return peer +} + +func (p *responseProtocolPeer) close() { + p.stop.Do(func() { + _ = p.conn.Close() + p.adapter.close() + p.client.Stop() + }) +} + +func (p *responseProtocolPeer) beginRequest() { + p.t.Helper() + p.send(responseProtocolMessage{ + JSONRPC: "2.0", + ID: json.RawMessage("1"), + Method: "llmInference.httpRequestStart", + Params: mustMarshal(p.t, map[string]any{ + "requestId": "test", + "method": "GET", + "url": "http://unused.test", + "headers": map[string][]string{}, + }), + }) + p.send(responseProtocolMessage{ + JSONRPC: "2.0", + ID: json.RawMessage("2"), + Method: "llmInference.httpRequestChunk", + Params: mustMarshal(p.t, map[string]any{ + "requestId": "test", + "data": "", + "end": true, + }), + }) +} + +func (p *responseProtocolPeer) cancelRequest() { + p.t.Helper() + p.send(responseProtocolMessage{ + JSONRPC: "2.0", + ID: json.RawMessage("3"), + Method: "llmInference.httpRequestChunk", + Params: mustMarshal(p.t, map[string]any{ + "requestId": "test", + "data": "", + "cancel": true, + }), + }) +} + +func (p *responseProtocolPeer) send(message any) { + p.t.Helper() + data, err := json.Marshal(message) + if err != nil { + p.t.Fatal(err) + } + if err := p.conn.SetWriteDeadline(time.Now().Add(responseProtocolTimeout)); err != nil { + p.t.Fatal(err) + } + if _, err := fmt.Fprintf(p.conn, "Content-Length: %d\r\n\r\n", len(data)); err != nil { + p.t.Fatal(err) + } + if _, err := p.conn.Write(data); err != nil { + p.t.Fatal(err) + } +} + +func (p *responseProtocolPeer) nextRequest() responseProtocolMessage { + p.t.Helper() + message, ok := p.nextRequestWithin(responseProtocolTimeout) + if !ok { + p.t.Fatal("timed out waiting for runtime request") + } + return message +} + +func (p *responseProtocolPeer) nextRequestWithin(timeout time.Duration) (responseProtocolMessage, bool) { + p.t.Helper() + deadline := time.Now().Add(timeout) + for { + if err := p.conn.SetReadDeadline(deadline); err != nil { + p.t.Fatal(err) + } + message, err := readResponseProtocolMessage(p.reader) + if err != nil { + if netErr, ok := err.(net.Error); ok && netErr.Timeout() { + return responseProtocolMessage{}, false + } + p.t.Fatal(err) + } + if message.Method != "" { + return message, true + } + } +} + +func (p *responseProtocolPeer) ack(message responseProtocolMessage) { + p.respondAccepted(message, true) +} + +func (p *responseProtocolPeer) respondAccepted(message responseProtocolMessage, accepted bool) { + p.t.Helper() + p.send(struct { + JSONRPC string `json:"jsonrpc"` + ID json.RawMessage `json:"id"` + Result map[string]bool `json:"result"` + }{ + JSONRPC: "2.0", + ID: message.ID, + Result: map[string]bool{"accepted": accepted}, + }) +} + +func (p *responseProtocolPeer) reject(message responseProtocolMessage, text string) { + p.t.Helper() + p.send(struct { + JSONRPC string `json:"jsonrpc"` + ID json.RawMessage `json:"id"` + Error map[string]any `json:"error"` + }{ + JSONRPC: "2.0", + ID: message.ID, + Error: map[string]any{"code": -32603, "message": text}, + }) +} + +func (p *responseProtocolPeer) startAndAck() { + p.t.Helper() + p.beginRequest() + head := p.nextRequest() + if head.Method != "llmInference.httpResponseStart" { + p.t.Fatalf("first request method = %q, want httpResponseStart", head.Method) + } + p.ack(head) +} + +func readResponseProtocolMessage(reader *bufio.Reader) (responseProtocolMessage, error) { + contentLength := 0 + for { + line, err := reader.ReadString('\n') + if err != nil { + return responseProtocolMessage{}, err + } + if line == "\r\n" { + break + } + if value, ok := strings.CutPrefix(line, "Content-Length:"); ok { + contentLength, err = strconv.Atoi(strings.TrimSpace(value)) + if err != nil { + return responseProtocolMessage{}, err + } + } + } + if contentLength == 0 { + return responseProtocolMessage{}, errors.New("missing Content-Length") + } + data := make([]byte, contentLength) + if _, err := io.ReadFull(reader, data); err != nil { + return responseProtocolMessage{}, err + } + var message responseProtocolMessage + if err := json.Unmarshal(data, &message); err != nil { + return responseProtocolMessage{}, err + } + return message, nil +} + +func mustMarshal(t *testing.T, value any) json.RawMessage { + t.Helper() + data, err := json.Marshal(value) + if err != nil { + t.Fatal(err) + } + return data +} + +func responseChunkParams(t *testing.T, message responseProtocolMessage) map[string]any { + t.Helper() + if message.Method != "llmInference.httpResponseChunk" { + t.Fatalf("request method = %q, want httpResponseChunk", message.Method) + } + var params map[string]any + if err := json.Unmarshal(message.Params, ¶ms); err != nil { + t.Fatal(err) + } + return params +} + +func responseChunkData(t *testing.T, message responseProtocolMessage) []byte { + t.Helper() + params := responseChunkParams(t, message) + if params["end"] != false { + t.Fatalf("chunk end = %v, want false", params["end"]) + } + if params["binary"] != true { + t.Fatalf("chunk binary = %v, want true", params["binary"]) + } + data, err := base64.StdEncoding.DecodeString(params["data"].(string)) + if err != nil { + t.Fatal(err) + } + if len(data) == 0 || len(data) > httpResponseReadAheadSize { + t.Fatalf("chunk size = %d, want 1..%d", len(data), httpResponseReadAheadSize) + } + return data +} + +func waitForBodyClose(t *testing.T, body *trackedResponseBody) { + t.Helper() + select { + case <-body.closed: + case <-time.After(responseProtocolTimeout): + t.Fatal("response body was not closed") + } +} + +func TestHTTPResponseReaderCloseStopsZeroProgressSource(t *testing.T) { + body := &zeroProgressResponseBody{closed: make(chan struct{})} + reader := newHTTPResponseReader(body) + closed := make(chan struct{}) + + go func() { + reader.Close() + close(closed) + }() + + select { + case <-closed: + case <-time.After(responseProtocolTimeout): + t.Fatal("reader close blocked on a zero-progress source") + } + + select { + case <-body.closed: + default: + t.Fatal("response body was not closed") + } +} + +func TestHTTPResponseReadsAheadAndCoalescesUnderOneOutstandingRPC(t *testing.T) { + body, writer := newTrackedResponseBody() + peer := newResponseProtocolPeer(t, body) + peer.startAndAck() + + if _, err := writer.Write([]byte("first")); err != nil { + t.Fatal(err) + } + first := peer.nextRequest() + if got := string(responseChunkData(t, first)); got != "first" { + t.Fatalf("first chunk = %q", got) + } + + writeDone := make(chan error, 1) + go func() { + for range 32 { + if _, err := writer.Write(make([]byte, 1024)); err != nil { + writeDone <- err + return + } + } + _, err := writer.Write([]byte("last")) + writeDone <- err + }() + + select { + case err := <-writeDone: + t.Fatalf("writer completed before read-ahead space was released: %v", err) + case <-time.After(50 * time.Millisecond): + } + if _, ok := peer.nextRequestWithin(50 * time.Millisecond); ok { + t.Fatal("sent another data RPC while the first acknowledgement was withheld") + } + + peer.ack(first) + combined := peer.nextRequest() + if got := len(responseChunkData(t, combined)); got != httpResponseReadAheadSize { + t.Fatalf("coalesced chunk size = %d, want %d", got, httpResponseReadAheadSize) + } + select { + case err := <-writeDone: + if err != nil { + t.Fatal(err) + } + case <-time.After(responseProtocolTimeout): + t.Fatal("source did not resume after acknowledgement") + } + if _, ok := peer.nextRequestWithin(50 * time.Millisecond); ok { + t.Fatal("sent another data RPC while the coalesced chunk acknowledgement was withheld") + } + + peer.ack(combined) + last := peer.nextRequest() + if got := string(responseChunkData(t, last)); got != "last" { + t.Fatalf("last chunk = %q", got) + } + if err := writer.Close(); err != nil { + t.Fatal(err) + } + peer.ack(last) + end := peer.nextRequest() + if params := responseChunkParams(t, end); params["end"] != true || params["error"] != nil { + t.Fatalf("unexpected terminal chunk: %s", end.Params) + } + peer.ack(end) +} + +func TestHTTPResponseFlushesPartialBytesImmediately(t *testing.T) { + body, writer := newTrackedResponseBody() + peer := newResponseProtocolPeer(t, body) + peer.startAndAck() + + if _, err := writer.Write([]byte{0xf0}); err != nil { + t.Fatal(err) + } + chunk := peer.nextRequest() + if got := responseChunkData(t, chunk); len(got) != 1 || got[0] != 0xf0 { + t.Fatalf("chunk = %v, want the single source byte", got) + } + peer.ack(chunk) + if err := writer.Close(); err != nil { + t.Fatal(err) + } + end := peer.nextRequest() + if params := responseChunkParams(t, end); params["end"] != true { + t.Fatalf("terminal end = %v, want true", params["end"]) + } + peer.ack(end) +} + +func TestHTTPResponseCompletionReleasesExchangeContext(t *testing.T) { + body, writer := newTrackedResponseBody() + peer := newResponseProtocolPeer(t, body) + peer.startAndAck() + + peer.adapter.mu.Lock() + exchange := peer.adapter.pending["test"] + peer.adapter.mu.Unlock() + if exchange == nil { + t.Fatal("pending exchange was not registered") + } + + if err := writer.Close(); err != nil { + t.Fatal(err) + } + end := peer.nextRequest() + if params := responseChunkParams(t, end); params["end"] != true || params["error"] != nil { + t.Fatalf("unexpected terminal chunk: %s", end.Params) + } + peer.ack(end) + + select { + case <-exchange.ctx.Done(): + case <-time.After(responseProtocolTimeout): + t.Fatal("completed exchange context was not released") + } +} + +func TestHTTPResponseUpstreamErrorFollowsBufferedBytesAndAcknowledgement(t *testing.T) { + body, writer := newTrackedResponseBody() + peer := newResponseProtocolPeer(t, body) + peer.startAndAck() + + if _, err := writer.Write([]byte("first")); err != nil { + t.Fatal(err) + } + first := peer.nextRequest() + if got := string(responseChunkData(t, first)); got != "first" { + t.Fatalf("first chunk = %q", got) + } + + sourceDone := make(chan struct{}) + go func() { + _, _ = writer.Write([]byte("partial")) + _ = writer.CloseWithError(errors.New("upstream failed")) + close(sourceDone) + }() + select { + case <-sourceDone: + case <-time.After(responseProtocolTimeout): + t.Fatal("source error was not read ahead") + } + if _, ok := peer.nextRequestWithin(50 * time.Millisecond); ok { + t.Fatal("upstream error overtook the outstanding data RPC") + } + + peer.ack(first) + partial := peer.nextRequest() + if got := string(responseChunkData(t, partial)); got != "partial" { + t.Fatalf("partial chunk = %q", got) + } + if _, ok := peer.nextRequestWithin(50 * time.Millisecond); ok { + t.Fatal("upstream error overtook buffered response bytes") + } + + peer.ack(partial) + terminal := peer.nextRequest() + params := responseChunkParams(t, terminal) + errorValue, ok := params["error"].(map[string]any) + if params["end"] != true || !ok || errorValue["message"] != "upstream failed" { + t.Fatalf("unexpected terminal error: %s", terminal.Params) + } + peer.ack(terminal) +} + +func TestHTTPResponseCancellationClosesSourceWithoutDataAcknowledgement(t *testing.T) { + body, writer := newTrackedResponseBody() + peer := newResponseProtocolPeer(t, body) + peer.startAndAck() + + if _, err := writer.Write([]byte("first")); err != nil { + t.Fatal(err) + } + first := peer.nextRequest() + if got := string(responseChunkData(t, first)); got != "first" { + t.Fatalf("first chunk = %q", got) + } + + peer.cancelRequest() + terminal := peer.nextRequest() + params := responseChunkParams(t, terminal) + errorValue, ok := params["error"].(map[string]any) + if params["end"] != true || !ok || errorValue["code"] != "cancelled" { + t.Fatalf("unexpected cancellation chunk: %s", terminal.Params) + } + waitForBodyClose(t, body) + peer.ack(terminal) +} + +func TestHTTPResponseCancellationSendsHeadBeforeTerminalError(t *testing.T) { + requestCancelled := make(chan struct{}) + peer := newResponseProtocolPeerWithTransport(t, roundTripFunc(func(request *http.Request) (*http.Response, error) { + <-request.Context().Done() + close(requestCancelled) + return nil, request.Context().Err() + })) + + peer.beginRequest() + peer.cancelRequest() + select { + case <-requestCancelled: + case <-time.After(responseProtocolTimeout): + t.Fatal("upstream request was not cancelled") + } + + head := peer.nextRequest() + if head.Method != "llmInference.httpResponseStart" { + t.Fatalf("first response method = %q, want httpResponseStart", head.Method) + } + var headParams map[string]any + if err := json.Unmarshal(head.Params, &headParams); err != nil { + t.Fatal(err) + } + if headParams["status"] != float64(499) { + t.Fatalf("response status = %v, want 499", headParams["status"]) + } + peer.ack(head) + + terminal := peer.nextRequest() + params := responseChunkParams(t, terminal) + errorValue, ok := params["error"].(map[string]any) + if params["end"] != true || !ok || errorValue["code"] != "cancelled" { + t.Fatalf("unexpected cancellation chunk: %s", terminal.Params) + } + peer.ack(terminal) +} + +func TestHTTPResponseRPCRejectionAndConnectionLossCloseSource(t *testing.T) { + t.Run("JSON-RPC rejection", func(t *testing.T) { + body, writer := newTrackedResponseBody() + peer := newResponseProtocolPeer(t, body) + peer.startAndAck() + + if _, err := writer.Write([]byte("first")); err != nil { + t.Fatal(err) + } + first := peer.nextRequest() + peer.reject(first, "write rejected") + waitForBodyClose(t, body) + + terminal := peer.nextRequest() + params := responseChunkParams(t, terminal) + errorValue, ok := params["error"].(map[string]any) + if params["end"] != true || !ok || !strings.Contains(errorValue["message"].(string), "write rejected") { + t.Fatalf("unexpected rejection terminal: %s", terminal.Params) + } + peer.ack(terminal) + }) + + t.Run("accepted false", func(t *testing.T) { + body, writer := newTrackedResponseBody() + peer := newResponseProtocolPeer(t, body) + peer.startAndAck() + + if _, err := writer.Write([]byte("first")); err != nil { + t.Fatal(err) + } + first := peer.nextRequest() + peer.respondAccepted(first, false) + waitForBodyClose(t, body) + + terminal := peer.nextRequest() + params := responseChunkParams(t, terminal) + errorValue, ok := params["error"].(map[string]any) + if params["end"] != true || !ok || !strings.Contains(errorValue["message"].(string), "was rejected") { + t.Fatalf("unexpected rejection terminal: %s", terminal.Params) + } + peer.ack(terminal) + }) + + t.Run("connection loss", func(t *testing.T) { + body, writer := newTrackedResponseBody() + peer := newResponseProtocolPeer(t, body) + peer.startAndAck() + + if _, err := writer.Write([]byte("first")); err != nil { + t.Fatal(err) + } + _ = peer.nextRequest() + if err := peer.conn.Close(); err != nil { + t.Fatal(err) + } + waitForBodyClose(t, body) + }) +} diff --git a/go/http_response_reader.go b/go/http_response_reader.go new file mode 100644 index 0000000000..ce0cb0db95 --- /dev/null +++ b/go/http_response_reader.go @@ -0,0 +1,147 @@ +package copilot + +import ( + "context" + "errors" + "io" + "runtime" + "sync" +) + +const httpResponseReadAheadSize = 32 * 1024 + +// httpResponseReader keeps one bounded buffer of response bytes ready while +// the preceding chunk waits for its runtime acknowledgement. +type httpResponseReader struct { + body io.ReadCloser + + mu sync.Mutex + buffered []byte + err error + + notify chan struct{} + space chan struct{} + stop chan struct{} + done chan struct{} + + closeOnce sync.Once +} + +func newHTTPResponseReader(body io.ReadCloser) *httpResponseReader { + r := &httpResponseReader{ + body: body, + buffered: make([]byte, 0, httpResponseReadAheadSize), + notify: make(chan struct{}, 1), + space: make(chan struct{}, 1), + stop: make(chan struct{}), + done: make(chan struct{}), + } + go r.readLoop() + return r +} + +func (r *httpResponseReader) readLoop() { + defer close(r.done) + scratch := make([]byte, httpResponseReadAheadSize) + for { + select { + case <-r.stop: + return + default: + } + + r.mu.Lock() + remaining := httpResponseReadAheadSize - len(r.buffered) + stopped := r.err != nil + r.mu.Unlock() + + if stopped { + return + } + if remaining == 0 { + select { + case <-r.space: + continue + case <-r.stop: + return + } + } + + n, err := r.body.Read(scratch[:remaining]) + r.mu.Lock() + if n > 0 { + r.buffered = append(r.buffered, scratch[:n]...) + } + if err != nil { + r.err = err + } + r.mu.Unlock() + + if n > 0 || err != nil { + signal(r.notify) + } + if err != nil { + return + } + if n == 0 { + // Broken readers are permitted to return (0, nil). Do not let one + // starve acknowledgement or cancellation processing. + runtime.Gosched() + } + } +} + +func (r *httpResponseReader) nextChunk(ctx context.Context, output *[]byte) (bool, error) { + for { + select { + case <-ctx.Done(): + return false, ctx.Err() + default: + } + + r.mu.Lock() + if len(r.buffered) > 0 { + replacement := (*output)[:0] + if cap(replacement) < httpResponseReadAheadSize { + replacement = make([]byte, 0, httpResponseReadAheadSize) + } + *output, r.buffered = r.buffered, replacement + r.mu.Unlock() + signal(r.space) + return true, nil + } + err := r.err + if err != nil { + r.err = nil + } + r.mu.Unlock() + + if err != nil { + if errors.Is(err, io.EOF) { + return false, nil + } + return false, err + } + + select { + case <-r.notify: + case <-ctx.Done(): + return false, ctx.Err() + } + } +} + +func (r *httpResponseReader) Close() { + r.closeOnce.Do(func() { + close(r.stop) + _ = r.body.Close() + <-r.done + }) +} + +func signal(ch chan struct{}) { + select { + case ch <- struct{}{}: + default: + } +} diff --git a/java/sdk/src/main/java/com/github/copilot/CopilotClient.java b/java/sdk/src/main/java/com/github/copilot/CopilotClient.java index ae41477c15..c7edfe65db 100644 --- a/java/sdk/src/main/java/com/github/copilot/CopilotClient.java +++ b/java/sdk/src/main/java/com/github/copilot/CopilotClient.java @@ -554,7 +554,6 @@ private Connection startCoreBody() { JsonRpcClient connectedRpc = rpc; Connection connection = new Connection(connectedRpc, process, new ServerRpc(connectedRpc::invoke), inProcessTransport == null ? null : inProcessTransport.host()); - connectedRpc.setCloseHandler(() -> sessions.values().forEach(CopilotSession::cancelPendingExternalTools)); // Register handlers for server-to-client calls RpcHandlerDispatcher dispatcher = new RpcHandlerDispatcher(sessions, lifecycleManager::dispatch, executor, @@ -564,11 +563,19 @@ private Connection startCoreBody() { // Register the LLM inference request handler when configured. com.github.copilot.CopilotRequestHandler requestHandler = this.options.getRequestHandler(); boolean hasLlmInference = requestHandler != null; + LlmInferenceAdapter llmAdapter = null; if (hasLlmInference) { - LlmInferenceAdapter llmAdapter = new LlmInferenceAdapter(requestHandler, - () -> connection.serverRpc().llmInference, executor); + llmAdapter = new LlmInferenceAdapter(requestHandler, () -> connection.serverRpc().llmInference, + executor); llmAdapter.registerHandlers(connectedRpc); } + LlmInferenceAdapter connectedLlmAdapter = llmAdapter; + connectedRpc.setCloseHandler(() -> { + sessions.values().forEach(CopilotSession::cancelPendingExternalTools); + if (connectedLlmAdapter != null) { + connectedLlmAdapter.cancelPending(); + } + }); // Register the GitHub telemetry forwarding handler when configured. Function> onGitHubTelemetry = this.options diff --git a/java/sdk/src/main/java/com/github/copilot/CopilotRequestHandler.java b/java/sdk/src/main/java/com/github/copilot/CopilotRequestHandler.java index 7b34b20e7e..d97a122a8d 100644 --- a/java/sdk/src/main/java/com/github/copilot/CopilotRequestHandler.java +++ b/java/sdk/src/main/java/com/github/copilot/CopilotRequestHandler.java @@ -16,6 +16,7 @@ import java.util.Set; import java.util.concurrent.CancellationException; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; /** * The base class for SDK consumers who want to observe or replace the LLM @@ -38,8 +39,6 @@ public class CopilotRequestHandler { private static final HttpClient SHARED_HTTP_CLIENT = HttpClient.newBuilder() .followRedirects(HttpClient.Redirect.NEVER).build(); - private static final int RESPONSE_CHUNK_SIZE = 32 * 1024; - static boolean isForbiddenRequestHeader(String name) { String lower = name.toLowerCase(Locale.ROOT); return FORBIDDEN_REQUEST_HEADERS.contains(lower) || lower.startsWith("sec-websocket-"); @@ -137,20 +136,57 @@ private static HttpRequest buildHttpRequest(LlmInferenceExchange exchange) throw private static void streamResponse(HttpResponse response, LlmInferenceExchange exchange) throws IOException { - exchange.startResponse(response.statusCode(), null, response.headers().map()); - try (InputStream body = response.body()) { - byte[] buffer = new byte[RESPONSE_CHUNK_SIZE]; - int n; - while ((n = body.read(buffer)) != -1) { - if (n > 0) { - exchange.writeResponseBinary(buffer, 0, n); + try (HttpResponseReader reader = new HttpResponseReader(response.body())) { + exchange.cancellation().thenRun(reader::cancel); + if (!awaitRpc(exchange.startResponseAsync(response.statusCode(), null, response.headers().map()), + exchange.cancellation())) { + exchange.errorResponse("Request cancelled by runtime", "cancelled"); + return; + } + + while (true) { + HttpResponseReader.Result result; + try { + result = reader.next(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IOException("Interrupted while reading HTTP response", e); + } + if (exchange.cancellation().isDone()) { + exchange.errorResponse("Request cancelled by runtime", "cancelled"); + return; + } + if (result.data() != null) { + if (!awaitRpc(exchange.writeResponseBinaryAsync(result.data()), exchange.cancellation())) { + exchange.errorResponse("Request cancelled by runtime", "cancelled"); + return; + } + } else if (result.error() != null) { + exchange.errorResponse(result.error().getMessage(), null); + return; + } else if (result.end()) { + exchange.endResponse(); + return; } } } catch (IOException e) { exchange.errorResponse(e.getMessage(), null); - return; } - exchange.endResponse(); + } + + private static boolean awaitRpc(CompletableFuture rpc, CompletableFuture cancellation) + throws IOException { + try { + CompletableFuture.anyOf(rpc, cancellation).join(); + if (cancellation.isDone()) { + return false; + } + rpc.join(); + return true; + } catch (CompletionException | CancellationException e) { + Throwable cause = e.getCause() != null ? e.getCause() : e; + throw new IOException(cause.getMessage(), cause); + } } private void handleWebSocket(LlmInferenceExchange exchange) throws Exception { diff --git a/java/sdk/src/main/java/com/github/copilot/HttpResponseReader.java b/java/sdk/src/main/java/com/github/copilot/HttpResponseReader.java new file mode 100644 index 0000000000..29654be122 --- /dev/null +++ b/java/sdk/src/main/java/com/github/copilot/HttpResponseReader.java @@ -0,0 +1,156 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +package com.github.copilot; + +import java.io.IOException; +import java.io.InputStream; +import java.io.UncheckedIOException; +import java.util.Arrays; + +/** + * Reads a blocking HTTP response body on a daemon thread and keeps at most one + * response chunk buffered ahead of the RPC writer. + */ +final class HttpResponseReader implements AutoCloseable { + + static final int CHUNK_SIZE = 32 * 1024; + + // Match the read-ahead bound so a source fragment as large as the bound is + // taken in one read; a smaller scratch would split it into several reads + // and several source round-trips for no benefit. + private static final int READ_SIZE = CHUNK_SIZE; + + record Result(byte[] data, IOException error, boolean end) { + } + + private final InputStream source; + private final byte[] buffered = new byte[CHUNK_SIZE]; + private final Thread thread; + + private int bufferedCount; + private boolean closed; + private boolean end; + private IOException error; + + HttpResponseReader(InputStream source) { + this.source = source; + thread = new Thread(this::readLoop, "llm-http-response-reader"); + thread.setDaemon(true); + thread.start(); + } + + synchronized Result next() throws InterruptedException { + while (bufferedCount == 0 && !end && error == null && !closed) { + wait(); + } + if (bufferedCount > 0) { + byte[] data = Arrays.copyOf(buffered, bufferedCount); + bufferedCount = 0; + notifyAll(); + return new Result(data, null, false); + } + if (error != null) { + IOException result = error; + error = null; + return new Result(null, result, false); + } + return new Result(null, null, true); + } + + private void readLoop() { + byte[] readBuffer = new byte[READ_SIZE]; + try { + while (true) { + int capacity; + synchronized (this) { + while (bufferedCount == CHUNK_SIZE && !closed) { + wait(); + } + if (closed) { + return; + } + capacity = Math.min(readBuffer.length, CHUNK_SIZE - bufferedCount); + } + + int count = source.read(readBuffer, 0, capacity); + if (count < 0) { + synchronized (this) { + end = true; + notifyAll(); + } + return; + } + if (count == 0) { + Thread.yield(); + continue; + } + + synchronized (this) { + if (closed) { + return; + } + System.arraycopy(readBuffer, 0, buffered, bufferedCount, count); + bufferedCount += count; + notifyAll(); + } + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } catch (IOException e) { + synchronized (this) { + if (!closed) { + error = e; + notifyAll(); + } + } + } catch (RuntimeException e) { + synchronized (this) { + if (!closed) { + error = toIOException(e); + notifyAll(); + } + } + } finally { + synchronized (this) { + end = true; + notifyAll(); + } + } + } + + private static IOException toIOException(RuntimeException error) { + if (error instanceof UncheckedIOException unchecked) { + IOException cause = unchecked.getCause(); + return cause.getMessage() != null ? cause : new IOException(unchecked.toString(), cause); + } + return new IOException(error.getMessage() != null ? error.getMessage() : error.toString(), error); + } + + void cancel() { + synchronized (this) { + if (closed) { + return; + } + closed = true; + notifyAll(); + } + try { + source.close(); + } catch (IOException ignored) { + // A pending read is already being abandoned. + } + thread.interrupt(); + } + + @Override + public void close() { + cancel(); + try { + thread.join(1000); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } +} diff --git a/java/sdk/src/main/java/com/github/copilot/JsonRpcClient.java b/java/sdk/src/main/java/com/github/copilot/JsonRpcClient.java index f78ce00425..2258963b14 100644 --- a/java/sdk/src/main/java/com/github/copilot/JsonRpcClient.java +++ b/java/sdk/src/main/java/com/github/copilot/JsonRpcClient.java @@ -141,7 +141,13 @@ public CompletableFuture invoke(String method, Object params, Class re long timingNanos = System.nanoTime(); long id = requestIdCounter.incrementAndGet(); var future = new CompletableFuture(); - pendingRequests.put(id, future); + synchronized (closeHandlerLock) { + if (closeNotified) { + future.completeExceptionally(new IOException("Client closed")); + } else { + pendingRequests.put(id, future); + } + } var request = new JsonRpcRequest(); request.setJsonrpc("2.0"); @@ -149,11 +155,13 @@ public CompletableFuture invoke(String method, Object params, Class re request.setMethod(method); request.setParams(params); - try { - sendMessage(request); - } catch (IOException e) { - pendingRequests.remove(id); - future.completeExceptionally(e); + if (!future.isDone()) { + try { + sendMessage(request); + } catch (IOException e) { + pendingRequests.remove(id); + future.completeExceptionally(e); + } } return future.thenApply(result -> { @@ -350,6 +358,8 @@ private void notifyClose() { } closeNotified = true; handler = closeHandler; + pendingRequests.forEach((id, future) -> future.completeExceptionally(new IOException("Client closed"))); + pendingRequests.clear(); } if (handler != null) { try { @@ -426,10 +436,6 @@ public void close() { readerExecutor.shutdownNow(); notifyClose(); - // Cancel all pending requests - pendingRequests.forEach((id, future) -> future.completeExceptionally(new IOException("Client closed"))); - pendingRequests.clear(); - try { if (socket != null) { socket.close(); diff --git a/java/sdk/src/main/java/com/github/copilot/LlmInferenceAdapter.java b/java/sdk/src/main/java/com/github/copilot/LlmInferenceAdapter.java index 3e741a56bc..dc00839ddf 100644 --- a/java/sdk/src/main/java/com/github/copilot/LlmInferenceAdapter.java +++ b/java/sdk/src/main/java/com/github/copilot/LlmInferenceAdapter.java @@ -52,6 +52,11 @@ void registerHandlers(JsonRpcClient rpc) { (rpcId, params) -> handleRequestChunk(rpc, rpcId, params)); } + void cancelPending() { + pending.values().forEach(LlmInferenceExchange::pushCancel); + pending.clear(); + } + private LlmInferenceExchange getOrCreateExchange(String requestId) { // The runtime dispatches httpRequestStart and httpRequestChunk frames // independently. Even though the current reader dispatches them in diff --git a/java/sdk/src/main/java/com/github/copilot/LlmInferenceExchange.java b/java/sdk/src/main/java/com/github/copilot/LlmInferenceExchange.java index 67933e40b6..bf7eab2f88 100644 --- a/java/sdk/src/main/java/com/github/copilot/LlmInferenceExchange.java +++ b/java/sdk/src/main/java/com/github/copilot/LlmInferenceExchange.java @@ -177,6 +177,11 @@ byte[] drainBody() throws InterruptedException { // --- Response emit (driven by the handler) --- void startResponse(int status, String statusText, Map> headers) throws IOException { + join(startResponseAsync(status, statusText, headers)); + } + + CompletableFuture startResponseAsync(int status, String statusText, Map> headers) + throws IOException { synchronized (lock) { if (started) { throw new IOException("LLM inference response startResponse() called twice"); @@ -187,7 +192,7 @@ void startResponse(int status, String statusText, Map> head started = true; } var params = new LlmInferenceHttpResponseStartParams(requestId, (long) status, statusText, headers); - join(api().httpResponseStart(params)); + return api().httpResponseStart(params).thenApply(ignored -> null); } void writeResponseText(String text) throws IOException { @@ -203,6 +208,10 @@ void writeResponseBinary(byte[] data, int offset, int length) throws IOException writeChunk(new String(encoded.array(), 0, encoded.limit(), StandardCharsets.ISO_8859_1), true); } + CompletableFuture writeResponseBinaryAsync(byte[] data) throws IOException { + return writeChunkAsync(Base64.getEncoder().encodeToString(data), true); + } + void endResponse() throws IOException { synchronized (lock) { if (finished) { @@ -227,6 +236,10 @@ void errorResponse(String message, String code) throws IOException { } private void writeChunk(String data, boolean binary) throws IOException { + join(writeChunkAsync(data, binary)); + } + + private CompletableFuture writeChunkAsync(String data, boolean binary) throws IOException { synchronized (lock) { if (cancelled) { throw new IOException("LLM inference request was cancelled by the runtime"); @@ -241,7 +254,7 @@ private void writeChunk(String data, boolean binary) throws IOException { } var params = new LlmInferenceHttpResponseChunkParams(requestId, data, binary ? Boolean.TRUE : null, Boolean.FALSE, null); - join(api().httpResponseChunk(params)); + return api().httpResponseChunk(params).thenApply(ignored -> null); } private ServerLlmInferenceApi api() throws IOException { diff --git a/java/sdk/src/test/java/com/github/copilot/HttpResponseForwardingTest.java b/java/sdk/src/test/java/com/github/copilot/HttpResponseForwardingTest.java new file mode 100644 index 0000000000..52607a4950 --- /dev/null +++ b/java/sdk/src/test/java/com/github/copilot/HttpResponseForwardingTest.java @@ -0,0 +1,492 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +package com.github.copilot; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +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.assertTrue; + +import java.io.IOException; +import java.io.InputStream; +import java.io.UncheckedIOException; +import java.net.URI; +import java.net.http.HttpClient; +import java.net.http.HttpHeaders; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.util.Arrays; +import java.util.Base64; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; +import javax.net.ssl.SSLSession; + +import org.junit.jupiter.api.Test; + +import com.github.copilot.generated.rpc.LlmInferenceHttpResponseChunkParams; +import com.github.copilot.generated.rpc.RpcCaller; +import com.github.copilot.generated.rpc.ServerRpc; + +class HttpResponseForwardingTest { + + private static final Duration DEADLINE = Duration.ofSeconds(5); + + @Test + void readsAheadAndCoalescesUnderOneWithheldAck() throws Exception { + try (TestFlow flow = new TestFlow()) { + flow.source.feed("first"); + flow.ack(flow.next("llmInference.httpResponseStart")); + + PendingCall first = flow.nextData(); + assertArrayEquals(bytes("first"), data(first)); + + byte[] expected = new byte[HttpResponseReader.CHUNK_SIZE]; + for (int i = 0; i < 32; i++) { + byte[] fragment = new byte[1024]; + Arrays.fill(fragment, (byte) i); + System.arraycopy(fragment, 0, expected, i * fragment.length, fragment.length); + flow.source.feed(fragment); + } + flow.source.awaitBytesRead(5 + HttpResponseReader.CHUNK_SIZE); + flow.source.feed("last"); + + assertNull(flow.poll(), "No second data RPC may overtake the outstanding write"); + assertEquals(1, flow.caller.outstandingData.get()); + assertEquals(1, flow.caller.maximumOutstandingData.get()); + + flow.ack(first); + PendingCall combined = flow.nextData(); + assertArrayEquals(expected, data(combined)); + assertEquals(1, flow.caller.outstandingData.get()); + + flow.source.finish(); + flow.ack(combined); + PendingCall last = flow.nextData(); + assertArrayEquals(bytes("last"), data(last)); + flow.ack(last); + + PendingCall end = flow.next("llmInference.httpResponseChunk"); + assertTrue(chunk(end).end()); + assertNull(chunk(end).error()); + flow.ack(end); + flow.awaitComplete(); + assertEquals(1, flow.caller.maximumOutstandingData.get()); + } + } + + @Test + void flushesPartialBytesWithoutWaitingForFutureInput() throws Exception { + try (TestFlow flow = new TestFlow()) { + flow.ack(flow.next("llmInference.httpResponseStart")); + flow.source.feed("partial"); + + PendingCall partial = flow.nextData(); + assertArrayEquals(bytes("partial"), data(partial)); + + flow.source.finish(); + flow.ack(partial); + PendingCall end = flow.next("llmInference.httpResponseChunk"); + assertTrue(chunk(end).end()); + flow.ack(end); + flow.awaitComplete(); + } + } + + @Test + void upstreamErrorFollowsBufferedBytesAndOutstandingAck() throws Exception { + try (TestFlow flow = new TestFlow()) { + flow.source.feed("first"); + flow.ack(flow.next("llmInference.httpResponseStart")); + PendingCall first = flow.nextData(); + + flow.source.feed("partial"); + flow.source.fail(new IOException("upstream failed")); + flow.source.awaitFailureRead(); + assertNull(flow.poll(), "Buffered bytes must not overtake the outstanding write"); + + flow.ack(first); + PendingCall partial = flow.nextData(); + assertArrayEquals(bytes("partial"), data(partial)); + assertNull(flow.poll(), "The upstream error must follow the buffered bytes"); + + flow.ack(partial); + PendingCall error = flow.next("llmInference.httpResponseChunk"); + assertTrue(chunk(error).end()); + assertEquals("upstream failed", chunk(error).error().message()); + flow.ack(error); + flow.awaitComplete(); + assertTrue(flow.source.closed.get()); + } + } + + @Test + void uncheckedUpstreamErrorFollowsBufferedBytesAndOutstandingAck() throws Exception { + try (TestFlow flow = new TestFlow()) { + flow.source.feed("first"); + flow.ack(flow.next("llmInference.httpResponseStart")); + PendingCall first = flow.nextData(); + + flow.source.feed("partial"); + flow.source.fail(new UncheckedIOException(new IOException("unchecked upstream failed"))); + flow.source.awaitFailureRead(); + assertNull(flow.poll(), "Buffered bytes must not overtake the outstanding write"); + + flow.ack(first); + PendingCall partial = flow.nextData(); + assertArrayEquals(bytes("partial"), data(partial)); + assertNull(flow.poll(), "The upstream error must follow the buffered bytes"); + + flow.ack(partial); + PendingCall error = flow.next("llmInference.httpResponseChunk"); + assertTrue(chunk(error).end()); + assertEquals("unchecked upstream failed", chunk(error).error().message()); + flow.ack(error); + flow.awaitComplete(); + assertTrue(flow.source.closed.get()); + } + } + + @Test + void messageLessUncheckedUpstreamErrorHasTerminalMessage() throws Exception { + try (TestFlow flow = new TestFlow()) { + flow.ack(flow.next("llmInference.httpResponseStart")); + flow.source.fail(new IllegalStateException()); + flow.source.awaitFailureRead(); + + PendingCall error = flow.next("llmInference.httpResponseChunk"); + assertTrue(chunk(error).end()); + assertTrue(chunk(error).error().message().contains(IllegalStateException.class.getName())); + flow.ack(error); + flow.awaitComplete(); + assertTrue(flow.source.closed.get()); + } + } + + @Test + void cancellationClosesSourceWithoutWaitingForDataAck() throws Exception { + try (TestFlow flow = new TestFlow()) { + flow.source.feed("first"); + flow.ack(flow.next("llmInference.httpResponseStart")); + PendingCall first = flow.nextData(); + + flow.exchange.pushCancel(); + flow.source.awaitClosed(); + + PendingCall cancelled = flow.next("llmInference.httpResponseChunk"); + assertTrue(chunk(cancelled).end()); + assertEquals("cancelled", chunk(cancelled).error().code()); + flow.ack(cancelled); + flow.awaitComplete(); + flow.ack(first); + } + } + + @Test + void cancellationClosesSourceWithoutWaitingForHeadAck() throws Exception { + try (TestFlow flow = new TestFlow()) { + PendingCall head = flow.next("llmInference.httpResponseStart"); + + flow.exchange.pushCancel(); + flow.source.awaitClosed(); + + PendingCall cancelled = flow.next("llmInference.httpResponseChunk"); + assertTrue(chunk(cancelled).end()); + assertEquals("cancelled", chunk(cancelled).error().code()); + flow.ack(cancelled); + flow.awaitComplete(); + flow.ack(head); + } + } + + @Test + void rpcRejectionClosesSourceAndReportsError() throws Exception { + assertRpcFailureClosesSource(new JsonRpcException(-32603, "write rejected"), "write rejected"); + } + + @Test + void connectionLossClosesSourceAndReportsError() throws Exception { + assertRpcFailureClosesSource(new IOException("Client closed"), "Client closed"); + } + + private static void assertRpcFailureClosesSource(Exception failure, String expectedMessage) throws Exception { + try (TestFlow flow = new TestFlow()) { + flow.source.feed("first"); + flow.ack(flow.next("llmInference.httpResponseStart")); + PendingCall first = flow.nextData(); + + flow.reject(first, failure); + flow.source.awaitClosed(); + + PendingCall error = flow.next("llmInference.httpResponseChunk"); + assertTrue(chunk(error).end()); + assertTrue(chunk(error).error().message().contains(expectedMessage)); + flow.ack(error); + flow.awaitComplete(); + } + } + + private static byte[] data(PendingCall call) { + LlmInferenceHttpResponseChunkParams chunk = chunk(call); + assertFalse(chunk.end()); + assertEquals(Boolean.TRUE, chunk.binary()); + return Base64.getDecoder().decode(chunk.data()); + } + + private static LlmInferenceHttpResponseChunkParams chunk(PendingCall call) { + return (LlmInferenceHttpResponseChunkParams) call.params(); + } + + private static byte[] bytes(String value) { + return value.getBytes(StandardCharsets.UTF_8); + } + + private record PendingCall(String method, Object params, CompletableFuture result) { + } + + private static final class RecordingCaller implements RpcCaller { + + private final BlockingQueue calls = new LinkedBlockingQueue<>(); + private final AtomicInteger outstandingData = new AtomicInteger(); + private final AtomicInteger maximumOutstandingData = new AtomicInteger(); + + @Override + public CompletableFuture invoke(String method, Object params, Class resultType) { + CompletableFuture result = new CompletableFuture<>(); + if (params instanceof LlmInferenceHttpResponseChunkParams chunk && !chunk.end()) { + int outstanding = outstandingData.incrementAndGet(); + maximumOutstandingData.accumulateAndGet(outstanding, Math::max); + result.whenComplete((ignored, error) -> outstandingData.decrementAndGet()); + } + calls.add(new PendingCall(method, params, result)); + @SuppressWarnings("unchecked") + CompletableFuture typed = (CompletableFuture) (CompletableFuture) result; + return typed; + } + } + + private static final class TestFlow implements AutoCloseable { + + private final QueueInputStream source = new QueueInputStream(); + private final RecordingCaller caller = new RecordingCaller(); + private final LlmInferenceExchange exchange; + private final CompletableFuture task; + + TestFlow() { + ServerRpc rpc = new ServerRpc(caller); + exchange = new LlmInferenceExchange("test", () -> rpc.llmInference); + exchange.setMethod("GET"); + exchange.setContext(new CopilotRequestContext("test", null, null, null, null, CopilotRequestTransport.HTTP, + "http://unused.test", Map.of(), exchange.cancellation())); + exchange.pushEnd(); + + CopilotRequestHandler handler = new CopilotRequestHandler() { + @Override + protected HttpResponse sendRequest(HttpRequest request, CopilotRequestContext context) { + return new StubResponse(source, request); + } + }; + task = CompletableFuture.runAsync(() -> { + try { + handler.handle(exchange); + } catch (Exception e) { + throw new RuntimeException(e); + } + }); + } + + PendingCall next(String method) throws Exception { + PendingCall call = caller.calls.poll(DEADLINE.toMillis(), TimeUnit.MILLISECONDS); + assertNotNull(call, "Timed out waiting for " + method); + assertEquals(method, call.method()); + return call; + } + + PendingCall nextData() throws Exception { + PendingCall call = next("llmInference.httpResponseChunk"); + assertFalse(chunk(call).end()); + return call; + } + + PendingCall poll() throws InterruptedException { + return caller.calls.poll(100, TimeUnit.MILLISECONDS); + } + + void ack(PendingCall call) { + call.result().complete(null); + } + + void reject(PendingCall call, Exception error) { + call.result().completeExceptionally(error); + } + + void awaitComplete() throws Exception { + task.get(DEADLINE.toMillis(), TimeUnit.MILLISECONDS); + } + + @Override + public void close() throws Exception { + source.close(); + if (!task.isDone()) { + exchange.pushCancel(); + PendingCall call; + while ((call = caller.calls.poll()) != null) { + call.result().complete(null); + } + } + task.get(DEADLINE.toMillis(), TimeUnit.MILLISECONDS); + } + } + + private static final class QueueInputStream extends InputStream { + + private static final Object END = new Object(); + + private final BlockingQueue items = new LinkedBlockingQueue<>(); + private final AtomicLong bytesRead = new AtomicLong(); + private final AtomicBoolean closed = new AtomicBoolean(); + private final CompletableFuture failureRead = new CompletableFuture<>(); + private final CompletableFuture closeObserved = new CompletableFuture<>(); + private final AtomicReference current = new AtomicReference<>(); + private int currentOffset; + + void feed(String value) { + feed(bytes(value)); + } + + void feed(byte[] value) { + items.add(value); + } + + void finish() { + items.add(END); + } + + void fail(IOException error) { + items.add(error); + } + + void fail(RuntimeException error) { + items.add(error); + } + + void awaitBytesRead(long expected) throws Exception { + long deadline = System.nanoTime() + DEADLINE.toNanos(); + while (bytesRead.get() < expected && System.nanoTime() < deadline) { + Thread.sleep(5); + } + assertEquals(expected, bytesRead.get()); + } + + void awaitFailureRead() throws Exception { + failureRead.get(DEADLINE.toMillis(), TimeUnit.MILLISECONDS); + } + + void awaitClosed() throws Exception { + closeObserved.get(DEADLINE.toMillis(), TimeUnit.MILLISECONDS); + } + + @Override + public int read() throws IOException { + byte[] single = new byte[1]; + int count = read(single, 0, 1); + return count < 0 ? -1 : Byte.toUnsignedInt(single[0]); + } + + @Override + public int read(byte[] buffer, int offset, int length) throws IOException { + while (true) { + if (closed.get()) { + throw new IOException("stream closed"); + } + byte[] data = current.get(); + if (data != null) { + int count = Math.min(length, data.length - currentOffset); + System.arraycopy(data, currentOffset, buffer, offset, count); + currentOffset += count; + if (currentOffset == data.length) { + current.set(null); + currentOffset = 0; + } + bytesRead.addAndGet(count); + return count; + } + + Object item; + try { + item = items.take(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IOException("interrupted", e); + } + if (item == END) { + return -1; + } + if (item instanceof IOException e) { + failureRead.complete(null); + throw e; + } + if (item instanceof RuntimeException e) { + failureRead.complete(null); + throw e; + } + current.set((byte[]) item); + } + } + + @Override + public void close() { + if (closed.compareAndSet(false, true)) { + items.add(END); + closeObserved.complete(null); + } + } + } + + private record StubResponse(InputStream body, HttpRequest request) implements HttpResponse { + + @Override + public int statusCode() { + return 200; + } + + @Override + public Optional> previousResponse() { + return Optional.empty(); + } + + @Override + public HttpHeaders headers() { + return HttpHeaders.of(Map.of("content-type", List.of("application/octet-stream")), (name, value) -> true); + } + + @Override + public Optional sslSession() { + return Optional.empty(); + } + + @Override + public URI uri() { + return request.uri(); + } + + @Override + public HttpClient.Version version() { + return HttpClient.Version.HTTP_1_1; + } + } +} diff --git a/java/sdk/src/test/java/com/github/copilot/JsonRpcClientTest.java b/java/sdk/src/test/java/com/github/copilot/JsonRpcClientTest.java index 13668cd97f..85faba7c76 100644 --- a/java/sdk/src/test/java/com/github/copilot/JsonRpcClientTest.java +++ b/java/sdk/src/test/java/com/github/copilot/JsonRpcClientTest.java @@ -215,6 +215,38 @@ void testCloseHandlerRunsWhenRegisteredAfterClose() throws Exception { } } + @Test + void testInvokeFailsAfterRemoteClose() throws Exception { + try (var pair = createSocketPair()) { + var closed = new CompletableFuture(); + pair.client.setCloseHandler(() -> closed.complete(null)); + + pair.serverSide.close(); + closed.get(5, TimeUnit.SECONDS); + + var future = pair.client.invoke("test", Map.of(), JsonNode.class); + var ex = assertThrows(ExecutionException.class, () -> future.get(5, TimeUnit.SECONDS)); + assertInstanceOf(IOException.class, ex.getCause()); + } + } + + @Test + void testPendingInvokeFailsBeforeRemoteCloseHandlerRuns() throws Exception { + try (var pair = createSocketPair()) { + var future = pair.client.invoke("test", Map.of(), JsonNode.class); + readRpcMessage(pair.serverSide.getInputStream()); + + var closeHandlerObservedFailure = new CompletableFuture(); + pair.client.setCloseHandler(() -> closeHandlerObservedFailure.complete(future.isCompletedExceptionally())); + + pair.serverSide.close(); + + assertTrue(closeHandlerObservedFailure.get(5, TimeUnit.SECONDS)); + var ex = assertThrows(ExecutionException.class, () -> future.get(5, TimeUnit.SECONDS)); + assertInstanceOf(IOException.class, ex.getCause()); + } + } + // ---- invoke() edge cases ---- @Test diff --git a/nodejs/src/client.ts b/nodejs/src/client.ts index 4b610888b6..39a637a31d 100644 --- a/nodejs/src/client.ts +++ b/nodejs/src/client.ts @@ -420,6 +420,7 @@ export class CopilotClient { private cliProcess: ChildProcess | null = null; private ffiHost: FfiRuntimeHost | null = null; private connection: MessageConnection | null = null; + private requestAdapter: ReturnType | null = null; private messageWriter: TeardownResilientStreamMessageWriter | null = null; private connectionClosed: boolean = false; private socket: Socket | null = null; @@ -816,13 +817,14 @@ export class CopilotClient { const handlers: import("./generated/rpc.js").ClientGlobalApiHandlers = {}; handlers.extensionLaunchProvider = this.extensionLaunchProvider; if (this.requestHandler) { - handlers.llmInference = createCopilotRequestAdapter(this.requestHandler, () => { + this.requestAdapter = createCopilotRequestAdapter(this.requestHandler, () => { if (!this.connection) { return undefined; } this._rpc ??= createServerRpc(this.connection); return this._rpc; }); + handlers.llmInference = this.requestAdapter; } if (this.onGitHubTelemetry) { const onGitHubTelemetry = this.onGitHubTelemetry; @@ -1069,6 +1071,7 @@ export class CopilotClient { } this.sessions.clear(); this.githubTokenProviders.clear(); + this.requestAdapter?.cancelPending(); // Ask SDK-owned runtimes to flush and clean up before we tear down // their transport/process. External runtimes may be shared, so only @@ -1257,6 +1260,7 @@ export class CopilotClient { } this.sessions.clear(); this.githubTokenProviders.clear(); + this.requestAdapter?.cancelPending(); // Force close connection. Suppress writer failures first so teardown // write rejections don't surface as unhandled rejections. @@ -3077,6 +3081,7 @@ export class CopilotClient { } this.sessions.clear(); this.githubTokenProviders.clear(); + this.requestAdapter?.cancelPending(); }; this.connection.onClose(markDisconnected); this.connection.onError(() => { diff --git a/nodejs/src/copilotRequestHandler.ts b/nodejs/src/copilotRequestHandler.ts index ccfe6591c6..d36deb5e3d 100644 --- a/nodejs/src/copilotRequestHandler.ts +++ b/nodejs/src/copilotRequestHandler.ts @@ -16,6 +16,7 @@ type ServerRpc = ReturnType; const sharedTextDecoder = new TextDecoder("utf-8", { fatal: false }); const sharedTextEncoder = new TextEncoder(); +const HTTP_RESPONSE_READ_AHEAD_BYTES = 32 * 1024; const kBridge = Symbol("copilotWebSocketResponseBridge"); const kCompletion = Symbol("copilotWebSocketCompletion"); @@ -341,7 +342,7 @@ export class CopilotRequestHandler { export function createCopilotRequestAdapter( handler: CopilotRequestHandler, getServerRpc: () => ServerRpc | undefined -): LlmInferenceHandler { +): LlmInferenceHandler & { cancelPending(): void } { const pending = new Map(); function getOrCreate(requestId: string): CopilotRequestExchange { @@ -378,7 +379,9 @@ export function createCopilotRequestAdapter( const message = err instanceof Error ? err.message : String(err); await finalize(exchange, 502, message); } finally { - pending.delete(exchange.requestId); + if (pending.get(exchange.requestId) === exchange) { + pending.delete(exchange.requestId); + } } } @@ -401,6 +404,13 @@ export function createCopilotRequestAdapter( routeChunk(getOrCreate(params.requestId), params); return {}; }, + cancelPending(): void { + const exchanges = [...pending.values()]; + pending.clear(); + for (const exchange of exchanges) { + exchange.pushCancel("RPC connection closed"); + } + }, }; } @@ -708,32 +718,189 @@ async function drainAsync(stream: AsyncIterable): Promise { - await exchange.startResponse({ - status: response.status, - statusText: response.statusText || undefined, - headers: headersToMultiMap(response.headers), - }); + const reader = response.body + ? new HttpResponseReader(response.body, exchange.signal) + : undefined; + try { + await exchange.startResponse({ + status: response.status, + statusText: response.statusText || undefined, + headers: headersToMultiMap(response.headers), + }); + + if (reader) { + for (;;) { + const chunk = await reader.nextChunk(); + if (!chunk) { + break; + } + await exchange.writeResponse(chunk); + } + } - const body = response.body; - if (!body) { await exchange.endResponse(); - return; + } finally { + await reader?.dispose(); } +} - const reader = body.getReader(); - try { - for (;;) { - const { value, done } = await reader.read(); - if (done) { - break; +/** + * Pulls independently of response RPC acknowledgements, bounding how far ahead + * of the runtime the source may run. Frames are held by reference and only + * copied when more than one has accumulated, so a consumer that keeps up pays + * no copy at all. + */ +class HttpResponseReader { + readonly #reader: ReadableStreamDefaultReader; + readonly #signal: AbortSignal; + readonly #frames: Uint8Array[] = []; + readonly #pump: Promise; + readonly #onAbort: () => void; + #queued = 0; + #done = false; + #cancelled = false; + #hasError = false; + #error: unknown; + #dataWaker: (() => void) | undefined; + #spaceWaker: (() => void) | undefined; + #cancelPromise: Promise | undefined; + + constructor(body: ReadableStream, signal: AbortSignal) { + this.#reader = body.getReader(); + this.#signal = signal; + this.#onAbort = () => { + void this.cancel(signal.reason); + }; + signal.addEventListener("abort", this.#onAbort, { once: true }); + this.#pump = this.#pumpSource(); + if (signal.aborted) { + this.#onAbort(); + } + } + + async nextChunk(): Promise { + while (this.#frames.length === 0 && !this.#done && !this.#hasError) { + await new Promise((resolve) => { + this.#dataWaker = resolve; + }); + } + + if (this.#frames.length > 0) { + const chunk = this.#takeQueued(); + this.#wakeSpace(); + return chunk; + } + + if (this.#hasError) { + this.#hasError = false; + throw this.#error; + } + + return undefined; + } + + #takeQueued(): Uint8Array { + const frames = this.#frames; + const total = this.#queued; + this.#queued = 0; + + // The consumer kept up, so a single frame is waiting: forward the + // source's own buffer instead of copying it through a staging buffer. + if (frames.length === 1) { + return frames.pop()!; + } + + const chunk = new Uint8Array(total); + let offset = 0; + for (const frame of frames) { + chunk.set(frame, offset); + offset += frame.byteLength; + } + frames.length = 0; + return chunk; + } + + async dispose(): Promise { + this.#signal.removeEventListener("abort", this.#onAbort); + if (!this.#done) { + await this.cancel(); + } else if (this.#cancelPromise) { + await this.#cancelPromise; + } + await this.#pump; + } + + cancel(reason?: unknown): Promise { + if (this.#done) { + return Promise.resolve(); + } + if (!this.#cancelPromise) { + this.#cancelled = true; + this.#done = true; + this.#hasError = true; + this.#error = + reason instanceof Error ? reason : new Error("HTTP response body cancelled."); + this.#frames.length = 0; + this.#queued = 0; + this.#wakeData(); + this.#wakeSpace(); + this.#cancelPromise = this.#reader.cancel(reason).catch(() => undefined); + } + return this.#cancelPromise; + } + + async #pumpSource(): Promise { + try { + while (!this.#cancelled) { + while (this.#queued >= HTTP_RESPONSE_READ_AHEAD_BYTES && !this.#cancelled) { + await new Promise((resolve) => { + this.#spaceWaker = resolve; + }); + } + if (this.#cancelled) { + return; + } + + const { value, done } = await this.#reader.read(); + if (done) { + return; + } + if (value.byteLength > 0) { + this.#frames.push(value); + this.#queued += value.byteLength; + this.#wakeData(); + continue; + } + + // A source that only ever yields empty frames would never reach + // the read-ahead bound, so yield explicitly to keep timers and + // socket I/O — including the acknowledgements this loop is + // racing — from being starved by the microtask queue. + await new Promise((resolve) => setImmediate(resolve)); } - if (value && value.byteLength > 0) { - await exchange.writeResponse(value); + } catch (error) { + if (!this.#cancelled) { + this.#hasError = true; + this.#error = error; } + } finally { + this.#done = true; + this.#wakeData(); + this.#wakeSpace(); + this.#reader.releaseLock(); } - await exchange.endResponse(); - } finally { - reader.releaseLock(); + } + + #wakeData(): void { + const waker = this.#dataWaker; + this.#dataWaker = undefined; + waker?.(); + } + + #wakeSpace(): void { + const waker = this.#spaceWaker; + this.#spaceWaker = undefined; + waker?.(); } } diff --git a/nodejs/test/client.test.ts b/nodejs/test/client.test.ts index f8f3848b11..8175b4b5ab 100644 --- a/nodejs/test/client.test.ts +++ b/nodejs/test/client.test.ts @@ -4286,6 +4286,19 @@ describe("CopilotClient", () => { }); describe("shutdown", () => { + it.each(["stop", "forceStop"] as const)( + "%s cancels pending inference requests before transport teardown", + async (method) => { + const client = new CopilotClient({ autoStart: false }); + const cancelPending = vi.fn(); + (client as any).requestAdapter = { cancelPending }; + + await client[method](); + + expect(cancelPending).toHaveBeenCalledTimes(1); + } + ); + it.each(["stop", "forceStop"] as const)( "%s waits for the initial in-process cleanup attempt", async (method) => { diff --git a/nodejs/test/copilot-request-handler.test.ts b/nodejs/test/copilot-request-handler.test.ts new file mode 100644 index 0000000000..ad5b5a6962 --- /dev/null +++ b/nodejs/test/copilot-request-handler.test.ts @@ -0,0 +1,368 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +import { describe, expect, it } from "vitest"; +import { + CopilotRequestHandler, + type CopilotRequestContext, + createCopilotRequestAdapter, +} from "../src/copilotRequestHandler.js"; +import type { + LlmInferenceHandler, + LlmInferenceHttpResponseChunkRequest, + LlmInferenceHttpResponseStartRequest, +} from "../src/generated/rpc.js"; + +const READ_AHEAD_BYTES = 32 * 1024; + +interface Deferred { + promise: Promise; + resolve: (value: T) => void; + reject: (reason?: unknown) => void; +} + +function deferred(): Deferred { + let resolve!: (value: T) => void; + let reject!: (reason?: unknown) => void; + const promise = new Promise((resolvePromise, rejectPromise) => { + resolve = resolvePromise; + reject = rejectPromise; + }); + return { promise, resolve, reject }; +} + +async function waitFor(predicate: () => boolean, message: string): Promise { + const deadline = Date.now() + 5_000; + while (!predicate()) { + if (Date.now() >= deadline) { + throw new Error(message); + } + await new Promise((resolve) => setTimeout(resolve, 1)); + } +} + +async function withTimeout(promise: Promise, message: string): Promise { + return Promise.race([ + promise, + new Promise((_, reject) => setTimeout(() => reject(new Error(message)), 5_000)), + ]); +} + +class StaticResponseHandler extends CopilotRequestHandler { + constructor(private readonly response: Response) { + super(); + } + + protected override sendRequest( + _request: Request, + _ctx: CopilotRequestContext + ): Promise { + return Promise.resolve(this.response); + } +} + +class ProtocolPeer { + readonly starts: LlmInferenceHttpResponseStartRequest[] = []; + readonly chunks: LlmInferenceHttpResponseChunkRequest[] = []; + readonly dataAcks: Array> = []; + readonly terminal = deferred(); + outstandingDataRpcs = 0; + maxOutstandingDataRpcs = 0; + onStart: (() => void) | undefined; + + readonly rpc = { + llmInference: { + httpResponseStart: async ( + params: LlmInferenceHttpResponseStartRequest + ): Promise<{ accepted: boolean }> => { + this.starts.push(params); + this.onStart?.(); + return { accepted: true }; + }, + httpResponseChunk: ( + params: LlmInferenceHttpResponseChunkRequest + ): Promise<{ accepted: boolean }> => { + this.chunks.push(params); + if (params.end) { + this.terminal.resolve(params); + return Promise.resolve({ accepted: true }); + } + + this.outstandingDataRpcs++; + this.maxOutstandingDataRpcs = Math.max( + this.maxOutstandingDataRpcs, + this.outstandingDataRpcs + ); + const ack = deferred<{ accepted: boolean }>(); + this.dataAcks.push(ack); + return ack.promise.finally(() => { + this.outstandingDataRpcs--; + }); + }, + }, + }; +} + +function createAdapter( + handler: CopilotRequestHandler, + getRpc: () => ProtocolPeer["rpc"] | undefined +): LlmInferenceHandler { + return createCopilotRequestAdapter(handler, () => getRpc() as never); +} + +async function startGet(adapter: LlmInferenceHandler, requestId: string): Promise { + await adapter.httpRequestStart({ + requestId, + method: "GET", + url: "https://example.test/inference", + headers: {}, + }); + await adapter.httpRequestChunk({ requestId, data: "", end: true }); +} + +function frameBytes(frame: LlmInferenceHttpResponseChunkRequest): Uint8Array { + return frame.binary + ? new Uint8Array(Buffer.from(frame.data, "base64")) + : new TextEncoder().encode(frame.data); +} + +describe("CopilotRequestHandler HTTP response protocol", () => { + it("reads ahead, coalesces tiny source chunks, and keeps one data RPC outstanding", async () => { + const totalBytes = 40_000; + let sourceReads = 0; + const expected = Uint8Array.from({ length: totalBytes }, (_, index) => index % 251); + const body = new ReadableStream( + { + pull(controller) { + if (sourceReads === totalBytes) { + controller.close(); + return; + } + controller.enqueue(expected.subarray(sourceReads, sourceReads + 1)); + sourceReads++; + }, + }, + { highWaterMark: 0 } + ); + const peer = new ProtocolPeer(); + const adapter = createAdapter( + new StaticResponseHandler(new Response(body, { status: 200 })), + () => peer.rpc + ); + + await startGet(adapter, "coalesce"); + await waitFor(() => peer.dataAcks.length === 1, "first response chunk was not sent"); + + const firstLength = frameBytes(peer.chunks[0]).byteLength; + expect(firstLength).toBeGreaterThan(0); + expect(firstLength).toBeLessThan(READ_AHEAD_BYTES); + await waitFor( + () => sourceReads === firstLength + READ_AHEAD_BYTES, + "source did not read ahead while the first ACK was withheld" + ); + await new Promise((resolve) => setTimeout(resolve, 20)); + + expect(sourceReads - firstLength).toBe(READ_AHEAD_BYTES); + expect(peer.dataAcks).toHaveLength(1); + expect(peer.outstandingDataRpcs).toBe(1); + expect(peer.maxOutstandingDataRpcs).toBe(1); + + peer.dataAcks[0].resolve({ accepted: true }); + await waitFor(() => peer.dataAcks.length === 2, "coalesced response chunk was not sent"); + expect(frameBytes(peer.chunks[1])).toHaveLength(READ_AHEAD_BYTES); + expect(peer.outstandingDataRpcs).toBe(1); + await waitFor( + () => sourceReads === totalBytes, + "source did not finish reading behind the second ACK" + ); + + peer.dataAcks[1].resolve({ accepted: true }); + await waitFor(() => peer.dataAcks.length === 3, "final partial chunk was not sent"); + peer.dataAcks[2].resolve({ accepted: true }); + await withTimeout(peer.terminal.promise, "terminal response chunk was not sent"); + + const dataFrames = peer.chunks.filter((chunk) => !chunk.end); + expect(dataFrames).toHaveLength(3); + expect(peer.maxOutstandingDataRpcs).toBe(1); + expect(Buffer.concat(dataFrames.map((frame) => Buffer.from(frameBytes(frame))))).toEqual( + Buffer.from(expected) + ); + }); + + it("cancels the source promptly when the runtime cancels with an ACK withheld", async () => { + const sourceCancelled = deferred(); + const body = new ReadableStream( + { + pull(controller) { + controller.enqueue(Uint8Array.of(1)); + }, + cancel() { + sourceCancelled.resolve(); + }, + }, + { highWaterMark: 0 } + ); + const peer = new ProtocolPeer(); + const adapter = createAdapter( + new StaticResponseHandler(new Response(body, { status: 200 })), + () => peer.rpc + ); + + await startGet(adapter, "runtime-cancel"); + await waitFor(() => peer.dataAcks.length === 1, "response chunk was not sent"); + await adapter.httpRequestChunk({ + requestId: "runtime-cancel", + data: "", + cancel: true, + cancelReason: "turn aborted", + }); + + await withTimeout( + sourceCancelled.promise, + "runtime cancellation did not cancel the source" + ); + expect(peer.dataAcks).toHaveLength(1); + peer.dataAcks[0].reject(new Error("runtime rejected cancelled request")); + await withTimeout(peer.terminal.promise, "cancelled response did not settle"); + }); + + it("cancels the source when a response chunk RPC rejects", async () => { + const sourceCancelled = deferred(); + const body = new ReadableStream( + { + pull(controller) { + controller.enqueue(Uint8Array.of(1)); + }, + cancel() { + sourceCancelled.resolve(); + }, + }, + { highWaterMark: 0 } + ); + const peer = new ProtocolPeer(); + const adapter = createAdapter( + new StaticResponseHandler(new Response(body, { status: 200 })), + () => peer.rpc + ); + + await startGet(adapter, "rpc-rejection"); + await waitFor(() => peer.dataAcks.length === 1, "response chunk was not sent"); + peer.dataAcks[0].reject(new Error("response RPC rejected")); + + await withTimeout(sourceCancelled.promise, "RPC rejection did not cancel the source"); + const terminal = await withTimeout( + peer.terminal.promise, + "RPC rejection did not report a terminal error" + ); + expect(terminal.error?.message).toBe("response RPC rejected"); + }); + + it("cancels the source when the RPC connection disappears", async () => { + const sourceCancelled = deferred(); + const body = new ReadableStream( + { + pull(controller) { + controller.enqueue(Uint8Array.of(1)); + }, + cancel() { + sourceCancelled.resolve(); + }, + }, + { highWaterMark: 0 } + ); + const peer = new ProtocolPeer(); + let connected = true; + peer.onStart = () => { + connected = false; + }; + const adapter = createAdapter( + new StaticResponseHandler(new Response(body, { status: 200 })), + () => (connected ? peer.rpc : undefined) + ); + + await startGet(adapter, "connection-loss"); + await withTimeout(sourceCancelled.promise, "connection loss did not cancel the source"); + expect(peer.dataAcks).toHaveLength(0); + }); + + it("cancels an idle source when the RPC connection closes", async () => { + const readStarted = deferred(); + const sourceCancelled = deferred(); + const body = new ReadableStream( + { + async pull() { + readStarted.resolve(); + await new Promise(() => {}); + }, + cancel() { + sourceCancelled.resolve(); + }, + }, + { highWaterMark: 0 } + ); + const peer = new ProtocolPeer(); + const adapter = createCopilotRequestAdapter( + new StaticResponseHandler(new Response(body, { status: 200 })), + () => peer.rpc as never + ); + + await startGet(adapter, "idle-connection-loss"); + await withTimeout(readStarted.promise, "source read did not start"); + adapter.cancelPending(); + + await withTimeout(sourceCancelled.promise, "connection loss did not cancel the source"); + expect(peer.dataAcks).toHaveLength(0); + const terminal = await withTimeout( + peer.terminal.promise, + "connection loss did not settle the response" + ); + expect(terminal.error?.code).toBe("cancelled"); + }); + + it("flushes buffered bytes before reporting an upstream error", async () => { + const releaseSecondChunk = deferred(); + let pull = 0; + const body = new ReadableStream( + { + async pull(controller) { + if (pull === 0) { + controller.enqueue(Uint8Array.of(1, 2)); + } else if (pull === 1) { + await releaseSecondChunk.promise; + controller.enqueue(Uint8Array.of(3, 4)); + } else { + controller.error(new Error("upstream failed")); + } + pull++; + }, + }, + { highWaterMark: 0 } + ); + const peer = new ProtocolPeer(); + const adapter = createAdapter( + new StaticResponseHandler(new Response(body, { status: 200 })), + () => peer.rpc + ); + + await startGet(adapter, "upstream-error"); + await waitFor(() => peer.dataAcks.length === 1, "first response chunk was not sent"); + expect(frameBytes(peer.chunks[0])).toEqual(Uint8Array.of(1, 2)); + + releaseSecondChunk.resolve(); + await waitFor(() => pull === 3, "upstream error was not observed during read-ahead"); + peer.dataAcks[0].resolve({ accepted: true }); + await waitFor(() => peer.dataAcks.length === 2, "buffered response chunk was not sent"); + expect(frameBytes(peer.chunks[1])).toEqual(Uint8Array.of(3, 4)); + + peer.dataAcks[1].resolve({ accepted: true }); + const terminal = await withTimeout( + peer.terminal.promise, + "upstream error was not reported" + ); + expect(peer.chunks).toHaveLength(3); + expect(terminal.end).toBe(true); + expect(terminal.error?.message).toBe("upstream failed"); + expect(peer.maxOutstandingDataRpcs).toBe(1); + }); +}); diff --git a/nodejs/tsconfig.test.json b/nodejs/tsconfig.test.json index 8d24d6dfd8..33490b84a8 100644 --- a/nodejs/tsconfig.test.json +++ b/nodejs/tsconfig.test.json @@ -8,6 +8,7 @@ "include": [ "src/**/*", "test/dependency-policy.test.ts", + "test/copilot-request-handler.test.ts", "test/ffiRuntimeHost.test.ts", "test/session-event-types.test.ts", "test/message-source.test.ts" diff --git a/python/copilot/client.py b/python/copilot/client.py index e8110d656a..7e91bee7e8 100644 --- a/python/copilot/client.py +++ b/python/copilot/client.py @@ -66,7 +66,11 @@ CanvasProviderIdentity, ExtensionInfo, ) -from .copilot_request_handler import CopilotRequestHandler, create_copilot_request_adapter +from .copilot_request_handler import ( + CopilotRequestHandler, + _CopilotRequestAdapterHandler, + create_copilot_request_adapter, +) from .generated.rpc import ( ClientGlobalApiHandlers, ClientSessionApiHandlers, @@ -1779,6 +1783,7 @@ def __init__( _validate_session_fs_config(options.session_fs) self._session_fs_config = options.session_fs self._request_handler = options.request_handler + self._llm_inference_adapter: _CopilotRequestAdapterHandler | None = None def _resolve_runtime_entrypoint( self, @@ -2074,6 +2079,8 @@ async def stop(self) -> None: ... print(f"Cleanup error: {error.message}") """ errors: list[StopError] = [] + if self._llm_inference_adapter is not None: + self._llm_inference_adapter.cancel_pending() # Atomically take ownership of all sessions and clear the dict # so no other thread can access them @@ -2208,6 +2215,9 @@ async def force_stop(self) -> None: ... except asyncio.TimeoutError: ... await client.force_stop() """ + if self._llm_inference_adapter is not None: + self._llm_inference_adapter.cancel_pending() + # Clear sessions immediately without trying to destroy them with self._sessions_lock: sessions = list(self._sessions.values()) @@ -4862,9 +4872,9 @@ async def _set_session_fs_provider(self) -> None: def _register_client_global_handlers(self) -> None: if not self._client: return - llm_inference_adapter = None + self._llm_inference_adapter = None if self._request_handler is not None: - llm_inference_adapter = create_copilot_request_adapter( + self._llm_inference_adapter = create_copilot_request_adapter( self._request_handler, lambda: self._rpc.llm_inference if self._rpc is not None else None, ) @@ -4876,7 +4886,7 @@ def _register_client_global_handlers(self) -> None: ClientGlobalApiHandlers( hooks=_HooksAdapter(self._get_session), extension_launch_provider=self._options.extension_launch_provider, - llm_inference=llm_inference_adapter, + llm_inference=self._llm_inference_adapter, git_hub_telemetry=github_telemetry_adapter, git_hub_token=self._github_token_provider_adapter, ), @@ -4901,10 +4911,13 @@ def _handle_connection_close(self) -> None: with self._github_token_providers_lock: self._github_token_providers.clear() client = self._client + llm_inference_adapter = self._llm_inference_adapter loop = client._loop if client is not None else None if loop is not None and not loop.is_closed(): def cancel_pending_external_tools() -> None: + if llm_inference_adapter is not None: + llm_inference_adapter.cancel_pending() for session in sessions: session._cancel_pending_external_tools() diff --git a/python/copilot/copilot_request_handler.py b/python/copilot/copilot_request_handler.py index 893309ca76..051c6e9df0 100644 --- a/python/copilot/copilot_request_handler.py +++ b/python/copilot/copilot_request_handler.py @@ -23,7 +23,7 @@ import asyncio import base64 -from collections.abc import AsyncIterator, Callable +from collections.abc import AsyncIterator, Awaitable, Callable from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any @@ -61,6 +61,7 @@ ) _shared_http_client: httpx.AsyncClient | None = None +_HTTP_RESPONSE_READ_AHEAD_BYTES = 32 * 1024 def _get_shared_http_client() -> httpx.AsyncClient: @@ -559,6 +560,14 @@ async def _run(self, exchange: _CopilotRequestExchange) -> None: finally: self._pending.pop(exchange.request_id, None) + def cancel_pending(self) -> None: + exchanges = tuple(self._pending.values()) + self._pending.clear() + for exchange in exchanges: + exchange.cancelled = True + exchange.cancel_event.set() + exchange._queue.push(_BodyItem(cancel=True, cancel_reason="RPC connection closed")) + def _get_or_create(self, request_id: str) -> _CopilotRequestExchange: # The runtime dispatches httpRequestStart and httpRequestChunk frames # independently. get-or-create keeps the adapter correct regardless of @@ -727,15 +736,125 @@ async def _stream_response_to_exchange( ) if response.is_stream_consumed: # An in-memory response (built with ``content=``) has already buffered its - # body, so its raw stream cannot be iterated; forward the buffered bytes. - body = response.content - if body: - await exchange.write_response(body) + # body, so its raw stream cannot be iterated. + async def body_stream() -> AsyncIterator[bytes]: + if response.content: + yield response.content + + source = body_stream() else: - async for chunk in response.aiter_raw(): - if chunk: - await exchange.write_response(chunk) - await exchange.end_response() + source = response.aiter_raw() + + reader = _HttpResponseReader(source, response.aclose) + try: + while True: + chunk = await reader.next_chunk() + if chunk is None: + await exchange.end_response() + return + await exchange.write_response(chunk) + finally: + await reader.aclose() + + +class _HttpResponseReader: + """Bounded response-body read-ahead with immediate partial flushes.""" + + def __init__( + self, + source: AsyncIterator[bytes], + close_response: Callable[[], Awaitable[None]], + ) -> None: + self._source = source.__aiter__() + self._close_response = close_response + self._frames: list[bytes] = [] + self._queued = 0 + self._error: Exception | None = None + self._done = False + self._closed = False + self._condition = asyncio.Condition() + self._producer = asyncio.create_task(self._produce()) + + async def _produce(self) -> None: + try: + while True: + async with self._condition: + await self._condition.wait_for( + lambda: self._closed or self._queued < _HTTP_RESPONSE_READ_AHEAD_BYTES + ) + if self._closed: + return + + # Always-ready custom iterators must yield so an awaiting + # consumer can flush the first available frame immediately. + await asyncio.sleep(0) + frame = await anext(self._source) + if not frame: + continue + + async with self._condition: + if self._closed: + return + self._frames.append(frame) + self._queued += len(frame) + self._condition.notify_all() + except StopAsyncIteration: + # Normal source exhaustion. + pass + except asyncio.CancelledError: + raise + except Exception as exc: + async with self._condition: + if not self._closed: + self._error = exc + finally: + close_error: Exception | None = None + close = getattr(self._source, "aclose", None) + if close is not None: + try: + await close() + except Exception as exc: + close_error = exc + try: + await self._close_response() + except Exception as exc: + close_error = close_error or exc + async with self._condition: + if not self._closed and self._error is None: + self._error = close_error + self._done = True + self._condition.notify_all() + + async def next_chunk(self) -> bytes | None: + async with self._condition: + await self._condition.wait_for( + lambda: self._frames or self._error is not None or self._done + ) + if self._frames: + # A lone frame is forwarded as-is; joining is only needed once + # read-ahead has run past the chunk awaiting acknowledgement. + chunk = self._frames[0] if len(self._frames) == 1 else b"".join(self._frames) + self._frames.clear() + self._queued = 0 + self._condition.notify_all() + return chunk + if self._error is not None: + error = self._error + self._error = None + raise error + return None + + async def aclose(self) -> None: + async with self._condition: + self._closed = True + self._condition.notify_all() + if not self._producer.done(): + self._producer.cancel() + try: + await self._producer + except asyncio.CancelledError: + # Producer cancellation is expected during teardown. + pass def _headers_to_multi_map(headers: Any) -> LlmInferenceHeaders: diff --git a/python/test_client_start.py b/python/test_client_start.py index 626f0f6d96..a8ed47cc7b 100644 --- a/python/test_client_start.py +++ b/python/test_client_start.py @@ -1,7 +1,7 @@ """Startup concurrency regressions without a live CLI runtime.""" import asyncio -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, Mock import pytest @@ -83,3 +83,13 @@ async def block(): await client.start() client._start_cli_server.assert_awaited_once() client._verify_protocol_version.assert_awaited_once() + + +@pytest.mark.parametrize("method", ["stop", "force_stop"]) +async def test_shutdown_cancels_pending_inference_requests(client, method): + adapter = Mock() + client._llm_inference_adapter = adapter + + await getattr(client, method)() + + adapter.cancel_pending.assert_called_once_with() diff --git a/python/test_copilot_request_handler_response.py b/python/test_copilot_request_handler_response.py new file mode 100644 index 0000000000..cc055997c5 --- /dev/null +++ b/python/test_copilot_request_handler_response.py @@ -0,0 +1,342 @@ +import asyncio +import base64 + +import httpx +import pytest + +from copilot.copilot_request_handler import ( + CopilotRequestContext, + CopilotRequestHandler, + _CopilotRequestAdapterHandler, + _CopilotRequestExchange, + _stream_response_to_exchange, + create_copilot_request_adapter, +) +from copilot.generated.rpc import ( + LlmInferenceHTTPRequestChunkRequest, + LlmInferenceHTTPRequestStartRequest, +) + + +async def _wait_for(predicate, timeout: float = 2.0) -> None: + async with asyncio.timeout(timeout): + while not predicate(): + await asyncio.sleep(0) + + +class _ControlledStream(httpx.AsyncByteStream): + def __init__( + self, + chunks: list[bytes], + *, + error: Exception | None = None, + block_after_chunks: bool = False, + ) -> None: + self._chunks = chunks + self._error = error + self._block_after_chunks = block_after_chunks + self.read_count = 0 + self.exhausted = asyncio.Event() + self.read_blocked = asyncio.Event() + self.read_cancelled = asyncio.Event() + self.closed = asyncio.Event() + + async def __aiter__(self): + for chunk in self._chunks: + self.read_count += 1 + yield chunk + self.exhausted.set() + if self._error is not None: + raise self._error + if self._block_after_chunks: + self.read_blocked.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + self.read_cancelled.set() + raise + + async def aclose(self) -> None: + self.closed.set() + + +class _TaskAffineStream(httpx.AsyncByteStream): + def __init__(self) -> None: + self.owner: asyncio.Task[None] | None = None + self.closed = asyncio.Event() + + async def __aiter__(self): + owner = asyncio.current_task() + self.owner = owner + yield b"first" + if asyncio.current_task() is not owner: + raise RuntimeError("response iterator resumed from a different task") + yield b"second" + + async def aclose(self) -> None: + if self.closed.is_set(): + return + if asyncio.current_task() is not self.owner: + raise RuntimeError("response stream closed from a different task") + self.closed.set() + + +class _BufferedTaskAffineStream(httpx.AsyncByteStream): + def __init__(self) -> None: + self.owner: asyncio.Task[None] | None = None + self.read_count = 0 + self.closed = asyncio.Event() + + async def __aiter__(self): + self.owner = asyncio.current_task() + while True: + self.read_count += 1 + yield b"x" * 1024 + + async def aclose(self) -> None: + if self.closed.is_set(): + return + if asyncio.current_task() is not self.owner: + raise RuntimeError("response stream closed from a different task") + self.closed.set() + + +class _WithheldAckRpc: + def __init__(self) -> None: + self.starts = [] + self.chunks = [] + self.acks: list[asyncio.Future[None]] = [] + self.outstanding_data = 0 + self.max_outstanding_data = 0 + + @property + def data_chunks(self): + return [chunk for chunk in self.chunks if not chunk.end] + + async def http_response_start(self, params): + self.starts.append(params) + + async def http_response_chunk(self, params): + self.chunks.append(params) + if params.end: + return + ack = asyncio.get_running_loop().create_future() + self.acks.append(ack) + self.outstanding_data += 1 + self.max_outstanding_data = max(self.max_outstanding_data, self.outstanding_data) + try: + _ = await ack + finally: + self.outstanding_data -= 1 + + def acknowledge(self, index: int) -> None: + self.acks[index].set_result(None) + + def reject(self, index: int, error: Exception) -> None: + self.acks[index].set_exception(error) + + +def _response(stream: httpx.AsyncByteStream) -> httpx.Response: + return httpx.Response( + 200, + headers={"content-type": "application/octet-stream"}, + stream=stream, + request=httpx.Request("GET", "https://example.test/response"), + ) + + +async def _pump(response: httpx.Response, exchange: _CopilotRequestExchange) -> None: + try: + await _stream_response_to_exchange(response, exchange) + finally: + await response.aclose() + + +def _decoded_data(rpc: _WithheldAckRpc) -> list[bytes]: + return [base64.b64decode(chunk.data) for chunk in rpc.data_chunks] + + +@pytest.mark.asyncio +async def test_reads_ahead_and_coalesces_with_one_data_rpc_outstanding() -> None: + stream = _ControlledStream([b"x" * 1024] * 70) + response = _response(stream) + rpc = _WithheldAckRpc() + exchange = _CopilotRequestExchange("request", lambda: rpc) + + pump = asyncio.create_task(_pump(response, exchange)) + await _wait_for(lambda: len(rpc.data_chunks) == 1 and stream.read_count == 33) + + assert _decoded_data(rpc) == [b"x" * 1024] + assert rpc.outstanding_data == 1 + assert rpc.max_outstanding_data == 1 + reads_at_capacity = stream.read_count + await asyncio.sleep(0.01) + assert stream.read_count == reads_at_capacity + assert len(rpc.data_chunks) == 1 + + rpc.acknowledge(0) + await _wait_for(lambda: len(rpc.data_chunks) == 2 and stream.read_count == 65) + assert len(_decoded_data(rpc)[1]) == 32 * 1024 + assert rpc.max_outstanding_data == 1 + + rpc.acknowledge(1) + await _wait_for(lambda: len(rpc.data_chunks) == 3 and stream.exhausted.is_set()) + assert len(_decoded_data(rpc)[2]) == 32 * 1024 + + rpc.acknowledge(2) + await _wait_for(lambda: len(rpc.data_chunks) == 4) + assert len(_decoded_data(rpc)[3]) == 5 * 1024 + rpc.acknowledge(3) + await asyncio.wait_for(pump, timeout=2) + + assert b"".join(_decoded_data(rpc)) == b"x" * (70 * 1024) + assert rpc.max_outstanding_data == 1 + assert rpc.chunks[-1].end is True + assert rpc.chunks[-1].error is None + assert stream.closed.is_set() + + +@pytest.mark.asyncio +async def test_response_iterator_is_advanced_by_one_persistent_task() -> None: + stream = _TaskAffineStream() + response = _response(stream) + rpc = _WithheldAckRpc() + exchange = _CopilotRequestExchange("request", lambda: rpc) + + pump = asyncio.create_task(_pump(response, exchange)) + await _wait_for(lambda: len(rpc.data_chunks) == 1) + rpc.acknowledge(0) + await _wait_for(lambda: len(rpc.data_chunks) == 2) + rpc.acknowledge(1) + await asyncio.wait_for(pump, timeout=2) + + assert b"".join(_decoded_data(rpc)) == b"firstsecond" + assert stream.closed.is_set() + + +@pytest.mark.asyncio +async def test_response_stream_is_closed_by_producer_when_read_ahead_is_full() -> None: + stream = _BufferedTaskAffineStream() + response = _response(stream) + rpc = _WithheldAckRpc() + exchange = _CopilotRequestExchange("request", lambda: rpc) + + pump = asyncio.create_task(_pump(response, exchange)) + await _wait_for(lambda: len(rpc.data_chunks) == 1 and stream.read_count == 33) + rpc.reject(0, ConnectionError("runtime connection lost")) + + with pytest.raises(ConnectionError, match="runtime connection lost"): + await asyncio.wait_for(pump, timeout=2) + assert stream.closed.is_set() + + +class _ResponseHandler(CopilotRequestHandler): + def __init__(self, response: httpx.Response) -> None: + self._response = response + + async def send_request( + self, request: httpx.Request, ctx: CopilotRequestContext + ) -> httpx.Response: + return self._response + + +async def _start_adapter_request( + stream: _ControlledStream, +) -> tuple[_WithheldAckRpc, _CopilotRequestAdapterHandler, asyncio.Task[None]]: + rpc = _WithheldAckRpc() + adapter = create_copilot_request_adapter(_ResponseHandler(_response(stream)), lambda: rpc) + request_id = "request" + await adapter.http_request_start( + LlmInferenceHTTPRequestStartRequest( + headers={}, + method="GET", + request_id=request_id, + url="https://example.test/response", + ) + ) + await adapter.http_request_chunk( + LlmInferenceHTTPRequestChunkRequest(data="", request_id=request_id, end=True) + ) + exchange = adapter._pending[request_id] + assert exchange.task is not None + return rpc, adapter, exchange.task + + +@pytest.mark.asyncio +async def test_runtime_cancellation_stops_pending_read_and_closes_source() -> None: + stream = _ControlledStream([b"first"], block_after_chunks=True) + rpc, adapter, task = await _start_adapter_request(stream) + await asyncio.wait_for(stream.read_blocked.wait(), timeout=2) + await _wait_for(lambda: len(rpc.data_chunks) == 1) + + await adapter.http_request_chunk( + LlmInferenceHTTPRequestChunkRequest( + data="", + request_id="request", + cancel=True, + cancel_reason="consumer stopped", + ) + ) + await asyncio.wait_for(task, timeout=2) + + assert stream.read_cancelled.is_set() + assert stream.closed.is_set() + assert rpc.outstanding_data == 0 + assert rpc.max_outstanding_data == 1 + assert rpc.chunks[-1].end is True + assert rpc.chunks[-1].error.code == "cancelled" + + +@pytest.mark.asyncio +async def test_upstream_error_follows_all_buffered_bytes() -> None: + stream = _ControlledStream( + [b"first", b"partial"], + error=RuntimeError("synthetic upstream failure"), + ) + rpc, _, task = await _start_adapter_request(stream) + await asyncio.wait_for(stream.exhausted.wait(), timeout=2) + await _wait_for(lambda: len(rpc.data_chunks) == 1) + + rpc.acknowledge(0) + await _wait_for(lambda: len(rpc.data_chunks) == 2) + rpc.acknowledge(1) + await asyncio.wait_for(task, timeout=2) + + assert _decoded_data(rpc) == [b"first", b"partial"] + assert rpc.chunks[-1].end is True + assert rpc.chunks[-1].error.message == "synthetic upstream failure" + assert stream.closed.is_set() + + +@pytest.mark.asyncio +async def test_rpc_rejection_stops_pending_read_and_closes_source() -> None: + stream = _ControlledStream([b"first"], block_after_chunks=True) + rpc, _, task = await _start_adapter_request(stream) + await asyncio.wait_for(stream.read_blocked.wait(), timeout=2) + await _wait_for(lambda: len(rpc.data_chunks) == 1) + + rpc.reject(0, ConnectionError("runtime connection lost")) + await asyncio.wait_for(task, timeout=2) + + assert stream.read_cancelled.is_set() + assert stream.closed.is_set() + assert rpc.outstanding_data == 0 + assert rpc.max_outstanding_data == 1 + assert rpc.chunks[-1].end is True + assert rpc.chunks[-1].error.message == "runtime connection lost" + + +@pytest.mark.asyncio +async def test_connection_loss_stops_idle_read_and_closes_source() -> None: + stream = _ControlledStream([], block_after_chunks=True) + rpc, adapter, task = await _start_adapter_request(stream) + await asyncio.wait_for(stream.read_blocked.wait(), timeout=2) + + adapter.cancel_pending() + await asyncio.wait_for(task, timeout=2) + + assert stream.read_cancelled.is_set() + assert stream.closed.is_set() + assert rpc.data_chunks == [] + assert rpc.chunks[-1].end is True + assert rpc.chunks[-1].error.code == "cancelled"