diff --git a/src/vs/workbench/api/common/extHostMcp.ts b/src/vs/workbench/api/common/extHostMcp.ts index e8ccfc48f4d..a163ff04aa0 100644 --- a/src/vs/workbench/api/common/extHostMcp.ts +++ b/src/vs/workbench/api/common/extHostMcp.ts @@ -112,6 +112,7 @@ export class ExtHostMcpService extends Disposable implements IExtHostMpcService servers.push({ id: ExtensionIdentifier.toKey(extension.identifier), label: item.label, + cacheNonce: item.version, launch: Convert.McpServerDefinition.from(item) }); } diff --git a/src/vs/workbench/contrib/mcp/browser/mcpCommands.ts b/src/vs/workbench/contrib/mcp/browser/mcpCommands.ts index 6f536d3a7c1..519b08137ef 100644 --- a/src/vs/workbench/contrib/mcp/browser/mcpCommands.ts +++ b/src/vs/workbench/contrib/mcp/browser/mcpCommands.ts @@ -247,6 +247,7 @@ export class MCPServerActionRendering extends Disposable implements IWorkbenchCo let thisState = DisplayedState.None; switch (server.toolsState.read(reader)) { case McpServerToolsState.Unknown: + case McpServerToolsState.Outdated: if (server.trusted.read(reader) === false) { thisState = DisplayedState.None; } else { @@ -324,7 +325,7 @@ export class MCPServerActionRendering extends Disposable implements IWorkbenchCo const { state, servers } = displayedState.get(); if (state === DisplayedState.NewTools) { - servers.forEach(server => server.start()); + servers.forEach(server => server.stop().then(() => server.start())); mcpService.activateCollections(); } else if (state === DisplayedState.Refreshing) { servers.at(-1)?.showOutput(); diff --git a/src/vs/workbench/contrib/mcp/common/mcpContextKeys.ts b/src/vs/workbench/contrib/mcp/common/mcpContextKeys.ts index 0d0aff9aa6d..3cbb821058c 100644 --- a/src/vs/workbench/contrib/mcp/common/mcpContextKeys.ts +++ b/src/vs/workbench/contrib/mcp/common/mcpContextKeys.ts @@ -56,7 +56,7 @@ export class McpContextKeysController extends Disposable implements IWorkbenchCo } const toolState = s.toolsState.read(r); - return toolState === McpServerToolsState.Unknown || toolState === McpServerToolsState.RefreshingFromUnknown; + return toolState === McpServerToolsState.Unknown || toolState === McpServerToolsState.Outdated || toolState === McpServerToolsState.RefreshingFromUnknown; })); })); } diff --git a/src/vs/workbench/contrib/mcp/common/mcpServer.ts b/src/vs/workbench/contrib/mcp/common/mcpServer.ts index 6545737fb4a..3bc03f06603 100644 --- a/src/vs/workbench/contrib/mcp/common/mcpServer.ts +++ b/src/vs/workbench/contrib/mcp/common/mcpServer.ts @@ -55,6 +55,7 @@ type ServerBootStateClassification = { }; interface IToolCacheEntry { + readonly nonce: string | undefined; /** Cached tools so we can show what's available before it's started */ readonly tools: readonly IValidatedMcpTool[]; } @@ -109,13 +110,13 @@ export class McpServerMetadataCache extends Disposable { } /** Gets cached tools for a server (used before a server is running) */ - getTools(definitionId: string): readonly IValidatedMcpTool[] | undefined { - return this.cache.get(definitionId)?.tools; + getTools(definitionId: string) { + return this.cache.get(definitionId); } /** Sets cached tools for a server */ - storeTools(definitionId: string, tools: readonly IValidatedMcpTool[]): void { - this.cache.set(definitionId, { ...this.cache.get(definitionId), tools }); + storeTools(definitionId: string, nonce: string | undefined, tools: readonly IValidatedMcpTool[]): void { + this.cache.set(definitionId, { ...this.cache.get(definitionId), nonce, tools }); this.didChange = true; } @@ -154,17 +155,33 @@ export class McpServer extends Disposable implements IMcpServer { private get toolsFromCache() { return this._toolCache.getTools(this.definition.id); } - private readonly toolsFromServerPromise = observableValue | undefined>(this, undefined); + private readonly toolsFromServerPromise = observableValue | undefined>(this, undefined); private readonly toolsFromServer = derived(reader => this.toolsFromServerPromise.read(reader)?.promiseResult.read(reader)?.data); public readonly tools: IObservable; public readonly toolsState = derived(reader => { + const currentNonce = () => this._mcpRegistry.collections.read(reader) + .find(c => c.id === this.collection.id) + ?.serverDefinitions.read(reader) + .find(d => d.id === this.definition.id) + ?.cacheNonce; + const stateWhenServingFromCache = () => { + if (!this.toolsFromCache) { + return McpServerToolsState.Unknown; + } + + return currentNonce() === this.toolsFromCache.nonce ? McpServerToolsState.Cached : McpServerToolsState.Outdated; + }; + const fromServer = this.toolsFromServerPromise.read(reader); const connectionState = this.connectionState.read(reader); const isIdle = McpConnectionState.canBeStarted(connectionState.state) && !fromServer; if (isIdle) { - return this.toolsFromCache ? McpServerToolsState.Cached : McpServerToolsState.Unknown; + return stateWhenServingFromCache(); } const fromServerResult = fromServer?.promiseResult.read(reader); @@ -172,7 +189,11 @@ export class McpServer extends Disposable implements IMcpServer { return this.toolsFromCache ? McpServerToolsState.RefreshingFromCached : McpServerToolsState.RefreshingFromUnknown; } - return fromServerResult.error ? (this.toolsFromCache ? McpServerToolsState.Cached : McpServerToolsState.Unknown) : McpServerToolsState.Live; + if (fromServerResult.error) { + return stateWhenServingFromCache(); + } + + return fromServerResult.data?.nonce === currentNonce() ? McpServerToolsState.Live : McpServerToolsState.Outdated; }); private readonly _loggerId: string; @@ -228,27 +249,20 @@ export class McpServer extends Disposable implements IMcpServer { // 2. Populate this.tools when we connect to a server. this._register(autorunWithStore((reader, store) => { - const cnx = this._connection.read(reader)?.handler.read(reader); - if (cnx) { - this.populateLiveData(cnx, store); + const cnx = this._connection.read(reader); + const handler = cnx?.handler.read(reader); + if (handler) { + this.populateLiveData(handler, cnx?.definition.cacheNonce, store); } else { this.resetLiveData(); } })); - // 3. Update the cache when tools update - this._register(autorun(reader => { - const tools = this.toolsFromServer.read(reader); - if (tools) { - this._toolCache.storeTools(definition.id, tools); - } - })); - - // 4. Publish tools + // 3. Publish tools const toolPrefix = this._mcpRegistry.collectionToolPrefix(this.collection); this.tools = derived(reader => { const serverTools = this.toolsFromServer.read(reader); - const definitions = serverTools ?? this.toolsFromCache ?? []; + const definitions = serverTools?.tools ?? this.toolsFromCache?.tools ?? []; const prefix = toolPrefix.read(reader); return definitions.map(def => new McpTool(this, prefix, def)).sort((a, b) => a.compare(b)); }); @@ -384,7 +398,7 @@ export class McpServer extends Disposable implements IMcpServer { return validated; } - private populateLiveData(handler: McpServerRequestHandler, store: DisposableStore) { + private populateLiveData(handler: McpServerRequestHandler, cacheNonce: string | undefined, store: DisposableStore) { const cts = new CancellationTokenSource(); store.add(toDisposable(() => cts.dispose(true))); @@ -394,11 +408,11 @@ export class McpServer extends Disposable implements IMcpServer { const toolPromise = handler.capabilities.tools ? handler.listTools({}, cts.token) : Promise.resolve([]); const toolPromiseSafe = toolPromise.then(async tools => { handler.logger.info(`Discovered ${tools.length} tools`); - return this._getValidatedTools(handler, tools); + return { tools: await this._getValidatedTools(handler, tools), nonce: cacheNonce }; }); this.toolsFromServerPromise.set(new ObservablePromise(toolPromiseSafe), tx); - return [toolPromise]; + return [toolPromiseSafe]; }; store.add(handler.onDidChangeToolList(() => { @@ -411,7 +425,9 @@ export class McpServer extends Disposable implements IMcpServer { promises = updateTools(tx); }); - Promise.all(promises!).then(([tools]) => { + Promise.all(promises!).then(([{ tools }]) => { + this._toolCache.storeTools(this.definition.id, cacheNonce, tools); + this._telemetryService.publicLog2('mcp/serverBoot', { supportsLogging: !!handler.capabilities.logging, supportsPrompts: !!handler.capabilities.prompts, diff --git a/src/vs/workbench/contrib/mcp/common/mcpTypes.ts b/src/vs/workbench/contrib/mcp/common/mcpTypes.ts index a590925c9a8..258040db89d 100644 --- a/src/vs/workbench/contrib/mcp/common/mcpTypes.ts +++ b/src/vs/workbench/contrib/mcp/common/mcpTypes.ts @@ -104,6 +104,8 @@ export interface McpServerDefinition { readonly roots?: URI[] | undefined; /** If set, allows configuration variables to be resolved in the {@link launch} with the given context */ readonly variableReplacement?: McpServerDefinitionVariableReplacement; + /** Nonce used for caching the server. Changing the nonce will indicate that tools need to be refreshed. */ + readonly cacheNonce?: string; readonly presentation?: { /** Sort order of the definition. */ @@ -117,6 +119,7 @@ export namespace McpServerDefinition { export interface Serialized { readonly id: string; readonly label: string; + readonly cacheNonce?: string; readonly launch: McpServerLaunch.Serialized; readonly variableReplacement?: McpServerDefinitionVariableReplacement.Serialized; } @@ -129,6 +132,7 @@ export namespace McpServerDefinition { return { id: def.id, label: def.label, + cacheNonce: def.cacheNonce, launch: McpServerLaunch.fromSerialized(def.launch), variableReplacement: def.variableReplacement ? McpServerDefinitionVariableReplacement.fromSerialized(def.variableReplacement) : undefined, }; @@ -233,6 +237,8 @@ export const enum McpServerToolsState { Unknown, /** Tools were read from the cache */ Cached, + /** Tools were read from the cache or live, but they may be outdated. */ + Outdated, /** Tools are refreshing for the first time */ RefreshingFromUnknown, /** Tools are refreshing and the current tools are cached */