Skip to content

Commit edfa43f

Browse files
vritant24vritant24
committed
feat: improve MCP server migration with enhanced error handling and file existence checks
1 parent 0383e07 commit edfa43f

6 files changed

Lines changed: 168 additions & 19 deletions

File tree

‎src/vs/workbench/contrib/chat/browser/aiCustomization/aiCustomizationManagementEditor.ts‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1416,7 +1416,10 @@ export class AICustomizationManagementEditor extends EditorPane {
14161416
const { migratedCount, failedServerNames } = await migrateMcpServers(
14171417
currentServers,
14181418
this.fileService,
1419-
error => this.logService.error('[MCP Migration] Failed to migrate a server', error),
1419+
error => {
1420+
this.logService.error('[MCP Migration] Failed to migrate a server', error);
1421+
onUnexpectedError(error);
1422+
},
14201423
);
14211424
const allFailedServerNames = [...noLongerEligibleServerNames, ...failedServerNames];
14221425
if (noLongerEligibleServerNames.length > 0) {
@@ -1494,7 +1497,9 @@ export class AICustomizationManagementEditor extends EditorPane {
14941497
: customization.name ?? basename(customization.uri);
14951498
const renderSelectionCheckbox = (row: HTMLElement, customization: CustomizationMigrationCandidate, onSelectionChange?: () => void): Checkbox => {
14961499
const checkboxContainer = DOM.append(row, $('.item-sync-checkbox.prompt-migration-checkbox'));
1497-
const checkboxTitle = localize('customizationMigrationSelectAriaLabel', "Select {0}", getCandidateName(customization));
1500+
const checkboxTitle = isMcpServerCustomizationMigrationCandidate(customization)
1501+
? localize('mcpMigrationSelectAriaLabel', "Select {0} from {1}", customization.name, this.labelService.getUriLabel(customization.sourceUri, { relative: true }))
1502+
: localize('customizationMigrationSelectAriaLabel', "Select {0}", getCandidateName(customization));
14981503
const checkbox = this.migrationPageDisposables.add(new Checkbox(checkboxTitle, this.isCustomizationSelectedForMigration(customization), defaultCheckboxStyles));
14991504
checkboxContainer.replaceChildren(checkbox.domNode);
15001505
this.migrationFirstFocusableElement ??= checkbox.domNode;

‎src/vs/workbench/contrib/chat/browser/aiCustomization/customizationMigration.ts‎

Lines changed: 17 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ import { URI } from '../../../../../base/common/uri.js';
88
import { VSBuffer } from '../../../../../base/common/buffer.js';
99
import { basename, dirname, getComparisonKey, isEqual } from '../../../../../base/common/resources.js';
1010
import { ResourceMap } from '../../../../../base/common/map.js';
11-
import { FileOperationResult, IFileService, IFileStatWithMetadata, toFileOperationResult } from '../../../../../platform/files/common/files.js';
11+
import { FileOperationError, FileOperationResult, IFileService, IFileStatWithMetadata, toFileOperationResult } from '../../../../../platform/files/common/files.js';
1212
import { getCleanPromptName, getPromptFileExtension, SKILL_FILENAME, VALID_SKILL_NAME_REGEX } from '../../common/promptSyntax/config/promptFileLocations.js';
1313
import { IHeaderAttribute, ParsedPromptFile, PromptFileParser, PromptHeaderAttributes } from '../../common/promptSyntax/promptFileParser.js';
1414
import { getCustomizationMigrationTargetType, IMcpServerCustomizationMigrationCandidate, MigratableConfiguration } from '../../common/promptSyntax/service/customizationMigrationService.js';
@@ -382,7 +382,8 @@ async function migrateMcpServerGroup(
382382
continue;
383383
}
384384

385-
if (Object.hasOwn(targetServers, candidate.name) && !equals(targetServers[candidate.name], migrationConfiguration)) {
385+
const targetConfiguration = canonicalizeMcpServerMigrationSourceConfiguration(targetServers[candidate.name]);
386+
if (Object.hasOwn(targetServers, candidate.name) && (!targetConfiguration || !equals(targetConfiguration, migrationConfiguration))) {
386387
failedServerNames.push(candidate.name);
387388
onMigrationError?.(new Error(`MCP server '${candidate.name}' already exists with a different configuration in ${group.targetUri.toString()}.`));
388389
continue;
@@ -422,6 +423,7 @@ async function migrateMcpServerGroup(
422423
if (writtenTarget) {
423424
try {
424425
if (target.exists) {
426+
await ensureFileExists(group.targetUri, fileService);
425427
await fileService.writeFile(group.targetUri, VSBuffer.fromString(target.content), {
426428
etag: writtenTarget.etag,
427429
mtime: writtenTarget.mtime,
@@ -441,6 +443,7 @@ async function migrateMcpServerGroup(
441443
} catch (verificationError) {
442444
// Two files cannot be updated atomically; restore the guarded source if another writer changed the target.
443445
try {
446+
await ensureFileExists(group.sourceUri, fileService);
444447
await fileService.writeFile(group.sourceUri, VSBuffer.fromString(source.content), {
445448
etag: writtenSource.etag,
446449
mtime: writtenSource.mtime,
@@ -462,7 +465,10 @@ async function verifyMigratedMcpServers(
462465
const target = await readTargetJsonDocument(targetUri, fileService);
463466
const targetServers = getObjectProperty(target.value, 'mcpServers')!;
464467
for (const candidate of candidates) {
465-
if (!equals(targetServers[candidate.name], canonicalizeMcpServerMigrationConfiguration(candidate.configuration))) {
468+
if (!equals(
469+
canonicalizeMcpServerMigrationSourceConfiguration(targetServers[candidate.name]),
470+
canonicalizeMcpServerMigrationConfiguration(candidate.configuration),
471+
)) {
466472
throw new Error(`MCP server '${candidate.name}' changed in ${targetUri.toString()} during migration.`);
467473
}
468474
}
@@ -506,8 +512,9 @@ async function readTargetJsonDocument(resource: URI, fileService: IFileService):
506512
}
507513
}
508514

509-
function writeJsonDocument(resource: URI, content: string, document: IJsonDocument, fileService: IFileService): Promise<IFileStatWithMetadata> {
515+
async function writeJsonDocument(resource: URI, content: string, document: IJsonDocument, fileService: IFileService): Promise<IFileStatWithMetadata> {
510516
if (document.exists) {
517+
await ensureFileExists(resource, fileService);
511518
return fileService.writeFile(resource, VSBuffer.fromString(content), {
512519
etag: document.etag,
513520
mtime: document.mtime,
@@ -516,6 +523,12 @@ function writeJsonDocument(resource: URI, content: string, document: IJsonDocume
516523
return fileService.createFile(resource, VSBuffer.fromString(content), { overwrite: false });
517524
}
518525

526+
async function ensureFileExists(resource: URI, fileService: IFileService): Promise<void> {
527+
if (!await fileService.exists(resource)) {
528+
throw new FileOperationError(`File was deleted during MCP migration: ${resource.toString()}`, FileOperationResult.FILE_NOT_FOUND);
529+
}
530+
}
531+
519532
function getObjectProperty(value: Record<string, unknown>, key: string): Record<string, unknown> | undefined {
520533
const property = value[key];
521534
return isJsonObject(property) ? property : undefined;

‎src/vs/workbench/contrib/chat/browser/aiCustomization/customizationMigrationServiceImpl.ts‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ import { equals } from '../../../../../base/common/objects.js';
1010
import { dirname } from '../../../../../base/common/resources.js';
1111
import { URI } from '../../../../../base/common/uri.js';
1212
import { localize } from '../../../../../nls.js';
13-
import { IFileService } from '../../../../../platform/files/common/files.js';
13+
import { FileOperationResult, IFileService, toFileOperationResult } from '../../../../../platform/files/common/files.js';
1414
import { isAgentHostSessionResource } from '../../common/chatSessionsService.js';
1515
import { ICustomizationHarnessService, ICustomizationSourceFolder } from '../../common/customizationHarnessService.js';
1616
import { getChatSessionType } from '../../common/model/chatUri.js';
@@ -193,7 +193,15 @@ export class CustomizationMigrationService implements ICustomizationMigrationSer
193193
}
194194

195195
private async readMcpServers(resource: URI): Promise<Record<string, unknown> | undefined> {
196-
const content = (await this.fileService.readFile(resource)).value.toString();
196+
let content: string;
197+
try {
198+
content = (await this.fileService.readFile(resource)).value.toString();
199+
} catch (error) {
200+
if (toFileOperationResult(error) === FileOperationResult.FILE_NOT_FOUND) {
201+
return undefined;
202+
}
203+
throw error;
204+
}
197205
const errors: ParseError[] = [];
198206
const value = parse(content, errors, { allowTrailingComma: true, allowEmptyContent: false });
199207
if (errors.length > 0 || !value || typeof value !== 'object' || Array.isArray(value)) {

‎src/vs/workbench/contrib/chat/test/browser/aiCustomization/aiCustomizationManagementEditor.test.ts‎

Lines changed: 12 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -365,16 +365,16 @@ suite('aiCustomizationManagementEditor', () => {
365365
const serverA: IMcpServerCustomizationMigrationCandidate = {
366366
type: CustomizationMigrationType.McpServers,
367367
id: 'mcp.config.ws0.server-a',
368-
name: 'server-a',
369-
sourceUri: URI.file('/workspace/.vscode/mcp.json'),
370-
targetUri: URI.file('/workspace/.mcp.json'),
371-
configuration: { type: McpServerType.LOCAL, command: 'server-a' },
368+
name: 'server',
369+
sourceUri: URI.file('/workspace-a/.vscode/mcp.json'),
370+
targetUri: URI.file('/workspace-a/.mcp.json'),
371+
configuration: { type: McpServerType.LOCAL, command: 'server' },
372372
};
373373
const serverB: IMcpServerCustomizationMigrationCandidate = {
374374
...serverA,
375375
id: 'mcp.config.ws0.server-b',
376-
name: 'server-b',
377-
configuration: { type: McpServerType.LOCAL, command: 'server-b' },
376+
sourceUri: URI.file('/workspace-b/.vscode/mcp.json'),
377+
targetUri: URI.file('/workspace-b/.mcp.json'),
378378
};
379379
editor.customizationsByMigrationCategory = new Map([[CustomizationMigrationCategoryId.McpServers, [serverA, serverB]]]);
380380
editor.activeMigrationCategoryId = CustomizationMigrationCategoryId.McpServers;
@@ -390,10 +390,13 @@ suite('aiCustomizationManagementEditor', () => {
390390
button: editor.migrationMigrateButton,
391391
}, {
392392
rowText: [
393-
'server-a/workspace/.vscode/mcp.json to /workspace/.mcp.json',
394-
'server-b/workspace/.vscode/mcp.json to /workspace/.mcp.json',
393+
'server/workspace-a/.vscode/mcp.json to /workspace-a/.mcp.json',
394+
'server/workspace-b/.vscode/mcp.json to /workspace-b/.mcp.json',
395+
],
396+
selectAriaLabel: [
397+
'Select server from /workspace-a/.vscode/mcp.json',
398+
'Select server from /workspace-b/.vscode/mcp.json',
395399
],
396-
selectAriaLabel: ['Select server-a', 'Select server-b'],
397400
button: { enabled: true, label: 'Migrate 1' },
398401
});
399402
editor.editorPreviewDisposables.dispose();

‎src/vs/workbench/contrib/chat/test/browser/aiCustomization/customizationMigrationServiceImpl.test.ts‎

Lines changed: 62 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ import { constObservable } from '../../../../../../base/common/observable.js';
1111
import { URI } from '../../../../../../base/common/uri.js';
1212
import { mock } from '../../../../../../base/test/common/mock.js';
1313
import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../../../base/test/common/utils.js';
14-
import { IFileService } from '../../../../../../platform/files/common/files.js';
14+
import { FileOperationError, FileOperationResult, IFileService } from '../../../../../../platform/files/common/files.js';
1515
import { McpServerType } from '../../../../../../platform/mcp/common/mcpPlatformTypes.js';
1616
import { CustomizationMigrationService } from '../../../browser/aiCustomization/customizationMigrationServiceImpl.js';
1717
import { IAgentHostActiveClientService } from '../../../browser/agentSessions/agentHost/agentHostActiveClientService.js';
@@ -429,6 +429,67 @@ suite('CustomizationMigrationService', () => {
429429
assert.deepStrictEqual(migration.candidates, []);
430430
});
431431

432+
test('treats a missing MCP source file as having no migration candidates', async () => {
433+
const root = URI.file('/workspace');
434+
const snapshot: IAgentHostMcpServerSupportSnapshot = {
435+
servers: [{
436+
id: 'mcp.config.ws0.server',
437+
name: 'server',
438+
collectionId: 'mcp.config.ws0',
439+
source: {
440+
group: 'local',
441+
kind: AgentHostMcpServerSourceKind.VscodeWorkspaceFolder,
442+
label: 'Workspace',
443+
collectionUri: URI.joinPath(root, '.vscode', 'mcp.json'),
444+
definitionLocation: undefined,
445+
remoteAuthority: null,
446+
extensionId: undefined,
447+
pluginUri: undefined,
448+
},
449+
enablement: { enabled: true, state: AgentHostMcpServerEnablementState.EnabledWorkspace },
450+
applicability: AgentHostMcpServerApplicability.Applicable,
451+
delivery: AgentHostMcpServerDelivery.ClientForwarded,
452+
compatibility: { kind: 'supported' },
453+
migrationConfiguration: { type: McpServerType.LOCAL, command: 'server' },
454+
}],
455+
discoveryComplete: true,
456+
coverage: {
457+
restrictedByMcpAccess: false,
458+
restrictedByCustomizationPolicy: false,
459+
},
460+
};
461+
const activeClientService = {
462+
acquireMcpServerSupportScope: () => ({
463+
support: constObservable(snapshot),
464+
isResolved: constObservable(true),
465+
whenResolved: () => Promise.resolve(),
466+
dispose: () => { },
467+
}),
468+
} as Partial<IAgentHostActiveClientService> as IAgentHostActiveClientService;
469+
const agentHostCustomizationService = new class extends mock<IAgentHostCustomizationService>() {
470+
override getWorkingDirectories() { return [root.toString()]; }
471+
}();
472+
const missingFileService = {
473+
readFile: async () => {
474+
throw new FileOperationError('missing', FileOperationResult.FILE_NOT_FOUND);
475+
},
476+
} as Partial<IFileService> as IFileService;
477+
const service = new CustomizationMigrationService(
478+
store.add(new TestPromptsService([])),
479+
new TestCustomizationHarnessService(),
480+
activeClientService,
481+
agentHostCustomizationService,
482+
missingFileService,
483+
);
484+
485+
const migration = await service.computeMigration(
486+
URI.from({ scheme: SessionType.AgentHostCopilot, path: '/session' }),
487+
CustomizationMigrationType.McpServers,
488+
);
489+
490+
assert.deepStrictEqual(migration.candidates, []);
491+
});
492+
432493
test('uses the session harness label in migration hints', async () => {
433494
const promptsService = store.add(new TestPromptsService([
434495
{ uri: URI.file('/workspace/.github/prompts/review.prompt.md'), storage: PromptsStorage.local, type: PromptsType.prompt, source: PromptFileSource.GitHubWorkspace },

‎src/vs/workbench/contrib/chat/test/browser/aiCustomization/mcpServerMigration.test.ts‎

Lines changed: 60 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,18 @@ class TargetChangingFileSystemProvider extends InMemoryFileSystemProvider {
5050
}
5151
}
5252

53+
class DeletingBeforeExistsFileService extends FileService {
54+
deleteBeforeExists: URI | undefined;
55+
56+
override async exists(resource: URI): Promise<boolean> {
57+
if (this.deleteBeforeExists && isEqual(resource, this.deleteBeforeExists)) {
58+
this.deleteBeforeExists = undefined;
59+
await this.del(resource);
60+
}
61+
return super.exists(resource);
62+
}
63+
}
64+
5365
suite('mcpServerMigration', () => {
5466
const store = ensureNoDisposablesAreLeakedInTestSuite();
5567

@@ -202,7 +214,7 @@ suite('mcpServerMigration', () => {
202214
const sourceUri = URI.joinPath(root, '.vscode', 'mcp.json');
203215
const targetUri = URI.joinPath(root, '.mcp.json');
204216
await fileService.writeFile(sourceUri, VSBuffer.fromString('{"servers":{"server":{"type":"stdio","command":"source"}}}'));
205-
await fileService.writeFile(targetUri, VSBuffer.fromString('{"mcpServers":{"server":{"type":"stdio","command":"source"}}}'));
217+
await fileService.writeFile(targetUri, VSBuffer.fromString('{"mcpServers":{"server":{"type":"stdio","command":"source","args":[]}}}'));
206218

207219
const result = await migrateMcpServers([
208220
candidate(root, 'server-id', 'server', { type: McpServerType.LOCAL, command: 'source', args: undefined }),
@@ -215,6 +227,53 @@ suite('mcpServerMigration', () => {
215227
}, {
216228
result: { migratedCount: 1, failedServerNames: [] },
217229
source: { servers: {} },
230+
target: { mcpServers: { server: { type: 'stdio', command: 'source', args: [] } } },
231+
});
232+
});
233+
234+
test('does not recreate an existing target deleted during migration', async () => {
235+
const fileService = store.add(new DeletingBeforeExistsFileService(new NullLogService()));
236+
const provider = store.add(new InMemoryFileSystemProvider());
237+
store.add(fileService.registerProvider(Schemas.file, provider));
238+
const root = URI.file('/workspace-deleted-target');
239+
const sourceUri = URI.joinPath(root, '.vscode', 'mcp.json');
240+
const targetUri = URI.joinPath(root, '.mcp.json');
241+
await fileService.writeFile(sourceUri, VSBuffer.fromString('{"servers":{"server":{"type":"stdio","command":"source"}}}'));
242+
await fileService.writeFile(targetUri, VSBuffer.fromString('{"mcpServers":{}}'));
243+
fileService.deleteBeforeExists = targetUri;
244+
245+
const result = await migrateMcpServers([candidate(root, 'server-id', 'server')], fileService);
246+
247+
assert.deepStrictEqual({
248+
result,
249+
source: parse((await fileService.readFile(sourceUri)).value.toString()),
250+
targetExists: await fileService.exists(targetUri),
251+
}, {
252+
result: { migratedCount: 0, failedServerNames: ['server'] },
253+
source: { servers: { server: { type: 'stdio', command: 'source' } } },
254+
targetExists: false,
255+
});
256+
});
257+
258+
test('does not recreate an existing source deleted during migration', async () => {
259+
const fileService = store.add(new DeletingBeforeExistsFileService(new NullLogService()));
260+
const provider = store.add(new InMemoryFileSystemProvider());
261+
store.add(fileService.registerProvider(Schemas.file, provider));
262+
const root = URI.file('/workspace-deleted-source');
263+
const sourceUri = URI.joinPath(root, '.vscode', 'mcp.json');
264+
const targetUri = URI.joinPath(root, '.mcp.json');
265+
await fileService.writeFile(sourceUri, VSBuffer.fromString('{"servers":{"server":{"type":"stdio","command":"source"}}}'));
266+
fileService.deleteBeforeExists = sourceUri;
267+
268+
const result = await migrateMcpServers([candidate(root, 'server-id', 'server')], fileService);
269+
270+
assert.deepStrictEqual({
271+
result,
272+
sourceExists: await fileService.exists(sourceUri),
273+
target: parse((await fileService.readFile(targetUri)).value.toString()),
274+
}, {
275+
result: { migratedCount: 0, failedServerNames: ['server'] },
276+
sourceExists: false,
218277
target: { mcpServers: { server: { type: 'stdio', command: 'source' } } },
219278
});
220279
});

0 commit comments

Comments
 (0)