sdkTestContext.ts

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

import fs, { realpathSync } from "fs";
import { rm } from "fs/promises";
import os from "os";
import { basename, dirname, join, resolve } from "path";
import { rimraf } from "rimraf";
import { fileURLToPath } from "url";
import { afterAll, afterEach, beforeEach, onTestFailed, TestContext } from "vitest";
import { CopilotClient, CopilotClientOptions, RuntimeConnection } from "../../../src";
import { CapiProxy } from "./CapiProxy";
import { formatError, retry } from "./sdkTestHelper";

export const isCI = process.env.GITHUB_ACTIONS === "true";
export const DEFAULT_GITHUB_TOKEN = "fake-token-for-e2e-tests";

const __filename = fileURLToPath(import.meta.url);
const __dirname = dirname(__filename);
const SNAPSHOTS_DIR = resolve(__dirname, "../../../../test/snapshots");

function getCliPathForTests(): string | undefined {
    if (process.env.COPILOT_CLI_PATH) {
        return process.env.COPILOT_CLI_PATH;
    }
    return undefined;
}

export async function createSdkTestContext({
    logLevel,
    useStdio,
    copilotClientOptions,
}: {
    logLevel?: "error" | "none" | "warning" | "info" | "debug" | "all";
    cliPath?: string;
    useStdio?: boolean;
    copilotClientOptions?: CopilotClientOptions;
} = {}) {
    const homeDir = realpathSync(fs.mkdtempSync(join(os.tmpdir(), "copilot-test-config-")));
    const copilotHomeDir = realpathSync(fs.mkdtempSync(join(os.tmpdir(), "copilot-test-home-")));
    const workDir = realpathSync(fs.mkdtempSync(join(os.tmpdir(), "copilot-test-work-")));

    const openAiEndpoint = new CapiProxy();
    const proxyUrl = await openAiEndpoint.start();
    await openAiEndpoint.setCopilotUserByToken(DEFAULT_GITHUB_TOKEN, {
        login: "e2e-test-user",
        copilot_plan: "individual_pro",
        is_mcp_enabled: true,
        endpoints: {
            api: proxyUrl,
            telemetry: "https://localhost:1/telemetry",
        },
        analytics_tracking_id: "e2e-test-tracking-id",
    });
    const authTokenToUse = isCI
        ? DEFAULT_GITHUB_TOKEN
        : (process.env.GITHUB_TOKEN ?? DEFAULT_GITHUB_TOKEN);

    const env = {
        ...process.env,
        ...openAiEndpoint.getProxyEnv(),
        COPILOT_API_URL: proxyUrl,
        // Route GitHub API calls (e.g. the MCP registry policy check) to the
        // replay proxy so MCP enablement stays hermetic. Without this the CLI
        // reaches the real api.github.com, which is slow/unreachable on macOS
        // CI runners and makes MCP servers time out before reaching connected.
        COPILOT_DEBUG_GITHUB_API_URL: proxyUrl,
        COPILOT_HOME: copilotHomeDir,
        COPILOT_SDK_AUTH_TOKEN: "",
        GH_CONFIG_DIR: homeDir,
        GH_TOKEN: "",
        GITHUB_TOKEN: "",

        // TODO: I'm not convinced the SDK should default to using whatever config you happen to have in your homedir.
        // The SDK config should be independent of the regular CLI app. Likewise it shouldn't mix sessions from the
        // SDK with those from the CLI app, at least not by default.
        XDG_CONFIG_HOME: homeDir,
        XDG_STATE_HOME: homeDir,
    };

    const userConn = copilotClientOptions?.connection;
    const cliPath = getCliPathForTests();
    let connection: RuntimeConnection;
    if (userConn) {
        // Caller supplied a RuntimeConnection — merge in the harness-managed
        // CLI path (and stay on the same transport variant). Strip `kind`
        // before forwarding to the factory opts since the factories don't
        // accept it in their argument shape.
        if (userConn.kind === "tcp") {
            const { kind: _k, ...tcp } = userConn;
            connection = RuntimeConnection.forTcp({
                ...tcp,
                path: tcp.path ?? cliPath,
            });
        } else if (userConn.kind === "stdio") {
            const { kind: _k, ...stdio } = userConn;
            connection = RuntimeConnection.forStdio({
                ...stdio,
                path: stdio.path ?? cliPath,
            });
        } else {
            connection = userConn;
        }
    } else {
        connection =
            useStdio === false
                ? RuntimeConnection.forTcp({ path: cliPath })
                : RuntimeConnection.forStdio({ path: cliPath });
    }

    const {
        connection: _ignoredConnection,
        env: userEnv,
        ...remainingClientOptions
    } = copilotClientOptions ?? {};
    const copilotClient = new CopilotClient({
        workingDirectory: workDir,
        env: { ...env, ...userEnv },
        logLevel: logLevel || "error",
        connection,
        gitHubToken: authTokenToUse,
        ...remainingClientOptions,
    });

    const harness = { homeDir, workDir, openAiEndpoint, copilotClient, env };

    // Track if any test fails to avoid writing corrupted snapshots
    let anyTestFailed = false;

    // Wire up to Vitest lifecycle
    beforeEach(async (testContext) => {
        // Must be inside beforeEach - vitest requires test context
        onTestFailed(() => {
            anyTestFailed = true;
        });

        await openAiEndpoint.updateConfig({
            filePath: getTrafficCapturePath(testContext),
            workDir,
            testInfo: {
                file: testContext.task.file.filepath,
                line: testContext.task.location?.line,
            },
        });
    });

    afterEach(async () => {
        // Empty directories but leave them in place for next test
        await rimraf([join(homeDir, "*"), join(workDir, "*")], { glob: true });
    });

    afterAll(async () => {
        await copilotClient.stop();
        await openAiEndpoint.stop(anyTestFailed);
        await rmDir("remove e2e test copilotHomeDir", copilotHomeDir);
        await rmDir("remove e2e test homeDir", homeDir);
        await rmDir("remove e2e test workDir", workDir);
    });

    return harness;
}

function getTrafficCapturePath(testContext: TestContext): string {
    const testFilePath = testContext.task.file.filepath;
    const suffix = ".test.ts";
    if (!testFilePath.endsWith(suffix)) {
        throw new Error(
            `Test file path does not end with expected suffix '${suffix}': ${testFilePath}`
        );
    }

    // Convert to snake_case for cross-SDK snapshot compatibility
    // Strip ".e2e" suffix so renamed "xxx.e2e.test.ts" still uses snapshot folder "xxx"
    let testFileName = basename(testFilePath, suffix).replace(/-/g, "_");
    if (testFileName.endsWith(".e2e")) {
        testFileName = testFileName.slice(0, -".e2e".length);
    }
    const taskNameAsFilename = testContext.task.name.replace(/[^a-z0-9]/gi, "_").toLowerCase();
    return join(SNAPSHOTS_DIR, testFileName, `${taskNameAsFilename}.yaml`);
}

async function rmDir(message: string, path: string): Promise<void> {
    // Use longer retries to tolerate Windows holding SQLite session-store.db
    // open briefly after the CLI subprocess exits. If the temp dir still can't
    // be removed (e.g. CLI background writer racing with cleanup), warn and
    // continue rather than failing the whole test run — the OS / CI runner
    // will reclaim the temp dir on shutdown.
    try {
        await retry(message, () => rm(path, { recursive: true, force: true }), 30, 1000);
    } catch (error) {
        console.warn(
            `WARN: ${message} failed; leaving temp dir for OS cleanup: ${formatError(error)}`
        );
    }
}