/*--------------------------------------------------------------------------------------------- * 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