MultiClientCommandsElicitationE2ETests.cs

/*---------------------------------------------------------------------------------------------
 *  Copyright (c) Microsoft Corporation. All rights reserved.
 *--------------------------------------------------------------------------------------------*/

using GitHub.Copilot.Test.Harness;
using Xunit;
using Xunit.Abstractions;

namespace GitHub.Copilot.Test.E2E;

/// <summary>
/// Custom fixture for multi-client commands/elicitation tests.
/// Uses TCP mode so a second (and third) client can connect to the same CLI process.
/// </summary>
public class MultiClientCommandsElicitationFixture : IAsyncLifetime
{
    public E2ETestContext Ctx { get; private set; } = null!;
    public CopilotClient Client1 { get; private set; } = null!;

    public const string SharedToken = "multi-client-cmd-shared-token";

    public async Task InitializeAsync()
    {
        Ctx = await E2ETestContext.CreateAsync();
        Client1 = Ctx.CreateClient(options: new CopilotClientOptions
        {
            Connection = RuntimeConnection.ForTcp(connectionToken: SharedToken),
        }, persistent: true);
    }

    public async Task DisposeAsync()
    {
        await Ctx.DisposeAsync();
    }
}

public class MultiClientCommandsElicitationE2ETests
    : IClassFixture<MultiClientCommandsElicitationFixture>, IAsyncLifetime
{
    private readonly MultiClientCommandsElicitationFixture _fixture;
    private readonly string _testName;
    private CopilotClient? _client2;
    private CopilotClient? _client3;

    private E2ETestContext Ctx => _fixture.Ctx;
    private CopilotClient Client1 => _fixture.Client1;

    public MultiClientCommandsElicitationE2ETests(
        MultiClientCommandsElicitationFixture fixture,
        ITestOutputHelper output)
    {
        _fixture = fixture;
        _testName = E2ETestBase.GetTestName(output);
    }

    public async Task InitializeAsync()
    {
        await Ctx.CleanupAfterTestAsync();
        await Ctx.ConfigureForTestAsync("multi_client", _testName);

        // Trigger connection so we can read the port
        var initSession = await Ctx.CreateSessionAsync(Client1, new SessionConfig
        {
            OnPermissionRequest = PermissionHandler.ApproveAll,
        });
        await initSession.DisposeAsync();

        var port = Client1.RuntimePort
            ?? throw new InvalidOperationException("Client1 is not using TCP mode; RuntimePort is null");

        _client2 = Ctx.CreateClient(options: new CopilotClientOptions
        {
            Connection = RuntimeConnection.ForUri($"localhost:{port}", connectionToken: MultiClientCommandsElicitationFixture.SharedToken),
        });
    }

    public async Task DisposeAsync()
    {
        try
        {
            if (_client3 is not null)
            {
                await _client3.ForceStopAsync();
            }

            if (_client2 is not null)
            {
                await _client2.ForceStopAsync();
            }
        }
        finally
        {
            _client3 = null;
            _client2 = null;
            await Ctx.CleanupAfterTestAsync();
        }
    }

    private CopilotClient Client2 => _client2
        ?? throw new InvalidOperationException("Client2 not initialized");

    [Fact]
    public async Task Client_Receives_Commands_Changed_When_Another_Client_Joins_With_Commands()
    {
        var session1 = await Ctx.CreateSessionAsync(Client1, new SessionConfig
        {
            OnPermissionRequest = PermissionHandler.ApproveAll,
        });

        // Wait for the commands.changed event deterministically
        var commandsChangedTcs = new TaskCompletionSource<CommandsChangedEvent>(
            TaskCreationOptions.RunContinuationsAsynchronously);

        using var sub = session1.On<SessionEvent>(evt =>
        {
            if (evt is CommandsChangedEvent changed)
            {
                commandsChangedTcs.TrySetResult(changed);
            }
        });

        // Client2 joins with commands
        var session2 = await Ctx.ResumeSessionAsync(Client2, session1.SessionId, new ResumeSessionConfig
        {
            OnPermissionRequest = PermissionHandler.ApproveAll,
            Commands =
            [
                new CommandDefinition
                {
                    Name = "deploy",
                    Description = "Deploy the app",
                    Handler = _ => Task.CompletedTask,
                },
            ],
            SuppressResumeEvent = true,
        });

        var commandsChanged = await commandsChangedTcs.Task.WaitAsync(TimeSpan.FromSeconds(15));

        Assert.NotNull(commandsChanged.Data.Commands);
        Assert.Contains(commandsChanged.Data.Commands, c =>
            c.Name == "deploy" && c.Description == "Deploy the app");

        await session2.DisposeAsync();
    }

    [Fact]
    public async Task Capabilities_Changed_Fires_When_Second_Client_Joins_With_Elicitation_Handler()
    {
        // Client1 creates session without elicitation
        var session1 = await Ctx.CreateSessionAsync(Client1, new SessionConfig
        {
            OnPermissionRequest = PermissionHandler.ApproveAll,
        });
        Assert.True(session1.Capabilities.Ui?.Elicitation != true,
            "Session without elicitation handler should not have elicitation capability");

        // Listen for capabilities.changed event
        var capChangedTcs = new TaskCompletionSource<CapabilitiesChangedEvent>(
            TaskCreationOptions.RunContinuationsAsynchronously);

        using var sub = session1.On<SessionEvent>(evt =>
        {
            if (evt is CapabilitiesChangedEvent capEvt)
            {
                capChangedTcs.TrySetResult(capEvt);
            }
        });

        // Client2 joins WITH elicitation handler — triggers capabilities.changed
        var session2 = await Ctx.ResumeSessionAsync(Client2, session1.SessionId, new ResumeSessionConfig
        {
            OnPermissionRequest = PermissionHandler.ApproveAll,
            OnElicitationRequest = _ => Task.FromResult(new ElicitationResult
            {
                Action = Rpc.UIElicitationResponseAction.Accept,
                Content = new Dictionary<string, object>(),
            }),
            SuppressResumeEvent = true,
        });

        var capEvent = await capChangedTcs.Task.WaitAsync(TimeSpan.FromSeconds(15));

        Assert.NotNull(capEvent.Data.Ui);
        Assert.True(capEvent.Data.Ui!.Elicitation);

        // Client1's capabilities should have been auto-updated
        Assert.True(session1.Capabilities.Ui?.Elicitation == true);

        await session2.DisposeAsync();
    }

    [Fact]
    public async Task Capabilities_Changed_Fires_When_Elicitation_Provider_Disconnects()
    {
        // Client1 creates session without elicitation
        var session1 = await Ctx.CreateSessionAsync(Client1, new SessionConfig
        {
            OnPermissionRequest = PermissionHandler.ApproveAll,
        });
        Assert.True(session1.Capabilities.Ui?.Elicitation != true,
            "Session without elicitation handler should not have elicitation capability");

        // Wait for elicitation to become available
        var capEnabledTcs = new TaskCompletionSource<bool>(
            TaskCreationOptions.RunContinuationsAsynchronously);

        using var subEnabled = session1.On<SessionEvent>(evt =>
        {
            if (evt is CapabilitiesChangedEvent { Data.Ui.Elicitation: true })
            {
                capEnabledTcs.TrySetResult(true);
            }
        });

        // Use a dedicated client (client3) so we can stop it without affecting client2
        var port = Client1.RuntimePort
            ?? throw new InvalidOperationException("Client1 RuntimePort is null");
        _client3 = Ctx.CreateClient(options: new CopilotClientOptions
        {
            Connection = RuntimeConnection.ForUri($"localhost:{port}", connectionToken: MultiClientCommandsElicitationFixture.SharedToken),
        });

        // Client3 joins WITH elicitation handler
        await Ctx.ResumeSessionAsync(_client3, session1.SessionId, new ResumeSessionConfig
        {
            OnPermissionRequest = PermissionHandler.ApproveAll,
            OnElicitationRequest = _ => Task.FromResult(new ElicitationResult
            {
                Action = Rpc.UIElicitationResponseAction.Accept,
                Content = new Dictionary<string, object>(),
            }),
            SuppressResumeEvent = true,
        });

        await capEnabledTcs.Task.WaitAsync(TimeSpan.FromSeconds(15));
        Assert.True(session1.Capabilities.Ui?.Elicitation == true);

        // Now listen for the capability being removed
        var capDisabledTcs = new TaskCompletionSource<bool>(
            TaskCreationOptions.RunContinuationsAsynchronously);

        using var subDisabled = session1.On<SessionEvent>(evt =>
        {
            if (evt is CapabilitiesChangedEvent { Data.Ui.Elicitation: false })
            {
                capDisabledTcs.TrySetResult(true);
            }
        });

        // Force-stop client3 — destroys the socket, triggering server-side cleanup
        await _client3.ForceStopAsync();
        _client3 = null;

        // Network teardown + server-side cleanup + capabilities recompute can take time on
        // slow CI runners. 30s is a defensive upper bound.
        await capDisabledTcs.Task.WaitAsync(TimeSpan.FromSeconds(30));
        Assert.True(session1.Capabilities.Ui?.Elicitation != true,
            "After elicitation provider disconnects, capability should be removed");
    }
}