MCPDownloadSnippetsExecuteStep.ts

/*---------------------------------------------------------------------------------------------
 *  Copyright (c) Microsoft Corporation. All rights reserved.
 *  Licensed under the MIT License. See License.txt in the project root for license information.
 *--------------------------------------------------------------------------------------------*/

import { AzExtFsExtra, AzureWizardExecuteStepWithActivityOutput, type IActionContext } from '@microsoft/vscode-azext-utils';
import * as path from 'path';
import { type Progress } from 'vscode';
import { hostFileName, ProjectLanguage, type GitHubFileMetadata } from '../../../constants';
import { localize } from "../../../localize";
import { feedUtils } from '../../../utils/feedUtils';
import { nonNullProp } from '../../../utils/nonNull';
import { parseJson } from '../../../utils/parseJson';
import { requestUtils } from '../../../utils/requestUtils';
import { type MCPProjectWizardContext } from '../IProjectWizardContext';

export class MCPDownloadSnippetsExecuteStep extends AzureWizardExecuteStepWithActivityOutput<MCPProjectWizardContext> {
    stepName: string = 'MCPDownloadSampleCodeExecuteStep';
    protected getTreeItemLabel(context: MCPProjectWizardContext): string {
        return localize('downloadSampleCode', 'Downloading {0} sample server code', nonNullProp(context, 'serverLanguage'));
    }
    protected getOutputLogSuccess(context: MCPProjectWizardContext): string {
        return localize('downloadSampleCodeSuccess', 'Successfully downloaded {0} sample server code', nonNullProp(context, 'serverLanguage'));
    }
    protected getOutputLogFail(context: MCPProjectWizardContext): string {
        return localize('downloadSampleCodeFail', 'Failed to download {0} sample server code', nonNullProp(context, 'serverLanguage'));
    }
    protected getOutputLogProgress(context: MCPProjectWizardContext): string {
        return localize('downloadingSampleCode', 'Downloading {0} sample server code...', nonNullProp(context, 'serverLanguage'));
    }

    public priority: number = 12;

    public async execute(context: MCPProjectWizardContext, progress: Progress<{ message?: string | undefined; increment?: number | undefined }>): Promise<void> {
        progress.report({ message: localize('downloadingSampleCodeExecute', 'Downloading sample server code...') });
        const sampleMcpRepoUrl: string = nonNullProp(context, 'sampleMcpRepoUrl');
        const sampleFiles: GitHubFileMetadata[] = await feedUtils.getJsonFeed(context, sampleMcpRepoUrl);
        await this.downloadFilesRecursively(context, sampleFiles, context.projectPath);
    }

    private static readonly languageExtensions: Partial<Record<ProjectLanguage, string>> = {
        [ProjectLanguage.TypeScript]: '.ts',
        [ProjectLanguage.Python]: '.py',
        [ProjectLanguage.CSharp]: '.cs',
    };

    private async downloadFilesRecursively(context: MCPProjectWizardContext, items: GitHubFileMetadata[], basePath: string): Promise<void> {
        const sourceExt = context.serverLanguage ? MCPDownloadSnippetsExecuteStep.languageExtensions[context.serverLanguage] : undefined;
        // Download all files and directories at this level in parallel
        await Promise.all(items.map(async (item) => {
            if (item.type === 'file') {
                const destPath = await MCPDownloadSnippetsExecuteStep.downloadSingleFile({
                    context, item,
                    destDirPath: basePath,
                    serverLanguage: context.serverLanguage,
                    projectName: path.basename(context.projectPath)
                });
                // Track the first source file matching the server language as the sample tool file
                if (!context.sampleToolFilePath && sourceExt && destPath.endsWith(sourceExt)) {
                    context.sampleToolFilePath = destPath;
                }
            } else if (item.type === 'dir') {
                // Create directory
                const dirPath: string = path.join(basePath, item.name);
                await AzExtFsExtra.ensureDir(dirPath);

                // Get directory contents
                const response = await requestUtils.sendRequestWithExtTimeout(context, { method: 'GET', url: item.url });
                const dirContents: GitHubFileMetadata[] = parseJson<GitHubFileMetadata[]>(nonNullProp(response, 'bodyAsText'));

                // Recursively download directory contents
                await this.downloadFilesRecursively(context, dirContents, dirPath);
            }
        }));
    }

    public shouldExecute(context: MCPProjectWizardContext): boolean {
        return context.includeSnippets === true;
    }

    public static async downloadSingleFile(options: {
        context: IActionContext,
        item: GitHubFileMetadata,
        destDirPath: string,
        projectName: string
        serverLanguage?: ProjectLanguage,
    }): Promise<string> {
        const { context, item, destDirPath, serverLanguage, projectName } = options;
        const fileUrl: string = item.download_url;
        let destinationPath: string = path.join(destDirPath, item.name);
        const response = await requestUtils.sendRequestWithExtTimeout(context, { method: 'GET', url: fileUrl });
        let fileContent = response.bodyAsText;
        if (serverLanguage === ProjectLanguage.CSharp) {
            if (item.name === hostFileName) {
                // for C#, we need to replace the host.json
                // "arguments": ["<path to the compiled DLL, e.g. HelloWorld.dll>"] with the name of the actual DLL
                const dllName: string = `${projectName}.dll`;
                fileContent = fileContent?.replace('<path to the compiled DLL, e.g. HelloWorld.dll>', dllName);
            } else if (item.name === 'dotnet.csproj') {
                // for C#, the project file needs to be named after the project folder
                const csprojFileName: string = `${projectName}.csproj`;
                destinationPath = path.join(destDirPath, csprojFileName);
            }
        }
        await AzExtFsExtra.writeFile(destinationPath, fileContent ?? '');
        return destinationPath;
    }
}