CopilotClientTest.java

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

package com.github.copilot;

import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.io.TempDir;

import com.github.copilot.generated.ExternalToolRequestedEvent;
import com.github.copilot.rpc.CloudSessionOptions;
import com.github.copilot.rpc.CopilotClientOptions;
import com.github.copilot.rpc.CopilotClientMode;
import com.github.copilot.rpc.CreateSessionRequest;
import com.github.copilot.rpc.CreateSessionResponse;
import com.github.copilot.rpc.DeleteSessionResponse;
import com.github.copilot.rpc.GitHubTokenProviderResult;
import com.github.copilot.rpc.PermissionHandler;
import com.github.copilot.rpc.PingResponse;
import com.github.copilot.rpc.ResumeSessionConfig;
import com.github.copilot.rpc.ResumeSessionResponse;
import com.github.copilot.rpc.SessionConfig;
import com.github.copilot.rpc.SessionLifecycleEvent;
import com.github.copilot.rpc.SessionLifecycleEventTypes;
import com.github.copilot.rpc.ToolDefinition;

import java.io.OutputStream;
import java.lang.reflect.Field;
import java.io.BufferedReader;
import java.io.FilterOutputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.CancellationException;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;

import static org.junit.jupiter.api.Assertions.*;
import static org.mockito.ArgumentMatchers.*;
import static org.mockito.Mockito.*;

/**
 * Tests for CopilotClient.
 *
 * Note: These tests require the Copilot CLI to be installed. Set the
 * COPILOT_CLI_PATH environment variable to the path to the CLI, or run 'npm
 * install' in the nodejs directory.
 */
public class CopilotClientTest {

    private static String cliPath;

    @BeforeAll
    static void setup() {
        cliPath = TestUtil.findCliPath();
    }

    @Test
    void testStopDoesNotThrowWhenRuntimeShutdownFails() throws Exception {
        var client = new CopilotClient(new CopilotClientOptions().setAutoStart(false));
        var rpc = mock(JsonRpcClient.class);
        when(rpc.invoke(eq("runtime.shutdown"), any(), eq(Void.class)))
                .thenReturn(CompletableFuture.failedFuture(new RuntimeException("shutdown failed")));
        var process = mock(Process.class);
        when(process.isAlive()).thenReturn(true);
        when(process.getOutputStream()).thenReturn(OutputStream.nullOutputStream());
        when(process.destroyForcibly()).thenReturn(process);
        when(process.waitFor(anyLong(), any(TimeUnit.class))).thenReturn(true);

        setConnectionFuture(client, rpc, process);

        assertDoesNotThrow(() -> client.stop().get());

        verify(rpc).invoke(eq("runtime.shutdown"), eq(Map.of()), eq(Void.class));
        verify(rpc).close();
        verify(process).destroyForcibly();
    }

    @Test
    void testForceStopAndExternalStopDoNotRequestRuntimeShutdown() throws Exception {
        var forceClient = new CopilotClient(new CopilotClientOptions().setAutoStart(false));
        var forceRpc = mock(JsonRpcClient.class);
        var process = mock(Process.class);
        when(process.isAlive()).thenReturn(true);
        when(process.destroyForcibly()).thenReturn(process);
        when(process.waitFor(anyLong(), any(TimeUnit.class))).thenReturn(true);
        setConnectionFuture(forceClient, forceRpc, process);

        forceClient.forceStop().get();

        verify(forceRpc, never()).invoke(eq("runtime.shutdown"), any(), eq(Void.class));
        verify(process).destroyForcibly();

        var externalClient = new CopilotClient(new CopilotClientOptions().setAutoStart(false));
        var externalRpc = mock(JsonRpcClient.class);
        setConnectionFuture(externalClient, externalRpc, null);

        externalClient.stop().get();

        verify(externalRpc, never()).invoke(eq("runtime.shutdown"), any(), eq(Void.class));
    }

    @Test
    @SuppressWarnings("unchecked")
    void testForceStopCancelsPendingExternalTools() throws Exception {
        var client = new CopilotClient(new CopilotClientOptions().setAutoStart(false));
        var rpc = mock(JsonRpcClient.class);
        setConnectionFuture(client, rpc, null);
        var session = new CopilotSession("force-stop-session", rpc);
        var toolFuture = new CompletableFuture<Object>();
        var started = new CountDownLatch(1);
        var lateStarted = new CountDownLatch(1);
        var invocations = new AtomicInteger();
        session.registerTools(List.of(ToolDefinition.create("blocked_tool", "Blocks", Map.of(), invocation -> {
            if (invocations.incrementAndGet() == 1) {
                started.countDown();
            } else {
                lateStarted.countDown();
            }
            return toolFuture;
        })));
        Field sessionsField = CopilotClient.class.getDeclaredField("sessions");
        sessionsField.setAccessible(true);
        var sessions = (Map<String, CopilotSession>) sessionsField.get(client);
        sessions.put(session.getSessionId(), session);

        var requested = new ExternalToolRequestedEvent();
        requested.setData(new ExternalToolRequestedEvent.ExternalToolRequestedEventData("request-force-stop",
                session.getSessionId(), "tool-call-force-stop", "blocked_tool", null, Map.of(), null, null, null));
        session.dispatchEvent(requested);
        assertTrue(started.await(1, TimeUnit.SECONDS));

        client.forceStop().get();

        assertThrows(CancellationException.class, () -> toolFuture.get(1, TimeUnit.SECONDS));
        assertTrue(sessions.isEmpty());

        var lateRequest = new ExternalToolRequestedEvent();
        lateRequest.setData(new ExternalToolRequestedEvent.ExternalToolRequestedEventData("request-after-force-stop",
                session.getSessionId(), "tool-call-after-force-stop", "blocked_tool", null, Map.of(), null, null,
                null));
        session.dispatchEvent(lateRequest);
        assertFalse(lateStarted.await(100, TimeUnit.MILLISECONDS));
    }

    @Test
    @SuppressWarnings("unchecked")
    void testDeleteSessionReleasesGitHubTokenProvider() throws Exception {
        var client = new CopilotClient(new CopilotClientOptions().setAutoStart(false));
        var rpc = mock(JsonRpcClient.class);
        when(rpc.invoke(eq("session.delete"), any(), eq(DeleteSessionResponse.class)))
                .thenReturn(CompletableFuture.completedFuture(new DeleteSessionResponse(true, null)));
        when(rpc.invoke(eq("session.detach"), any(), eq(CopilotSession.SessionDetachResponse.class)))
                .thenReturn(CompletableFuture.completedFuture(new CopilotSession.SessionDetachResponse(true, null)));
        setConnectionFuture(client, rpc, null);

        var registry = new GitHubTokenProviderRegistry();
        var registration = registry
                .register(args -> CompletableFuture.completedFuture(GitHubTokenProviderResult.cancelled()));
        var session = new CopilotSession("delete-session", rpc);
        session.setGitHubTokenProviderRegistration(registration);
        Field sessionsField = CopilotClient.class.getDeclaredField("sessions");
        sessionsField.setAccessible(true);
        ((Map<String, CopilotSession>) sessionsField.get(client)).put(session.getSessionId(), session);

        try {
            client.deleteSession(session.getSessionId()).join();
            assertNull(registry.get(registration.id()));
        } finally {
            session.close();
            client.close();
        }
    }

    @Test
    @AllowCopilotExperimental
    void cloudSessionWithSkillProviderIsRejectedBeforeConnectingOrCallingProvider() throws Exception {
        var providerCalls = new AtomicInteger();
        SkillProvider provider = new SkillProvider() {
            @Override
            public CompletableFuture<List<SkillProviderDescriptor>> listSkills() {
                providerCalls.incrementAndGet();
                return CompletableFuture.completedFuture(List.of());
            }

            @Override
            public CompletableFuture<String> readSkill(String name) {
                providerCalls.incrementAndGet();
                return CompletableFuture.completedFuture(null);
            }
        };

        try (var client = new CopilotClient(new CopilotClientOptions().setAutoStart(false))) {
            var config = new SessionConfig().setOnPermissionRequest(PermissionHandler.APPROVE_ALL)
                    .setCloud(new CloudSessionOptions()).setSkillProvider(provider);

            ExecutionException error = assertThrows(ExecutionException.class, () -> client.createSession(config).get());
            assertInstanceOf(IllegalArgumentException.class, error.getCause());
            assertEquals("Skill providers are not supported for cloud sessions.", error.getCause().getMessage());
            assertEquals(0, providerCalls.get());
        }
    }

    @Test
    @AllowCopilotExperimental
    void emptyModeCreateWithSkillProviderSendsFlagsAndRegistersBeforeRpc() throws Exception {
        var client = new CopilotClient(new CopilotClientOptions().setAutoStart(false).setMode(CopilotClientMode.EMPTY)
                .setCopilotHome(System.getProperty("java.io.tmpdir")));
        var rpc = mock(JsonRpcClient.class);
        setConnectionFuture(client, rpc, null);
        SkillProvider provider = testSkillProvider();
        var createRequest = new AtomicReference<CreateSessionRequest>();

        when(rpc.invoke(eq("session.create"), any(), eq(CreateSessionResponse.class))).thenAnswer(invocation -> {
            CreateSessionRequest request = invocation.getArgument(1);
            createRequest.set(request);
            assertTrue(request.getHasSkillProvider());
            assertFalse(request.getEnableSkills());
            assertSame(provider, registeredSkillProvider(client, "skill-provider-session"));
            return CompletableFuture
                    .completedFuture(new CreateSessionResponse("skill-provider-session", null, null, null));
        });
        when(rpc.invoke(eq("session.options.update"), any(), eq(Void.class)))
                .thenReturn(CompletableFuture.completedFuture(null));

        try {
            var config = new SessionConfig().setSessionId("skill-provider-session")
                    .setOnPermissionRequest(PermissionHandler.APPROVE_ALL).setAvailableTools(List.of())
                    .setSkillProvider(provider);

            CopilotSession session = client.createSession(config).get();

            assertSame(provider, session.getSkillProvider());
            assertNotNull(createRequest.get());
            assertTrue(createRequest.get().getHasSkillProvider());
            assertFalse(createRequest.get().getEnableSkills());
        } finally {
            client.forceStop().get();
        }
    }

    @Test
    @AllowCopilotExperimental
    void emptyModeResumeValidationFailureKeepsResidentSkillProvider() throws Exception {
        var client = new CopilotClient(new CopilotClientOptions().setAutoStart(false).setMode(CopilotClientMode.EMPTY)
                .setCopilotHome(System.getProperty("java.io.tmpdir")));
        var rpc = mock(JsonRpcClient.class);
        setConnectionFuture(client, rpc, null);
        SkillProvider resident = testSkillProvider();
        when(rpc.invoke(eq("session.create"), any(), eq(CreateSessionResponse.class))).thenReturn(
                CompletableFuture.completedFuture(new CreateSessionResponse("resident-session", null, null, null)));
        when(rpc.invoke(eq("session.options.update"), any(), eq(Void.class)))
                .thenReturn(CompletableFuture.completedFuture(null));

        try {
            client.createSession(new SessionConfig().setSessionId("resident-session")
                    .setOnPermissionRequest(PermissionHandler.APPROVE_ALL).setAvailableTools(List.of())
                    .setSkillProvider(resident)).get();

            var resumeConfig = new ResumeSessionConfig().setOnPermissionRequest(PermissionHandler.APPROVE_ALL)
                    .setSkillProvider(testSkillProvider());
            var error = assertThrows(ExecutionException.class,
                    () -> client.resumeSession("resident-session", resumeConfig).get());

            assertInstanceOf(IllegalArgumentException.class, error.getCause());
            assertSame(resident, registeredSkillProvider(client, "resident-session"));
            verify(rpc, never()).invoke(eq("session.resume"), any(), any());
        } finally {
            client.forceStop().get();
        }
    }

    @Test
    @AllowCopilotExperimental
    void failedResumeDoesNotEvictAConcurrentResume() throws Exception {
        var client = new CopilotClient(
                new CopilotClientOptions().setAutoStart(false).setCopilotHome(System.getProperty("java.io.tmpdir")));
        var rpc = mock(JsonRpcClient.class);
        setConnectionFuture(client, rpc, null);
        when(rpc.invoke(eq("session.create"), any(), eq(CreateSessionResponse.class))).thenReturn(
                CompletableFuture.completedFuture(new CreateSessionResponse("shared-session", null, null, null)));
        when(rpc.invoke(eq("session.options.update"), any(), eq(Void.class)))
                .thenReturn(CompletableFuture.completedFuture(null));
        var firstResume = new CompletableFuture<ResumeSessionResponse>();
        when(rpc.invoke(eq("session.resume"), any(), eq(ResumeSessionResponse.class))).thenReturn(firstResume,
                CompletableFuture.completedFuture(new ResumeSessionResponse("shared-session", null, null, null)));

        try {
            client.createSession(new SessionConfig().setSessionId("shared-session")
                    .setOnPermissionRequest(PermissionHandler.APPROVE_ALL).setSkillProvider(testSkillProvider())).get();
            var failing = client.resumeSession("shared-session", new ResumeSessionConfig()
                    .setOnPermissionRequest(PermissionHandler.APPROVE_ALL).setSkillProvider(testSkillProvider()));
            SkillProvider winner = testSkillProvider();
            client.resumeSession("shared-session", new ResumeSessionConfig()
                    .setOnPermissionRequest(PermissionHandler.APPROVE_ALL).setSkillProvider(winner)).get();

            firstResume.completeExceptionally(new RuntimeException("resume failed"));

            assertThrows(ExecutionException.class, failing::get);
            assertSame(winner, registeredSkillProvider(client, "shared-session"));
        } finally {
            client.forceStop().get();
        }
    }

    @Test
    void testClientConstruction() {
        var client = new CopilotClient();
        assertEquals(ConnectionState.DISCONNECTED, client.getState());
        client.close();
    }

    @Test
    void testClientConstructionWithOptions() {
        var options = new CopilotClientOptions().setCliPath("/path/to/cli").setLogLevel("debug").setAutoStart(false);

        var client = new CopilotClient(options);
        assertEquals(ConnectionState.DISCONNECTED, client.getState());
        client.close();
    }

    @Test
    void testCliUrlAutoCorrectsUseStdio() {
        var options = new CopilotClientOptions().setCliUrl("localhost:3000").setUseStdio(true);

        // Should NOT throw - useStdio is auto-corrected to false when cliUrl is set
        var client = new CopilotClient(options);
        assertFalse(options.isUseStdio(), "useStdio should be auto-corrected to false when cliUrl is set");
        client.close();
    }

    @Test
    void testCliUrlOnlyConstruction() {
        var options = new CopilotClientOptions().setCliUrl("localhost:4321");

        // Should work without explicitly setting useStdio to false
        var client = new CopilotClient(options);
        assertEquals(ConnectionState.DISCONNECTED, client.getState());
        assertFalse(options.isUseStdio(), "useStdio should be auto-corrected to false when cliUrl is set");
        client.close();
    }

    @Test
    void testBracketedIpv6CliUrlNormalizesHost() throws Exception {
        try (var client = new CopilotClient(new CopilotClientOptions().setCliUrl("[::1]:4321"))) {
            Field hostField = CopilotClient.class.getDeclaredField("optionsHost");
            hostField.setAccessible(true);
            assertEquals("::1", hostField.get(client));
        }
    }

    @Test
    void testCliUrlMutualExclusionWithCliPath() {
        var options = new CopilotClientOptions().setCliUrl("localhost:3000").setCliPath("/path/to/cli");

        assertThrows(IllegalArgumentException.class, () -> new CopilotClient(options));
    }

    @Test
    void testStartAndConnectUsingStdio() throws Exception {
        assertNotNull(cliPath, "Copilot CLI not found in PATH or COPILOT_CLI_PATH");

        try (var client = new CopilotClient(new CopilotClientOptions().setCliPath(cliPath).setUseStdio(true))) {
            client.start().get();
            assertEquals(ConnectionState.CONNECTED, client.getState());

            PingResponse pong = client.ping("test message").get();
            assertEquals("pong: test message", pong.message());
            assertNotNull(pong.timestamp());

            client.stop().get();
            assertEquals(ConnectionState.DISCONNECTED, client.getState());
        }
    }

    @Test
    void testShouldReportErrorWithStderrWhenCliFailsToStart() throws Exception {
        assertNotNull(cliPath, "Copilot CLI not found in PATH or COPILOT_CLI_PATH");

        var options = new CopilotClientOptions().setCliPath(cliPath)
                .setCliArgs(new String[]{"--nonexistent-flag-for-testing"}).setUseStdio(true);

        try (var client = new CopilotClient(options)) {
            Exception ex = assertThrows(Exception.class, () -> client.start().get());
            Throwable root = ex instanceof ExecutionException && ex.getCause() != null ? ex.getCause() : ex;
            String message = root.getMessage();
            assertNotNull(message);
            assertTrue(message.contains("--nonexistent-flag-for-testing"),
                    "Error should include the CLI's invalid-flag diagnostic: " + message);
        }
    }

    @Test
    void testStartupWriteFailureDrainsStderrBeforeKillingChild(@TempDir Path tempDir) throws Exception {
        var script = tempDir.resolve("startup-error.js");
        Files.writeString(script, """
                const fs = require('node:fs');
                fs.watch(__dirname, () => {
                    if (fs.existsSync(__dirname + '/release-stderr')) {
                        fs.writeSync(2, 'error: --nonexistent-flag-for-testing\\n');
                        process.exit(1);
                    }
                });
                process.stdout.write('started\\n');
                """);
        var options = new CopilotClientOptions().setCliPath(script.toString()).setUseStdio(true);
        var manager = spy(new CliServerManager(options));
        var child = new AtomicReference<Process>();
        var failedOutput = new AtomicReference<OutputStream>();
        var writeFailure = new IOException("Controlled startup write failure");
        doAnswer(invocation -> {
            CliServerManager.ProcessInfo info = (CliServerManager.ProcessInfo) invocation.callRealMethod();
            child.set(info.process());
            // Closing fd 0 does not reliably fail the first parent write on Windows.
            var output = new FilterOutputStream(info.process().getOutputStream()) {
                @Override
                public void write(int value) throws IOException {
                    throw writeFailure;
                }

                @Override
                public void write(byte[] bytes, int offset, int length) throws IOException {
                    throw writeFailure;
                }

                @Override
                public void close() throws IOException {
                    // No write succeeds, so there is nothing to flush on this failed transport.
                    out.close();
                }
            };
            failedOutput.set(output);
            var transportProcess = mock(Process.class, org.mockito.AdditionalAnswers.delegatesTo(info.process()));
            doReturn(output).when(transportProcess).getOutputStream();
            return new CliServerManager.ProcessInfo(transportProcess, info.port());
        }).when(manager).startCliServer();
        doAnswer(invocation -> {
            Process process = invocation.getArgument(0);
            var reader = new BufferedReader(new InputStreamReader(process.getInputStream(), StandardCharsets.UTF_8));
            assertEquals("started", reader.readLine());
            return invocation.callRealMethod();
        }).when(manager).connectToServer(any(Process.class), isNull(), isNull());
        doAnswer(invocation -> {
            // Release stderr only when startup drains it, while the child is still alive.
            assertTrue(child.get().isAlive(), "Startup must drain stderr before killing the controlled child.");
            Files.createFile(tempDir.resolve("release-stderr"));
            return invocation.callRealMethod();
        }).when(manager).awaitStderrReader();

        try (AutoCloseable childCleanup = () -> {
            var process = child.get();
            if (process != null && process.isAlive()) {
                process.destroyForcibly();
                assertTrue(process.waitFor(10, TimeUnit.SECONDS), "Could not reap the controlled child.");
            }
        }; var client = new CopilotClient(options); AutoCloseable outputCleanup = () -> {
            var output = failedOutput.get();
            if (output != null) {
                output.close();
            }
        }) {
            Field managerField = CopilotClient.class.getDeclaredField("serverManager");
            managerField.setAccessible(true);
            managerField.set(client, manager);

            var error = assertThrows(ExecutionException.class, () -> client.start().get(10, TimeUnit.SECONDS));
            var cause = error.getCause();
            assertTrue(cause.getMessage().contains("--nonexistent-flag-for-testing"),
                    "Startup lost the child's stderr: " + cause);
            assertInstanceOf(IOException.class, cause);
            Throwable original = cause;
            while (original.getCause() != null) {
                original = original.getCause();
            }
            assertSame(writeFailure, original, "Startup must preserve the actual failed transport write.");
            assertNotSame(cause, original, "The startup error must retain its original transport failure.");
            assertFalse(child.get().isAlive(), "Failed startup retained its child process.");
        }
    }

    @Test
    void testProtocolMismatchTerminatesLiveChildBeforeDrainingStderr(@TempDir Path tempDir) throws Exception {
        var script = tempDir.resolve("incompatible-server.js");
        Files.writeString(script, """
                let input = Buffer.alloc(0);
                process.stdin.on('data', chunk => {
                    input = Buffer.concat([input, chunk]);
                    const headerEnd = input.indexOf('\\r\\n\\r\\n');
                    if (headerEnd < 0) return;
                    const length = Number(/Content-Length:\\s*(\\d+)/i.exec(input.toString())[1]);
                    if (input.length < headerEnd + 4 + length) return;
                    const request = JSON.parse(input.subarray(headerEnd + 4, headerEnd + 4 + length));
                    const response = JSON.stringify({
                        jsonrpc: '2.0', id: request.id, result: { protocolVersion: 1 }
                    });
                    process.stdout.write(`Content-Length: ${Buffer.byteLength(response)}\\r\\n\\r\\n${response}`);
                    input = input.subarray(headerEnd + 4 + length);
                });
                setInterval(() => {}, 1000);
                """);
        var options = new CopilotClientOptions().setCliPath(script.toString()).setUseStdio(true);
        var manager = spy(new CliServerManager(options));
        var child = new AtomicReference<Process>();
        doAnswer(invocation -> {
            child.set(invocation.getArgument(0));
            return invocation.callRealMethod();
        }).when(manager).connectToServer(any(Process.class), isNull(), isNull());
        doAnswer(invocation -> {
            assertFalse(child.get().isAlive(), "A rejected live server must exit before its stderr reader is joined.");
            return invocation.callRealMethod();
        }).when(manager).awaitStderrReader();

        try (var client = new CopilotClient(options)) {
            Field managerField = CopilotClient.class.getDeclaredField("serverManager");
            managerField.setAccessible(true);
            managerField.set(client, manager);

            var error = assertThrows(ExecutionException.class, () -> client.start().get(10, TimeUnit.SECONDS));
            assertTrue(error.getCause().getMessage().contains("SDK protocol version mismatch"),
                    "Startup must report the incompatible protocol: " + error.getCause());
            assertFalse(child.get().isAlive(), "Failed startup retained its child process.");
            verify(manager).awaitStderrReader();
        } finally {
            if (child.get() != null && child.get().isAlive()) {
                child.get().destroyForcibly();
                assertTrue(child.get().waitFor(10, TimeUnit.SECONDS), "Could not reap the incompatible server.");
            }
        }
    }

    @Test
    void testStartAndConnectUsingTcp() throws Exception {
        assertNotNull(cliPath, "Copilot CLI not found in PATH or COPILOT_CLI_PATH");

        try (var client = new CopilotClient(new CopilotClientOptions().setCliPath(cliPath).setUseStdio(false))) {
            client.start().get();
            assertEquals(ConnectionState.CONNECTED, client.getState());

            PingResponse pong = client.ping("test message").get();
            assertEquals("pong: test message", pong.message());

            client.stop().get();
        }
    }

    @Test
    void testForceStopWithoutCleanup() throws Exception {
        assertNotNull(cliPath, "Copilot CLI not found in PATH or COPILOT_CLI_PATH");

        try (var client = new CopilotClient(new CopilotClientOptions().setCliPath(cliPath))) {
            client.createSession(new SessionConfig().setOnPermissionRequest(PermissionHandler.APPROVE_ALL)).get();
            client.forceStop().get();

            assertEquals(ConnectionState.DISCONNECTED, client.getState());
        }
    }

    @Test
    void testGitHubTokenOptionAccepted() {
        var options = new CopilotClientOptions().setCliPath("/path/to/cli").setGitHubToken("gho_test_token");

        assertEquals("gho_test_token", options.getGitHubToken());
    }

    @Test
    void testUseLoggedInUserDefaultsToNull() {
        var options = new CopilotClientOptions().setCliPath("/path/to/cli");

        assertTrue(options.getUseLoggedInUser().isEmpty());
    }

    @Test
    void testExplicitUseLoggedInUserFalse() {
        var options = new CopilotClientOptions().setCliPath("/path/to/cli").setUseLoggedInUser(false);

        assertEquals(Optional.of(false), options.getUseLoggedInUser());
    }

    @Test
    void testExplicitUseLoggedInUserTrueWithGitHubToken() {
        var options = new CopilotClientOptions().setCliPath("/path/to/cli").setGitHubToken("gho_test_token")
                .setUseLoggedInUser(true);

        assertEquals(Optional.of(true), options.getUseLoggedInUser());
    }

    @Test
    void testGitHubTokenWithCliUrlThrows() {
        var options = new CopilotClientOptions().setCliUrl("localhost:8080").setGitHubToken("gho_test_token");

        assertThrows(IllegalArgumentException.class, () -> new CopilotClient(options));
    }

    @Test
    void testUseLoggedInUserWithCliUrlThrows() {
        var options = new CopilotClientOptions().setCliUrl("localhost:8080").setUseLoggedInUser(false);

        assertThrows(IllegalArgumentException.class, () -> new CopilotClient(options));
    }

    @Test
    void testSessionIdleTimeoutSecondsDefaultsToNull() {
        var options = new CopilotClientOptions();

        assertTrue(options.getSessionIdleTimeoutSeconds().isEmpty());
    }

    @Test
    void testSessionIdleTimeoutSecondsOptionAccepted() {
        var options = new CopilotClientOptions().setSessionIdleTimeoutSeconds(600);

        assertEquals(600, options.getSessionIdleTimeoutSeconds().getAsInt());
    }

    @Test
    void testTcpConnectionTokenWithUseStdioThrows() {
        var options = new CopilotClientOptions().setUseStdio(true).setTcpConnectionToken("my-token");

        assertThrows(IllegalArgumentException.class, () -> new CopilotClient(options));
    }

    @Test
    void testTcpConnectionTokenAcceptedInTcpMode() {
        var options = new CopilotClientOptions().setUseStdio(false).setTcpConnectionToken("my-token");

        // Should not throw
        try (var client = new CopilotClient(options)) {
            assertNotNull(client);
        }
    }

    @Test
    void testCopilotHomeOptionSetOnOptions() {
        var options = new CopilotClientOptions().setCopilotHome("/custom/home");

        assertEquals("/custom/home", options.getCopilotHome());
    }

    // ===== onLifecycle tests =====

    /**
     * Gets the internal LifecycleEventManager from a CopilotClient via reflection
     * so we can dispatch events for testing.
     */
    private static LifecycleEventManager getLifecycleManager(CopilotClient client) throws Exception {
        Field f = CopilotClient.class.getDeclaredField("lifecycleManager");
        f.setAccessible(true);
        return (LifecycleEventManager) f.get(client);
    }

    private static SessionLifecycleEvent lifecycleEvent(String type) {
        var e = new SessionLifecycleEvent();
        e.setType(type);
        e.setSessionId("test-session-id");
        return e;
    }

    @Test
    void testOnLifecycleWildcardReceivesAllEvents() throws Exception {
        try (var client = new CopilotClient()) {
            var received = new ArrayList<SessionLifecycleEvent>();
            client.onLifecycle(received::add);

            LifecycleEventManager mgr = getLifecycleManager(client);
            mgr.dispatch(lifecycleEvent(SessionLifecycleEventTypes.CREATED));
            mgr.dispatch(lifecycleEvent(SessionLifecycleEventTypes.DELETED));

            assertEquals(2, received.size());
            assertEquals(SessionLifecycleEventTypes.CREATED, received.get(0).getType());
            assertEquals(SessionLifecycleEventTypes.DELETED, received.get(1).getType());
        }
    }

    @Test
    void testOnLifecycleTypedReceivesOnlyMatchingEvents() throws Exception {
        try (var client = new CopilotClient()) {
            var received = new ArrayList<SessionLifecycleEvent>();
            client.onLifecycle(SessionLifecycleEventTypes.CREATED, received::add);

            LifecycleEventManager mgr = getLifecycleManager(client);
            mgr.dispatch(lifecycleEvent(SessionLifecycleEventTypes.CREATED));
            mgr.dispatch(lifecycleEvent(SessionLifecycleEventTypes.DELETED));

            assertEquals(1, received.size());
            assertEquals(SessionLifecycleEventTypes.CREATED, received.get(0).getType());
        }
    }

    @Test
    void testOnLifecycleUnsubscribeStopsDelivery() throws Exception {
        try (var client = new CopilotClient()) {
            var received = new ArrayList<SessionLifecycleEvent>();
            AutoCloseable sub = client.onLifecycle(received::add);

            LifecycleEventManager mgr = getLifecycleManager(client);
            mgr.dispatch(lifecycleEvent(SessionLifecycleEventTypes.CREATED));
            assertEquals(1, received.size());

            sub.close();

            mgr.dispatch(lifecycleEvent(SessionLifecycleEventTypes.DELETED));
            assertEquals(1, received.size(), "Should not receive events after unsubscribe");
        }
    }

    @Test
    void testOnLifecycleTypedUnsubscribeStopsDelivery() throws Exception {
        try (var client = new CopilotClient()) {
            var received = new ArrayList<SessionLifecycleEvent>();
            AutoCloseable sub = client.onLifecycle(SessionLifecycleEventTypes.UPDATED, received::add);

            LifecycleEventManager mgr = getLifecycleManager(client);
            mgr.dispatch(lifecycleEvent(SessionLifecycleEventTypes.UPDATED));
            assertEquals(1, received.size());

            sub.close();

            mgr.dispatch(lifecycleEvent(SessionLifecycleEventTypes.UPDATED));
            assertEquals(1, received.size(), "Should not receive events after unsubscribe");
        }
    }

    @Test
    void testOnLifecycleMultipleHandlers() throws Exception {
        try (var client = new CopilotClient()) {
            var wildcard = new ArrayList<SessionLifecycleEvent>();
            var typed = new ArrayList<SessionLifecycleEvent>();

            client.onLifecycle(wildcard::add);
            client.onLifecycle(SessionLifecycleEventTypes.CREATED, typed::add);

            LifecycleEventManager mgr = getLifecycleManager(client);
            mgr.dispatch(lifecycleEvent(SessionLifecycleEventTypes.CREATED));

            assertEquals(1, wildcard.size());
            assertEquals(1, typed.size());
        }
    }

    // ===== getState() coverage =====

    @Test
    void testGetStateErrorAfterFailedStart() throws Exception {
        // Use a non-existent CLI path to trigger a startup failure
        var options = new CopilotClientOptions().setCliPath("/nonexistent/path/to/cli").setAutoStart(false);

        try (var client = new CopilotClient(options)) {
            // Manually start to trigger the error
            CompletableFuture<Void> startFuture = client.start();

            // Wait for the start to fail
            try {
                startFuture.get();
            } catch (ExecutionException e) {
                // Expected
            }

            assertEquals(ConnectionState.ERROR, client.getState());
        }
    }

    @Test
    void testGetStateConnectingDuringStart() throws Exception {
        // Use a non-existent CLI path; the future won't complete immediately
        var options = new CopilotClientOptions().setCliPath("/nonexistent/path/to/cli").setAutoStart(false);

        try (var client = new CopilotClient(options)) {
            // Start is async - grab state before completion
            client.start();

            // The state should be either CONNECTING or ERROR depending on timing
            ConnectionState state = client.getState();
            assertTrue(state == ConnectionState.CONNECTING || state == ConnectionState.ERROR,
                    "State should be CONNECTING or ERROR, was: " + state);
        }
    }

    // ===== ensureConnected throws when autoStart=false and not connected =====

    @Test
    void testEnsureConnectedThrowsWhenNotStartedAndAutoStartDisabled() {
        var options = new CopilotClientOptions().setAutoStart(false);

        try (var client = new CopilotClient(options)) {
            // Calling ping (which calls ensureConnected) without start() should throw
            assertThrows(IllegalStateException.class, () -> client.ping("test"));
        }
    }

    // ===== close() idempotency =====

    @Test
    void testCloseIsIdempotent() {
        var client = new CopilotClient();

        // First close
        client.close();
        // Second close should not throw
        assertDoesNotThrow(() -> client.close());
    }

    @Test
    void testCloseAfterFailedStart() throws Exception {
        var options = new CopilotClientOptions().setCliPath("/nonexistent/path/to/cli").setAutoStart(false);
        var client = new CopilotClient(options);

        CompletableFuture<Void> startFuture = client.start();
        try {
            startFuture.get();
        } catch (ExecutionException e) {
            // Expected
        }

        // close() after a failed start should not throw
        assertDoesNotThrow(() -> client.close());
    }

    // ===== stop() with no connection =====

    @Test
    void testStopWithNoConnectionCompletes() throws Exception {
        try (var client = new CopilotClient(new CopilotClientOptions().setAutoStart(false))) {
            // stop() without start() should complete without error
            client.stop().get();
            assertEquals(ConnectionState.DISCONNECTED, client.getState());
        }
    }

    @Test
    void testForceStopWithNoConnectionCompletes() throws Exception {
        try (var client = new CopilotClient(new CopilotClientOptions().setAutoStart(false))) {
            // forceStop() without start() should complete without error
            client.forceStop().get();
            assertEquals(ConnectionState.DISCONNECTED, client.getState());
        }
    }

    @Test
    void testCloseSessionAfterStoppingClientDoesNotThrow() throws Exception {
        assertNotNull(cliPath, "Copilot CLI not found in PATH or COPILOT_CLI_PATH");

        try (var client = new CopilotClient(new CopilotClientOptions().setCliPath(cliPath))) {
            var session = client
                    .createSession(new SessionConfig().setOnPermissionRequest(PermissionHandler.APPROVE_ALL)).get();

            // Stop the client first (which closes the RPC connection)
            client.stop().get();

            // Then close the session - should not throw even though RPC is closed
            assertDoesNotThrow(() -> session.close(), "Closing session after client.stop() should not throw exception");

            // Verify session is terminated
            assertThrows(IllegalStateException.class, () -> session.send("test"),
                    "Session should be terminated after close()");
        }
    }

    // ===== start() idempotency =====

    @Test
    void testStartIsIdempotentSingleConnectionAttempt() throws Exception {
        var options = new CopilotClientOptions().setCliPath("/nonexistent/path/to/cli").setAutoStart(false);

        try (var client = new CopilotClient(options)) {
            client.start();
            client.start();

            // Both calls should result in the same state (single connection attempt)
            ConnectionState state = client.getState();
            assertTrue(state == ConnectionState.CONNECTING || state == ConnectionState.ERROR,
                    "State should be CONNECTING or ERROR after start(), was: " + state);
        }
    }

    // ===== null options defaulting =====

    @Test
    void testNullOptionsDefaultsToEmpty() {
        try (var client = new CopilotClient(null)) {
            assertEquals(ConnectionState.DISCONNECTED, client.getState());
        }
    }

    // ===== OnListModels =====

    @Test
    void testListModels_WithCustomHandler_CallsHandler() throws Exception {
        var customModels = new ArrayList<com.github.copilot.rpc.ModelInfo>();
        var model = new com.github.copilot.rpc.ModelInfo();
        model.setId("my-custom-model");
        customModels.add(model);

        var callCount = new int[]{0};
        var options = new CopilotClientOptions().setOnListModels(() -> {
            callCount[0]++;
            return CompletableFuture.completedFuture(new ArrayList<>(customModels));
        });

        try (var client = new CopilotClient(options)) {
            var models = client.listModels().get();
            assertEquals(1, callCount[0]);
            assertEquals(1, models.size());
            assertEquals("my-custom-model", models.get(0).getId());
        }
    }

    @Test
    void testListModels_WithCustomHandler_CachesResults() throws Exception {
        var customModels = new ArrayList<com.github.copilot.rpc.ModelInfo>();
        var model = new com.github.copilot.rpc.ModelInfo();
        model.setId("cached-model");
        customModels.add(model);

        var callCount = new int[]{0};
        var options = new CopilotClientOptions().setOnListModels(() -> {
            callCount[0]++;
            return CompletableFuture.completedFuture(new ArrayList<>(customModels));
        });

        try (var client = new CopilotClient(options)) {
            client.listModels().get();
            client.listModels().get();
            assertEquals(1, callCount[0], "Handler should be called only once due to caching");
        }
    }

    @Test
    void testListModels_WithCustomHandler_WorksWithoutStart() throws Exception {
        var customModels = new ArrayList<com.github.copilot.rpc.ModelInfo>();
        var model = new com.github.copilot.rpc.ModelInfo();
        model.setId("no-start-model");
        customModels.add(model);

        var callCount = new int[]{0};
        var options = new CopilotClientOptions().setOnListModels(() -> {
            callCount[0]++;
            return CompletableFuture.completedFuture(new ArrayList<>(customModels));
        });

        // No start() needed when onListModels is provided
        try (var client = new CopilotClient(options)) {
            var models = client.listModels().get();
            assertEquals(1, callCount[0]);
            assertEquals(1, models.size());
            assertEquals("no-start-model", models.get(0).getId());
        }
    }

    private static void setConnectionFuture(CopilotClient client, JsonRpcClient rpc, Process process) throws Exception {
        var connectionClass = Class.forName("com.github.copilot.CopilotClient$Connection");
        var constructor = connectionClass.getDeclaredConstructor(JsonRpcClient.class, Process.class,
                com.github.copilot.generated.rpc.ServerRpc.class, AutoCloseable.class);
        constructor.setAccessible(true);
        var connection = constructor.newInstance(rpc, process, null, null);

        Field field = CopilotClient.class.getDeclaredField("connectionFuture");
        field.setAccessible(true);
        field.set(client, CompletableFuture.completedFuture(connection));
    }

    @AllowCopilotExperimental
    private static SkillProvider testSkillProvider() {
        return new SkillProvider() {
            @Override
            public CompletableFuture<List<SkillProviderDescriptor>> listSkills() {
                return CompletableFuture.completedFuture(List.of());
            }

            @Override
            public CompletableFuture<String> readSkill(String name) {
                return CompletableFuture.completedFuture(null);
            }
        };
    }

    @AllowCopilotExperimental
    @SuppressWarnings("unchecked")
    private static SkillProvider registeredSkillProvider(CopilotClient client, String sessionId) throws Exception {
        Field sessionsField = CopilotClient.class.getDeclaredField("sessions");
        sessionsField.setAccessible(true);
        Map<String, CopilotSession> sessions = (Map<String, CopilotSession>) sessionsField.get(client);
        CopilotSession session = sessions.get(sessionId);
        return session != null ? session.getSkillProvider() : null;
    }
}