mcp: persist extension server caches when they don't register

Author: connor4312

src/vs/workbench/contrib/mcp/common/discovery/extensionMcpDiscovery.ts

@@ -25,8 +25,13 @@ interface IServerCacheEntry {
 
 const _mcpExtensionPoint = extensionsRegistry.ExtensionsRegistry.registerExtensionPoint(mcpContributionPoint);
 
+const enum PersistWhen {
+	CollectionExists,
+	Always,
+}
+
 export class ExtensionMcpDiscovery extends Disposable implements IMcpDiscovery {
-	private readonly _extensionCollectionIdsToPersist = new Set<string>();
+	private readonly _extensionCollectionIdsToPersist = new Map<string, PersistWhen>();
 	private readonly cachedServers: { [collcetionId: string]: IServerCacheEntry };
 
 	constructor(
@@ -39,13 +44,17 @@ export class ExtensionMcpDiscovery extends Disposable implements IMcpDiscovery {
 
 		this._register(storageService.onWillSaveState(() => {
 			let updated = false;
-			for (const collectionId of this._extensionCollectionIdsToPersist) {
+			for (const [collectionId, behavior] of this._extensionCollectionIdsToPersist.entries()) {
 				const collection = this._mcpRegistry.collections.get().find(c => c.id === collectionId);
+				let defs = collection?.serverDefinitions.get();
 				if (!collection || collection.lazy) {
-					continue;
+					if (behavior === PersistWhen.Always) {
+						defs = [];
+					} else {
+						continue;
+					}
 				}
 
-				const defs = collection.serverDefinitions.get();
 				if (defs) {
 					updated = true;
 					this.cachedServers[collectionId] = { servers: defs.map(McpServerDefinition.toSerialized) };
@@ -77,7 +86,7 @@ export class ExtensionMcpDiscovery extends Disposable implements IMcpDiscovery {
 
 				for (const coll of collections.value) {
 					const id = extensionPrefixedIdentifier(collections.description.identifier, coll.id);
-					this._extensionCollectionIdsToPersist.add(id);
+					this._extensionCollectionIdsToPersist.set(id, PersistWhen.CollectionExists);
 
 					const serverDefs = this.cachedServers.hasOwnProperty(id) ? this.cachedServers[id].servers : undefined;
 					const dispo = this._mcpRegistry.registerCollection({
@@ -90,7 +99,10 @@ export class ExtensionMcpDiscovery extends Disposable implements IMcpDiscovery {
 						serverDefinitions: observableValue<McpServerDefinition[]>(this, serverDefs?.map(McpServerDefinition.fromSerialized) || []),
 						lazy: {
 							isCached: !!serverDefs,
-							load: () => this._activateExtensionServers(coll.id),
+							load: () => this._activateExtensionServers(coll.id).then(() => {
+								// persist (an empty collection) in case the extension doesn't end up publishing one
+								this._extensionCollectionIdsToPersist.set(id, PersistWhen.Always);
+							}),
 							removed: () => extensionCollections.deleteAndDispose(id),
 						},
 						source: collections.description.identifier