mcp: finish wiring up the 'version' on ext MCP collections (#247096)

This commit is contained in:
Connor Peet
2025-04-22 08:18:23 -07:00
committed by GitHub
parent 544674f984
commit 4ccec8917a
5 changed files with 50 additions and 26 deletions
@@ -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)
});
}
@@ -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();
@@ -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;
}));
}));
}
@@ -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<ObservablePromise<readonly IValidatedMcpTool[]> | undefined>(this, undefined);
private readonly toolsFromServerPromise = observableValue<ObservablePromise<{
readonly tools: IValidatedMcpTool[];
readonly nonce: string | undefined;
}> | undefined>(this, undefined);
private readonly toolsFromServer = derived(reader => this.toolsFromServerPromise.read(reader)?.promiseResult.read(reader)?.data);
public readonly tools: IObservable<readonly IMcpTool[]>;
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<ServerBootData, ServerBootClassification>('mcp/serverBoot', {
supportsLogging: !!handler.capabilities.logging,
supportsPrompts: !!handler.capabilities.prompts,
@@ -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 */