newWorkspaceFollowup.ts

/*---------------------------------------------------------------------------------------------
 *  Copyright (c) Microsoft Corporation. All rights reserved.
 *  Licensed under the MIT License. See License.txt in the project root for license information.
 *--------------------------------------------------------------------------------------------*/
import * as l10n from '@vscode/l10n';
import { ChatResponseFileTreePart, Disposable, MarkdownString, ProgressLocation, SaveDialogOptions, Tab, TabInputText, Uri, commands, env, interactive, window, workspace } from 'vscode';
import { IConversationOptions } from '../../../platform/chat/common/conversationOptions';
import { ILogService } from '../../../platform/log/common/logService';
import * as path from '../../../util/vs/base/common/path';
import { CopilotFileScheme, CopilotWorkspaceScheme, CreateFileCommand, CreateProjectCommand, GithubWorkspaceScheme, INewWorkspacePreviewContentManager, OpenFileCommand } from '../../intents/node/newIntent';
import { NewWorkspacePreviewFileSystemProvider } from '../../intents/vscode-node/newWorkspacePreviewFileSystemProvider';
import { NewWorkspaceTextDocumentProvider } from '../../intents/vscode-node/newWorkspaceTextDocumentProvider';
import { listFilesInResponseFileTree } from '../../prompt/common/fileTreeParser';

export function registerNewWorkspaceIntentCommand(previewContentManager: INewWorkspacePreviewContentManager, logService: ILogService, options: IConversationOptions) {
	const copilotWorkspaceProvider = new NewWorkspacePreviewFileSystemProvider(previewContentManager);
	const githubWorkspaceProvider = new NewWorkspacePreviewFileSystemProvider(previewContentManager);
	const copilotTextDocumentProvider = new NewWorkspaceTextDocumentProvider(previewContentManager);

	return Disposable.from(
		workspace.registerFileSystemProvider(CopilotWorkspaceScheme,
			copilotWorkspaceProvider,
			{ isReadonly: new MarkdownString(l10n.t('This file preview was generated by Copilot and may contain surprises or mistakes.\n\nAsk followup questions to refine it, then press Create Workspace.')) }),
		workspace.registerFileSystemProvider(GithubWorkspaceScheme,
			githubWorkspaceProvider,
			{ isReadonly: true }),
		workspace.registerTextDocumentContentProvider(CopilotFileScheme, copilotTextDocumentProvider),
		commands.registerCommand(CreateProjectCommand, async (fileTreePart: ChatResponseFileTreePart, workspaceRoot: Uri | undefined) => {
			const parentFolder = (await window.showOpenDialog({ defaultUri: workspaceRoot, title: path.basename(fileTreePart.baseUri.path), canSelectFolders: true, canSelectFiles: false, canSelectMany: false, openLabel: 'Select as Parent Folder' }))?.[0];
			if (!parentFolder) {
				return;
			}
			await createWorkspace(logService, workspaceRoot, parentFolder, fileTreePart);
		}),
		commands.registerCommand(OpenFileCommand, async (fileTreePart: ChatResponseFileTreePart) => {
			const pathStr = Uri.joinPath(fileTreePart.baseUri, fileTreePart.value[0].name).toString();
			const document = await workspace.openTextDocument(Uri.parse(pathStr));
			await window.showTextDocument(document, { preview: false });
		}),
		commands.registerCommand(CreateFileCommand, async (fileTreePart: ChatResponseFileTreePart) => {
			const options: SaveDialogOptions = {
				defaultUri: Uri.file(path.posix.join(workspace.workspaceFolders?.[0].uri.path ?? '', fileTreePart.value[0].name)),
				saveLabel: l10n.t('Save File'),
			};
			const uri = await window.showSaveDialog(options);
			if (uri) {
				const pathStr = Uri.joinPath(fileTreePart.baseUri, fileTreePart.value[0].name).toString();
				const document = await workspace.openTextDocument(Uri.parse(pathStr));
				await workspace.fs.writeFile(uri, Buffer.from(document.getText()));

				// Close out all previews since they won't properly restore
				const tabsToClose: Tab[] = [];
				window.tabGroups.all.forEach(group => {
					group.tabs.forEach((tab) => {
						if (tab.input instanceof TabInputText && tab.input.uri.scheme === CopilotFileScheme) {
							tabsToClose.push(tab);
						}
					});
				});
				window.tabGroups.close(tabsToClose, true);

				// re-open saved file
				const fileDoc = await workspace.openTextDocument(Uri.file(uri.fsPath));
				await window.showTextDocument(fileDoc);
			}
		}),
		copilotWorkspaceProvider,
		githubWorkspaceProvider,
		copilotTextDocumentProvider,
	);
}

async function createWorkspace(logService: ILogService, workspaceRoot: Uri | undefined, parentFolder: Uri, fileTreePart: ChatResponseFileTreePart) {
	// Close out all previews since they won't properly restore
	const tabsToClose: Tab[] = [];
	window.tabGroups.all.forEach(group => {
		group.tabs.forEach((tab) => {
			if (tab.input instanceof TabInputText && tab.input.uri.scheme === CopilotWorkspaceScheme) {
				tabsToClose.push(tab);
			}
		});
	});
	window.tabGroups.close(tabsToClose, true);

	// remove path separator from the beginning of the path
	const projectRoot = fileTreePart.baseUri.path.slice(1);
	const projectName = await getUniqueProjectName(parentFolder, projectRoot);
	const workspaceUri = Uri.joinPath(parentFolder, projectName);

	const files = listFilesInResponseFileTree(fileTreePart.value);
	if (files.length === 0) {
		return;
	}

	try {
		await window.withProgress({ location: ProgressLocation.Notification, cancellable: true }, async (progress, token) => {
			for (const file of files) {
				const relativeFilePath = path.relative(projectRoot, file);
				const fileUri = Uri.joinPath(workspaceUri, relativeFilePath);
				progress.report({ message: l10n.t(`Creating file {0}...`, fileUri.fsPath) });
				const content = await workspace.fs.readFile(Uri.joinPath(fileTreePart.baseUri, file));
				await workspace.fs.createDirectory(Uri.joinPath(fileUri, '..'));
				await workspace.fs.writeFile(fileUri, content);
			}
			await updateAiGeneratedWorkspacesFile(workspaceUri);
		});

		if (workspaceRoot && workspaceUri.fsPath.startsWith(workspaceRoot.fsPath + path.sep)) {
			// If the new workspace is a subfolder of the current workspace, do nothing
			return;
		}

		const message = l10n.t('Would you like to open the created workspace?');
		const open = l10n.t('Open');
		const openNewWindow = l10n.t('Open in New Window');
		const choices = [open, openNewWindow];
		const result = await window.showInformationMessage(message, { modal: true }, ...choices);
		if (result === open) {

			await interactive.transferActiveChat(workspaceUri);
			logService.info(
				'[newIntent] Opening folder: ' + workspaceUri.fsPath,
			);
			commands.executeCommand('vscode.openFolder', workspaceUri);
		} else if (result === openNewWindow) {
			commands.executeCommand('vscode.openFolder', workspaceUri, true);
		}
	}
	catch (error) {
		const errorMessage = l10n.t('Failed to create workspace: {0}', projectName);
		logService.error(error, errorMessage);
		window.showErrorMessage(errorMessage);
		await workspace.fs.delete(workspaceUri, { recursive: true });
	}
}


async function getUniqueProjectName(projectFolder: Uri, projectName: string): Promise<string> {
	let i = 0;
	let uniqueProjectNameNotFound = true;
	let newProjectName = projectName.replace(/^\W+/, '');
	while (uniqueProjectNameNotFound) {
		try {
			await workspace.fs.stat(Uri.joinPath(projectFolder, newProjectName));
			newProjectName = projectName + '-' + ++i;
		} catch {
			uniqueProjectNameNotFound = false;
		}
	}
	return newProjectName;
}

async function checkFileExists(filePath: Uri): Promise<boolean> {
	try {
		await workspace.fs.stat(filePath);
		return true;
	} catch (error) {
		return false;
	}
}

async function updateAiGeneratedWorkspacesFile(workspaceUris: Uri) {
	const aiGeneratedFilePath = getAiGeneratedWorkspacesFile();
	if (!aiGeneratedFilePath) {
		return;
	}

	if ((await checkFileExists(aiGeneratedFilePath))) {
		const fileContnet = await workspace.fs.readFile(aiGeneratedFilePath).then(b => { return new TextDecoder().decode(b); });
		const workspaces = JSON.parse(fileContnet) as string[];
		workspaces.push(workspaceUris.toString());
		await workspace.fs.writeFile(aiGeneratedFilePath, Buffer.from(JSON.stringify(workspaces, null, 2)));
	} else {
		await workspace.fs.writeFile(aiGeneratedFilePath, Buffer.from(JSON.stringify([workspaceUris.toString()], null, 2)));
	}
}

function getAiGeneratedWorkspacesFile(): Uri | undefined {
	const vscodeFolderName = env.appName.indexOf('Insider') > 0 || env.appName.indexOf('Code - OSS Dev') >= 0 ? 'Code - Insiders' : 'Code';
	const homeDir = Uri.file(process.env.HOME || (process.env.USERPROFILE ? process.env.USERPROFILE : ''));
	switch (process.platform) {
		case 'darwin':
			return Uri.joinPath(
				homeDir,
				'Library',
				'Application Support',
				vscodeFolderName,
				'User',
				'workspaceStorage',
				'aiGeneratedWorkspaces.json'
			);
		case 'linux':
			return Uri.joinPath(homeDir, '.config', vscodeFolderName, 'User', 'workspaceStorage', 'aiGeneratedWorkspaces.json');
		case 'win32':
			return process.env.APPDATA
				? Uri.joinPath(Uri.file(process.env.APPDATA), vscodeFolderName, 'User', 'workspaceStorage', 'aiGeneratedWorkspaces.json')
				: undefined;
		default:
			return;
	}
}