Skip to content

Commit 39fe821

Browse files
stephentoubCopilot
stephentoub
and
Copilot
authored
Coalesce intercepted HTTP response chunks (#2734)
* Coalesce intercepted HTTP response chunks Add bounded 32 KiB read-ahead across Node.js, Python, Go, .NET, and Java while preserving byte ordering, cancellation, and a single outstanding data RPC. Add protocol-level coverage for backpressure, cancellation, connection loss, and upstream errors. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> * Fix scenario event test concurrency Synchronize the resumed session event list while callbacks and assertions access it to avoid collection-modified failures in the Windows E2E matrix. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> * Fix response forwarding lifecycle edge cases Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> * Explain expected response reader exceptions Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> * Resolve .NET merge-build duplication Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
1 parent c310b24 commit 39fe821

25 files changed

Lines changed: 3657 additions & 106 deletions

‎dotnet/src/Client.cs‎

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -557,6 +557,7 @@ await InvokeRpcAsync<JsonElement>(
557557
public async Task StopAsync()
558558
{
559559
List<Exception> errors = [];
560+
CancelPendingExternalTools();
560561

561562
foreach (var session in _sessions.Values.ToArray())
562563
{
@@ -602,10 +603,7 @@ public async Task StopAsync()
602603
/// </example>
603604
public async Task ForceStopAsync()
604605
{
605-
foreach (var session in _sessions.Values)
606-
{
607-
session.CancelPendingExternalTools();
608-
}
606+
CancelPendingExternalTools();
609607
_sessions.Clear();
610608
ClearGitHubTokenProviders();
611609

@@ -2788,6 +2786,15 @@ private async Task CancelExternalToolsWhenConnectionClosesAsync(JsonRpc rpc)
27882786
{
27892787
return;
27902788
}
2789+
CancelPendingExternalTools();
2790+
}
2791+
2792+
private void CancelPendingExternalTools()
2793+
{
2794+
if (_clientGlobalApis?.LlmInference is LlmInferenceAdapter llmInferenceAdapter)
2795+
{
2796+
llmInferenceAdapter.CancelPending();
2797+
}
27912798
foreach (var session in _sessions.Values)
27922799
{
27932800
session.CancelPendingExternalTools();

‎dotnet/src/CopilotRequestHandler.cs‎

Lines changed: 229 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
using System.Diagnostics.CodeAnalysis;
99
using System.Net.WebSockets;
1010
using System.Runtime.CompilerServices;
11+
using System.Runtime.ExceptionServices;
1112
using System.Text;
1213
using System.Threading.Channels;
1314

@@ -462,6 +463,7 @@ private static Uri ToWebSocketUri(string url)
462463
[Experimental(Diagnostics.Experimental)]
463464
public class CopilotRequestHandler
464465
{
466+
private const int HttpResponseReadAheadSize = 32 * 1024;
465467
private static readonly HttpClient s_sharedHttpClient = new();
466468

467469
private readonly HttpClient _httpClient;
@@ -554,23 +556,45 @@ private static async Task<HttpRequestMessage> BuildHttpRequestAsync(LlmInference
554556

555557
private static async Task StreamResponseAsync(HttpResponseMessage response, LlmInferenceExchange exchange)
556558
{
557-
await exchange.StartResponseAsync(
558-
(int)response.StatusCode,
559-
response.ReasonPhrase,
560-
HeadersToMultiMap(response)).ConfigureAwait(false);
561-
562559
var ct = exchange.Context.CancellationToken;
563-
using var stream = await response.Content.ReadAsStreamAsync(ct).ConfigureAwait(false);
564-
var buffer = new byte[16 * 1024];
560+
await AwaitRpcAsync(
561+
exchange.StartResponseAsync(
562+
(int)response.StatusCode,
563+
response.ReasonPhrase,
564+
HeadersToMultiMap(response)),
565+
ct).ConfigureAwait(false);
566+
567+
var stream = await response.Content.ReadAsStreamAsync(ct).ConfigureAwait(false);
568+
await using var reader = new BoundedHttpResponseReader(stream, HttpResponseReadAheadSize, ct);
569+
var buffer = new byte[HttpResponseReadAheadSize];
565570
int read;
566-
while ((read = await stream.ReadAsync(buffer.AsMemory(), ct).ConfigureAwait(false)) > 0)
571+
while ((read = await reader.ReadChunkAsync(buffer, ct).ConfigureAwait(false)) > 0)
567572
{
568-
await exchange.WriteResponseAsync(new ReadOnlyMemory<byte>(buffer, 0, read)).ConfigureAwait(false);
573+
await AwaitRpcAsync(
574+
exchange.WriteResponseAsync(new ReadOnlyMemory<byte>(buffer, 0, read)),
575+
ct).ConfigureAwait(false);
569576
}
570577

571578
await exchange.EndResponseAsync().ConfigureAwait(false);
572579
}
573580

581+
private static async Task AwaitRpcAsync(Task rpcTask, CancellationToken cancellationToken)
582+
{
583+
try
584+
{
585+
await rpcTask.WaitAsync(Timeout.InfiniteTimeSpan, cancellationToken).ConfigureAwait(false);
586+
}
587+
catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
588+
{
589+
_ = rpcTask.ContinueWith(
590+
static completed => _ = completed.Exception,
591+
CancellationToken.None,
592+
TaskContinuationOptions.ExecuteSynchronously | TaskContinuationOptions.OnlyOnFaulted,
593+
TaskScheduler.Default);
594+
throw;
595+
}
596+
}
597+
574598
private async Task HandleWebSocketAsync(LlmInferenceExchange exchange)
575599
{
576600
var ctx = exchange.Context;
@@ -657,6 +681,191 @@ private static Dictionary<string, IReadOnlyList<string>> HeadersToMultiMap(HttpR
657681

658682
return result;
659683
}
684+
685+
/// <summary>
686+
/// Reads an HTTP response into a fixed-size ring while the current response
687+
/// chunk waits for its runtime acknowledgement. The producer reserves ring
688+
/// space before every read, so committed bytes plus an in-flight read never
689+
/// exceed the configured read-ahead bound.
690+
/// </summary>
691+
private sealed class BoundedHttpResponseReader : IAsyncDisposable
692+
{
693+
private readonly Stream _stream;
694+
private readonly byte[] _buffer;
695+
private readonly CancellationTokenSource _disposeCts;
696+
private readonly object _gate = new();
697+
private TaskCompletionSource<bool> _changed = CreateSignal();
698+
private readonly Task _pump;
699+
700+
private long _head;
701+
private long _committedTail;
702+
private long _reservedTail;
703+
private bool _completed;
704+
private ExceptionDispatchInfo? _error;
705+
706+
internal BoundedHttpResponseReader(Stream stream, int capacity, CancellationToken cancellationToken)
707+
{
708+
_stream = stream;
709+
_buffer = new byte[capacity];
710+
_disposeCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
711+
_pump = PumpAsync(_disposeCts.Token);
712+
}
713+
714+
internal async Task<int> ReadChunkAsync(byte[] destination, CancellationToken cancellationToken)
715+
{
716+
while (true)
717+
{
718+
Task waitTask;
719+
ExceptionDispatchInfo? error;
720+
lock (_gate)
721+
{
722+
var count = checked((int)(_committedTail - _head));
723+
if (count > 0)
724+
{
725+
var headIndex = (int)(_head % _buffer.Length);
726+
var firstCount = Math.Min(count, _buffer.Length - headIndex);
727+
Buffer.BlockCopy(_buffer, headIndex, destination, 0, firstCount);
728+
if (firstCount < count)
729+
{
730+
Buffer.BlockCopy(_buffer, 0, destination, firstCount, count - firstCount);
731+
}
732+
733+
_head += count;
734+
PulseLocked();
735+
return count;
736+
}
737+
738+
error = _error;
739+
if (error is null)
740+
{
741+
if (_completed)
742+
{
743+
cancellationToken.ThrowIfCancellationRequested();
744+
return 0;
745+
}
746+
747+
waitTask = _changed.Task;
748+
}
749+
else
750+
{
751+
waitTask = Task.CompletedTask;
752+
}
753+
}
754+
755+
if (error is not null)
756+
{
757+
error.Throw();
758+
}
759+
760+
await waitTask.WaitAsync(Timeout.InfiniteTimeSpan, cancellationToken).ConfigureAwait(false);
761+
}
762+
}
763+
764+
public async ValueTask DisposeAsync()
765+
{
766+
_disposeCts.Cancel();
767+
_stream.Dispose();
768+
try
769+
{
770+
await _pump.ConfigureAwait(false);
771+
}
772+
catch (OperationCanceledException) when (_disposeCts.IsCancellationRequested)
773+
{
774+
// Cancellation is the expected result of disposing an active pump.
775+
}
776+
finally
777+
{
778+
_disposeCts.Dispose();
779+
}
780+
}
781+
782+
private async Task PumpAsync(CancellationToken cancellationToken)
783+
{
784+
while (true)
785+
{
786+
Task? waitTask = null;
787+
long reservationStart = 0;
788+
int reservationLength = 0;
789+
int reservationIndex = 0;
790+
791+
lock (_gate)
792+
{
793+
var used = checked((int)(_reservedTail - _head));
794+
if (used == _buffer.Length)
795+
{
796+
waitTask = _changed.Task;
797+
}
798+
else
799+
{
800+
reservationStart = _reservedTail;
801+
reservationIndex = (int)(reservationStart % _buffer.Length);
802+
reservationLength = Math.Min(_buffer.Length - used, _buffer.Length - reservationIndex);
803+
_reservedTail += reservationLength;
804+
}
805+
}
806+
807+
if (waitTask is not null)
808+
{
809+
await waitTask.WaitAsync(Timeout.InfiniteTimeSpan, cancellationToken).ConfigureAwait(false);
810+
continue;
811+
}
812+
813+
int read;
814+
try
815+
{
816+
read = await _stream.ReadAsync(
817+
_buffer.AsMemory(reservationIndex, reservationLength),
818+
cancellationToken).ConfigureAwait(false);
819+
}
820+
catch (Exception ex)
821+
{
822+
lock (_gate)
823+
{
824+
_reservedTail = reservationStart;
825+
if (cancellationToken.IsCancellationRequested)
826+
{
827+
_completed = true;
828+
}
829+
else
830+
{
831+
_error = ExceptionDispatchInfo.Capture(ex);
832+
}
833+
834+
PulseLocked();
835+
}
836+
837+
return;
838+
}
839+
840+
lock (_gate)
841+
{
842+
_committedTail = reservationStart + read;
843+
_reservedTail = _committedTail;
844+
if (read == 0)
845+
{
846+
_completed = true;
847+
}
848+
849+
PulseLocked();
850+
}
851+
852+
if (read == 0)
853+
{
854+
return;
855+
}
856+
}
857+
}
858+
859+
private void PulseLocked()
860+
{
861+
var changed = _changed;
862+
_changed = CreateSignal();
863+
changed.TrySetResult(true);
864+
}
865+
866+
private static TaskCompletionSource<bool> CreateSignal() =>
867+
new(TaskCreationOptions.RunContinuationsAsynchronously);
868+
}
660869
}
661870

662871
/// <summary>
@@ -917,6 +1126,17 @@ public Task<LlmInferenceHttpRequestChunkResult> HttpRequestChunkAsync(LlmInferen
9171126
return Task.FromResult(new LlmInferenceHttpRequestChunkResult());
9181127
}
9191128

1129+
internal void CancelPending()
1130+
{
1131+
foreach (var (requestId, exchange) in _pending)
1132+
{
1133+
if (_pending.TryRemove(requestId, out _))
1134+
{
1135+
exchange.PushCancel("RPC connection closed");
1136+
}
1137+
}
1138+
}
1139+
9201140
private async Task RunAsync(LlmInferenceExchange exchange)
9211141
{
9221142
try

‎dotnet/test/E2E/ScenarioTestingEventSubscriptionsE2ETests.cs‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -141,7 +141,13 @@ await TestHelper.WaitForConditionAsync(
141141
});
142142

143143
var newEvents = new List<SessionEvent>();
144-
using var newSubscription = secondSession.On<SessionEvent>(newEvents.Add);
144+
using var newSubscription = secondSession.On<SessionEvent>(evt =>
145+
{
146+
lock (newEvents)
147+
{
148+
newEvents.Add(evt);
149+
}
150+
});
145151
var newInfo = new TaskCompletionSource<SessionInfoEvent>(TaskCreationOptions.RunContinuationsAsynchronously);
146152
using var infoSubscription = secondSession.On<SessionInfoEvent>(evt =>
147153
{
@@ -154,7 +160,10 @@ await TestHelper.WaitForConditionAsync(
154160
await newInfo.Task.WaitAsync(EventTimeout);
155161

156162
Assert.Equal(countAfterClose, Volatile.Read(ref oldEventCount));
157-
Assert.Contains(newEvents, evt => evt is SessionInfoEvent info && info.Data.Message == "SCENARIO_EVENT_SOURCE_TWO");
163+
lock (newEvents)
164+
{
165+
Assert.Contains(newEvents, evt => evt is SessionInfoEvent info && info.Data.Message == "SCENARIO_EVENT_SOURCE_TWO");
166+
}
158167
}
159168

160169
private static bool IsEventChannelClosed(CopilotSession session)

0 commit comments

Comments
 (0)