From ac9e4934d330b6f29a7dcf4bf8558a1f9837c4ca Mon Sep 17 00:00:00 2001 From: Charles Vien Date: Mon, 27 Jul 2026 12:19:33 -0700 Subject: [PATCH 1/2] discover MCP UI tools lazily for cloud runs --- packages/core/src/mcp-apps/mcp-apps.test.ts | 121 +++++++++++- packages/core/src/mcp-apps/mcp-apps.ts | 196 ++++++++++++++------ 2 files changed, 258 insertions(+), 59 deletions(-) diff --git a/packages/core/src/mcp-apps/mcp-apps.test.ts b/packages/core/src/mcp-apps/mcp-apps.test.ts index 5e86117696..18beb59c27 100644 --- a/packages/core/src/mcp-apps/mcp-apps.test.ts +++ b/packages/core/src/mcp-apps/mcp-apps.test.ts @@ -1,6 +1,6 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import { McpAppsService } from "./mcp-apps"; -import type { McpServerConnectionConfig } from "./schemas"; +import { McpAppsServiceEvent, type McpServerConnectionConfig } from "./schemas"; function makeLogger() { const scopedLog = { @@ -116,3 +116,122 @@ describe("McpAppsService config resolver", () => { expect(createConnection).toHaveBeenCalledTimes(2); }); }); + +const UI_MIME_TYPE = "text/html;profile=mcp-app"; +const REVIEW_URI = "ui://posthog/loops-review.html"; +const REVIEW_CSP = { connectDomains: ["https://us.posthog.com"] }; + +function makeClient() { + return { + listTools: vi.fn(async () => ({ + tools: [ + { + name: "loops-review", + _meta: { ui: { resourceUri: REVIEW_URI } }, + }, + { name: "loops-list" }, + ], + })), + listResources: vi.fn(async () => ({ + resources: [{ uri: REVIEW_URI, _meta: { ui: { csp: REVIEW_CSP } } }], + })), + readResource: vi.fn(async ({ uri }: { uri: string }) => ({ + contents: [{ uri, mimeType: UI_MIME_TYPE, text: "" }], + })), + }; +} + +function connectClient(service: McpAppsService, client = makeClient()) { + vi.spyOn(internals(service), "createConnection").mockImplementation( + async (c) => ({ name: c.name, client, transport: {} }), + ); + return client; +} + +describe("McpAppsService lazy discovery", () => { + let service: McpAppsService; + + beforeEach(() => { + service = makeService(); + }); + + it("discovers on first hasUiForTool when no session ran discovery", async () => { + service.setConfigResolver(async (name) => { + service.addServerConfigs([config(name)]); + }); + const client = connectClient(service); + + await expect( + service.hasUiForTool("mcp__posthog__loops-review"), + ).resolves.toBe(true); + expect(client.listTools).toHaveBeenCalledTimes(1); + }); + + it("emits DiscoveryComplete after a lazy discovery", async () => { + service.setServerConfigs([config("posthog")]); + connectClient(service); + const onComplete = vi.fn(); + service.on(McpAppsServiceEvent.DiscoveryComplete, onComplete); + + await service.hasUiForTool("mcp__posthog__loops-review"); + + expect(onComplete).toHaveBeenCalledWith({ + toolKeys: ["mcp__posthog__loops-review"], + }); + }); + + it.each([ + ["a non-MCP tool", "Bash"], + ["a malformed MCP key", "mcp__posthog"], + ])("returns false for %s without connecting", async (_label, toolKey) => { + const createConnection = vi.spyOn(internals(service), "createConnection"); + + await expect(service.hasUiForTool(toolKey)).resolves.toBe(false); + expect(createConnection).not.toHaveBeenCalled(); + }); + + it("answers UI-less tools from the discovered cache without re-listing", async () => { + service.setServerConfigs([config("posthog")]); + const client = connectClient(service); + + await service.hasUiForTool("mcp__posthog__loops-review"); + await expect( + service.hasUiForTool("mcp__posthog__loops-list"), + ).resolves.toBe(false); + expect(client.listTools).toHaveBeenCalledTimes(1); + }); + + it("rethrows discovery failures and backs off retries", async () => { + const resolver = vi.fn(async () => {}); + service.setConfigResolver(resolver); + + await expect( + service.hasUiForTool("mcp__posthog__loops-review"), + ).rejects.toThrow("No server config for: posthog"); + await expect( + service.hasUiForTool("mcp__posthog__loops-review"), + ).rejects.toThrow("UI tool discovery recently failed for: posthog"); + expect(resolver).toHaveBeenCalledTimes(1); + }); + + it("resolves lazily-discovered UI resources for a tool", async () => { + service.setServerConfigs([config("posthog")]); + connectClient(service); + + const resource = await service.getUiResourceForTool( + "mcp__posthog__loops-review", + ); + expect(resource?.uri).toBe(REVIEW_URI); + expect(resource?.html).toBe(""); + }); + + it("attaches discovered CSP metadata on direct URI fetches", async () => { + service.setConfigResolver(async (name) => { + service.addServerConfigs([config(name)]); + }); + connectClient(service); + + const resource = await service.getUiResourceByUri("posthog", REVIEW_URI); + expect(resource?.csp).toEqual(REVIEW_CSP); + }); +}); diff --git a/packages/core/src/mcp-apps/mcp-apps.ts b/packages/core/src/mcp-apps/mcp-apps.ts index 0e123e5111..830e5ca8a8 100644 --- a/packages/core/src/mcp-apps/mcp-apps.ts +++ b/packages/core/src/mcp-apps/mcp-apps.ts @@ -10,7 +10,7 @@ import { type IUrlLauncher, URL_LAUNCHER_SERVICE, } from "@posthog/platform/url-launcher"; -import { TypedEventEmitter } from "@posthog/shared"; +import { parseMcpToolName, TypedEventEmitter } from "@posthog/shared"; import { inject, injectable } from "inversify"; import { BUILTIN_POSTHOG_SERVER_NAME, @@ -52,6 +52,7 @@ function summarizeResult(result: unknown): Record { const UI_MIME_TYPE = "text/html;profile=mcp-app"; const MAX_HTML_SIZE = 5 * 1024 * 1024; // 5MB +const DISCOVERY_FAILURE_BACKOFF_MS = 60_000; interface ServerConnection { name: string; @@ -70,6 +71,9 @@ export class McpAppsService extends TypedEventEmitter { private pendingConnections = new Map>(); private pendingFetches = new Map>(); private resourceMetaCache = new Map(); + private discoveredServers = new Set(); + private pendingDiscoveries = new Map>(); + private discoveryFailedAt = new Map(); private readonly log: ScopedLogger; constructor( @@ -124,11 +128,21 @@ export class McpAppsService extends TypedEventEmitter { * emits DiscoveryComplete. */ async handleDiscovery(serverNames: string[]): Promise { - await Promise.allSettled( - serverNames - .filter((name) => this.serverConfigs.has(name)) - .map((name) => this.discoverServerUiTools(name)), + const names = serverNames.filter((name) => this.serverConfigs.has(name)); + const results = await Promise.allSettled( + names.map((name) => this.discoverServer(name)), ); + results.forEach((result, i) => { + if (result.status === "rejected") { + this.log.warn("Failed to discover UI tools for server", { + serverName: names[i], + error: + result.reason instanceof Error + ? result.reason.message + : String(result.reason), + }); + } + }); const toolKeys = [...this.toolAssociations.keys()]; this.log.info("Discovery complete", { @@ -145,69 +159,112 @@ export class McpAppsService extends TypedEventEmitter { /** * Connect to a single server and call listTools() to discover which * tools have _meta.ui fields. The connection is kept for later reuse - * (proxy calls, resource reads, lazy HTML fetches). + * (proxy calls, resource reads, lazy HTML fetches). Throws on connection + * or listTools failure — callers decide whether that is fatal. */ private async discoverServerUiTools(serverName: string): Promise { - try { - const conn = await this.getOrCreateConnection(serverName); + const conn = await this.getOrCreateConnection(serverName); + + const [toolsList, resourcesList] = await Promise.all([ + conn.client.listTools(), + conn.client.listResources().catch((err) => { + this.log.warn("listResources failed during discovery", { + serverName, + error: err instanceof Error ? err.message : String(err), + }); + return null; + }), + ]); + + this.log.info("discoverServerUiTools: listed tools", { + serverName, + toolNames: toolsList.tools.map((t) => t.name), + hasExecTool: + serverName === BUILTIN_POSTHOG_SERVER_NAME && + toolsList.tools.some((t) => t.name === EXEC_TOOL_NAME), + resourceUris: resourcesList?.resources.map((r) => r.uri), + }); + + for (const tool of toolsList.tools) { + if ( + serverName === BUILTIN_POSTHOG_SERVER_NAME && + tool.name === EXEC_TOOL_NAME + ) { + this.toolDefinitions.set(POSTHOG_EXEC_TOOL_KEY, tool); + } + + const uiMeta = (tool as McpToolUiMeta)._meta?.ui; + if (!uiMeta?.resourceUri) continue; - const [toolsList, resourcesList] = await Promise.all([ - conn.client.listTools(), - conn.client.listResources().catch((err) => { - this.log.warn("listResources failed during discovery", { - serverName, - error: err instanceof Error ? err.message : String(err), - }); - return null; - }), - ]); - - this.log.info("discoverServerUiTools: listed tools", { + const toolKey = `mcp__${serverName}__${tool.name}`; + this.toolAssociations.set(toolKey, { + toolKey, serverName, - toolNames: toolsList.tools.map((t) => t.name), - hasExecTool: - serverName === BUILTIN_POSTHOG_SERVER_NAME && - toolsList.tools.some((t) => t.name === EXEC_TOOL_NAME), - resourceUris: resourcesList?.resources.map((r) => r.uri), + toolName: tool.name, + resourceUri: uiMeta.resourceUri, + visibility: uiMeta.visibility, }); + this.toolDefinitions.set(toolKey, tool); + } - for (const tool of toolsList.tools) { - if ( - serverName === BUILTIN_POSTHOG_SERVER_NAME && - tool.name === EXEC_TOOL_NAME - ) { - this.toolDefinitions.set(POSTHOG_EXEC_TOOL_KEY, tool); + // Cache resource metadata (CSP, permissions) for use in fetchUiResource + if (resourcesList) { + for (const resource of resourcesList.resources) { + const meta = resource as McpResourceUiMeta; + if (meta._meta?.ui) { + this.resourceMetaCache.set(resource.uri, meta); } + } + } + } - const uiMeta = (tool as McpToolUiMeta)._meta?.ui; - if (!uiMeta?.resourceUri) continue; - - const toolKey = `mcp__${serverName}__${tool.name}`; - this.toolAssociations.set(toolKey, { - toolKey, - serverName, - toolName: tool.name, - resourceUri: uiMeta.resourceUri, - visibility: uiMeta.visibility, - }); - this.toolDefinitions.set(toolKey, tool); + /** + * Run discovery for one server, deduplicating concurrent attempts. Always + * lists fresh (no discovered-cache short-circuit) so session starts pick up + * server-side tool changes. + */ + private discoverServer(serverName: string): Promise { + const pending = this.pendingDiscoveries.get(serverName); + if (pending) return pending; + + const discovery = (async () => { + try { + await this.discoverServerUiTools(serverName); + this.discoveredServers.add(serverName); + this.discoveryFailedAt.delete(serverName); + } catch (err) { + this.discoveryFailedAt.set(serverName, Date.now()); + throw err; + } finally { + this.pendingDiscoveries.delete(serverName); } + })(); + this.pendingDiscoveries.set(serverName, discovery); + return discovery; + } - // Cache resource metadata (CSP, permissions) for use in fetchUiResource - if (resourcesList) { - for (const resource of resourcesList.resources) { - const meta = resource as McpResourceUiMeta; - if (meta._meta?.ui) { - this.resourceMetaCache.set(resource.uri, meta); - } - } + /** + * Lazily discover a server's UI tools on first use. Cloud runs never start a + * local agent session, so handleDiscovery never fires for them and the + * association map stays empty — the review-card path then fails silently. + * Failures rethrow (with a short backoff against hammering the config + * resolver) so callers' queries surface an error and retry instead of + * caching a permanent miss. + */ + private async ensureServerDiscovered(serverName: string): Promise { + if (this.discoveredServers.has(serverName)) return; + + if (!this.pendingDiscoveries.has(serverName)) { + const failedAt = this.discoveryFailedAt.get(serverName); + if (failedAt && Date.now() - failedAt < DISCOVERY_FAILURE_BACKOFF_MS) { + throw new Error(`UI tool discovery recently failed for: ${serverName}`); } - } catch (err) { - this.log.warn("Failed to discover UI tools for server", { - serverName, - error: err instanceof Error ? err.message : String(err), - }); } + + await this.discoverServer(serverName); + this.emit(McpAppsServiceEvent.DiscoveryComplete, { + toolKeys: [...this.toolAssociations.keys()], + } satisfies McpAppsDiscoveryCompleteEvent); } /** @@ -294,7 +351,13 @@ export class McpAppsService extends TypedEventEmitter { * Fetch the UI resource for a registration-discovered tool, by its tool key. */ async getUiResourceForTool(toolKey: string): Promise { - const association = this.toolAssociations.get(toolKey); + let association = this.toolAssociations.get(toolKey); + if (!association) { + const mcp = parseMcpToolName(toolKey); + if (!mcp) return null; + await this.ensureServerDiscovered(mcp.server); + association = this.toolAssociations.get(toolKey); + } if (!association) { this.log.debug("getUiResourceForTool: no association found", { toolKey }); return null; @@ -370,6 +433,11 @@ export class McpAppsService extends TypedEventEmitter { serverName: string, resourceUri: string, ): Promise { + // Best-effort warm of resourceMetaCache so CSP/permissions attach on paths + // where handleDiscovery never ran (cloud runs fetching by result URI). The + // read below decides success on its own. + await this.ensureServerDiscovered(serverName).catch(() => undefined); + let resourceResult: Awaited>; try { const conn = await this.getOrCreateConnection(serverName); @@ -432,7 +500,11 @@ export class McpAppsService extends TypedEventEmitter { return resource; } - hasUiForTool(toolKey: string): boolean { + async hasUiForTool(toolKey: string): Promise { + if (this.toolAssociations.has(toolKey)) return true; + const mcp = parseMcpToolName(toolKey); + if (!mcp) return false; + await this.ensureServerDiscovered(mcp.server); const has = this.toolAssociations.has(toolKey); this.log.debug("hasUiForTool", { toolKey, result: has }); return has; @@ -543,6 +615,9 @@ export class McpAppsService extends TypedEventEmitter { this.toolDefinitions.clear(); this.pendingConnections.clear(); this.pendingFetches.clear(); + this.discoveredServers.clear(); + this.pendingDiscoveries.clear(); + this.discoveryFailedAt.clear(); // Re-discover using stored server configs const serverNames = [...this.serverConfigs.keys()]; @@ -568,6 +643,8 @@ export class McpAppsService extends TypedEventEmitter { }); } this.connections.delete(serverName); + this.discoveredServers.delete(serverName); + this.discoveryFailedAt.delete(serverName); // Clean up associations and cached resources for this server const urisToEvict = new Set(); @@ -601,5 +678,8 @@ export class McpAppsService extends TypedEventEmitter { this.serverConfigs.clear(); this.pendingConnections.clear(); this.pendingFetches.clear(); + this.discoveredServers.clear(); + this.pendingDiscoveries.clear(); + this.discoveryFailedAt.clear(); } } From 68b361fa539e688be390bc1cb94ae44680975fd7 Mon Sep 17 00:00:00 2001 From: Charles Vien Date: Mon, 27 Jul 2026 12:34:20 -0700 Subject: [PATCH 2/2] address review findings on lazy discovery --- packages/core/src/mcp-apps/mcp-apps.test.ts | 110 +++++++++++++++++++- packages/core/src/mcp-apps/mcp-apps.ts | 83 +++++++++++---- 2 files changed, 170 insertions(+), 23 deletions(-) diff --git a/packages/core/src/mcp-apps/mcp-apps.test.ts b/packages/core/src/mcp-apps/mcp-apps.test.ts index 18beb59c27..70f3a6a44a 100644 --- a/packages/core/src/mcp-apps/mcp-apps.test.ts +++ b/packages/core/src/mcp-apps/mcp-apps.test.ts @@ -123,6 +123,7 @@ const REVIEW_CSP = { connectDomains: ["https://us.posthog.com"] }; function makeClient() { return { + close: vi.fn(async () => {}), listTools: vi.fn(async () => ({ tools: [ { @@ -175,11 +176,39 @@ describe("McpAppsService lazy discovery", () => { await service.hasUiForTool("mcp__posthog__loops-review"); - expect(onComplete).toHaveBeenCalledWith({ + expect(onComplete).toHaveBeenCalledExactlyOnceWith({ toolKeys: ["mcp__posthog__loops-review"], }); }); + it("dedupes concurrent lazy discoveries and emits once", async () => { + service.setServerConfigs([config("posthog")]); + const client = connectClient(service); + const onComplete = vi.fn(); + service.on(McpAppsServiceEvent.DiscoveryComplete, onComplete); + + const [reviewHasUi, listHasUi] = await Promise.all([ + service.hasUiForTool("mcp__posthog__loops-review"), + service.hasUiForTool("mcp__posthog__loops-list"), + ]); + + expect(reviewHasUi).toBe(true); + expect(listHasUi).toBe(false); + expect(client.listTools).toHaveBeenCalledTimes(1); + expect(onComplete).toHaveBeenCalledTimes(1); + }); + + it("does not emit DiscoveryComplete when lazy discovery fails", async () => { + service.setConfigResolver(vi.fn(async () => {})); + const onComplete = vi.fn(); + service.on(McpAppsServiceEvent.DiscoveryComplete, onComplete); + + await expect( + service.hasUiForTool("mcp__posthog__loops-review"), + ).rejects.toThrow(); + expect(onComplete).not.toHaveBeenCalled(); + }); + it.each([ ["a non-MCP tool", "Bash"], ["a malformed MCP key", "mcp__posthog"], @@ -190,6 +219,19 @@ describe("McpAppsService lazy discovery", () => { expect(createConnection).not.toHaveBeenCalled(); }); + it.each([ + ["a non-MCP tool", "Bash"], + ["a malformed MCP key", "mcp__posthog"], + ])( + "getUiResourceForTool returns null for %s without connecting", + async (_label, toolKey) => { + const createConnection = vi.spyOn(internals(service), "createConnection"); + + await expect(service.getUiResourceForTool(toolKey)).resolves.toBeNull(); + expect(createConnection).not.toHaveBeenCalled(); + }, + ); + it("answers UI-less tools from the discovered cache without re-listing", async () => { service.setServerConfigs([config("posthog")]); const client = connectClient(service); @@ -214,6 +256,57 @@ describe("McpAppsService lazy discovery", () => { expect(resolver).toHaveBeenCalledTimes(1); }); + it("retries discovery after the failure backoff expires", async () => { + vi.useFakeTimers(); + try { + const resolver = vi.fn(async () => {}); + service.setConfigResolver(resolver); + + await expect( + service.hasUiForTool("mcp__posthog__loops-review"), + ).rejects.toThrow("No server config for: posthog"); + vi.advanceTimersByTime(61_000); + await expect( + service.hasUiForTool("mcp__posthog__loops-review"), + ).rejects.toThrow("No server config for: posthog"); + expect(resolver).toHaveBeenCalledTimes(2); + } finally { + vi.useRealTimers(); + } + }); + + it("skips re-listing when handleDiscovery already discovered the server", async () => { + service.setServerConfigs([config("posthog")]); + const client = connectClient(service); + + await service.handleDiscovery(["posthog"]); + await expect( + service.hasUiForTool("mcp__posthog__loops-list"), + ).resolves.toBe(false); + expect(client.listTools).toHaveBeenCalledTimes(1); + }); + + it("re-lists on handleDiscovery even when already discovered", async () => { + service.setServerConfigs([config("posthog")]); + const client = connectClient(service); + + await service.hasUiForTool("mcp__posthog__loops-review"); + await service.handleDiscovery(["posthog"]); + expect(client.listTools).toHaveBeenCalledTimes(2); + }); + + it("re-discovers after disconnectServer clears server state", async () => { + service.setServerConfigs([config("posthog")]); + const client = connectClient(service); + + await service.hasUiForTool("mcp__posthog__loops-review"); + await service.disconnectServer("posthog"); + await expect( + service.hasUiForTool("mcp__posthog__loops-review"), + ).resolves.toBe(true); + expect(client.listTools).toHaveBeenCalledTimes(2); + }); + it("resolves lazily-discovered UI resources for a tool", async () => { service.setServerConfigs([config("posthog")]); connectClient(service); @@ -234,4 +327,19 @@ describe("McpAppsService lazy discovery", () => { const resource = await service.getUiResourceByUri("posthog", REVIEW_URI); expect(resource?.csp).toEqual(REVIEW_CSP); }); + + it("returns an uncached resource when the metadata warm-up fails", async () => { + service.setServerConfigs([config("posthog")]); + const client = makeClient(); + client.listTools.mockRejectedValue(new Error("listTools broken")); + connectClient(service, client); + + const first = await service.getUiResourceByUri("posthog", REVIEW_URI); + expect(first?.html).toBe(""); + expect(first?.csp).toBeUndefined(); + + const second = await service.getUiResourceByUri("posthog", REVIEW_URI); + expect(second?.html).toBe(""); + expect(client.readResource).toHaveBeenCalledTimes(2); + }); }); diff --git a/packages/core/src/mcp-apps/mcp-apps.ts b/packages/core/src/mcp-apps/mcp-apps.ts index 830e5ca8a8..0169efab12 100644 --- a/packages/core/src/mcp-apps/mcp-apps.ts +++ b/packages/core/src/mcp-apps/mcp-apps.ts @@ -254,7 +254,11 @@ export class McpAppsService extends TypedEventEmitter { private async ensureServerDiscovered(serverName: string): Promise { if (this.discoveredServers.has(serverName)) return; - if (!this.pendingDiscoveries.has(serverName)) { + // Only the caller that starts the discovery emits DiscoveryComplete — + // joiners would otherwise re-emit once per caller and stampede the + // renderer's query invalidations. + const startedHere = !this.pendingDiscoveries.has(serverName); + if (startedHere) { const failedAt = this.discoveryFailedAt.get(serverName); if (failedAt && Date.now() - failedAt < DISCOVERY_FAILURE_BACKOFF_MS) { throw new Error(`UI tool discovery recently failed for: ${serverName}`); @@ -262,9 +266,25 @@ export class McpAppsService extends TypedEventEmitter { } await this.discoverServer(serverName); - this.emit(McpAppsServiceEvent.DiscoveryComplete, { - toolKeys: [...this.toolAssociations.keys()], - } satisfies McpAppsDiscoveryCompleteEvent); + if (startedHere) { + this.emit(McpAppsServiceEvent.DiscoveryComplete, { + toolKeys: [...this.toolAssociations.keys()], + } satisfies McpAppsDiscoveryCompleteEvent); + } + } + + /** + * Look up a tool's UI association, lazily discovering its server on a miss. + */ + private async resolveAssociation( + toolKey: string, + ): Promise { + const existing = this.toolAssociations.get(toolKey); + if (existing) return existing; + const mcp = parseMcpToolName(toolKey); + if (!mcp) return undefined; + await this.ensureServerDiscovered(mcp.server); + return this.toolAssociations.get(toolKey); } /** @@ -351,13 +371,7 @@ export class McpAppsService extends TypedEventEmitter { * Fetch the UI resource for a registration-discovered tool, by its tool key. */ async getUiResourceForTool(toolKey: string): Promise { - let association = this.toolAssociations.get(toolKey); - if (!association) { - const mcp = parseMcpToolName(toolKey); - if (!mcp) return null; - await this.ensureServerDiscovered(mcp.server); - association = this.toolAssociations.get(toolKey); - } + const association = await this.resolveAssociation(toolKey); if (!association) { this.log.debug("getUiResourceForTool: no association found", { toolKey }); return null; @@ -436,7 +450,17 @@ export class McpAppsService extends TypedEventEmitter { // Best-effort warm of resourceMetaCache so CSP/permissions attach on paths // where handleDiscovery never ran (cloud runs fetching by result URI). The // read below decides success on its own. - await this.ensureServerDiscovered(serverName).catch(() => undefined); + const warmed = await this.ensureServerDiscovered(serverName).then( + () => true, + (err) => { + this.log.warn("UI resource metadata warm-up failed", { + serverName, + uri: resourceUri, + error: err instanceof Error ? err.message : String(err), + }); + return false; + }, + ); let resourceResult: Awaited>; try { @@ -489,23 +513,25 @@ export class McpAppsService extends TypedEventEmitter { serverName, }; - this.resourceCache.set(resourceUri, resource); - this.log.info("Lazily fetched and cached UI resource", { + // A failed warm-up with no known metadata may have produced a CSP-less + // copy; leave it uncached so a later fetch can attach the real CSP. + const cacheable = warmed || resourceMeta !== undefined; + if (cacheable) { + this.resourceCache.set(resourceUri, resource); + } + this.log.info("Lazily fetched UI resource", { serverName, uri: resourceUri, htmlLength: textContent.text.length, hasCsp: !!resource.csp, + cached: cacheable, }); return resource; } async hasUiForTool(toolKey: string): Promise { - if (this.toolAssociations.has(toolKey)) return true; - const mcp = parseMcpToolName(toolKey); - if (!mcp) return false; - await this.ensureServerDiscovered(mcp.server); - const has = this.toolAssociations.has(toolKey); + const has = !!(await this.resolveAssociation(toolKey)); this.log.debug("hasUiForTool", { toolKey, result: has }); return has; } @@ -631,6 +657,16 @@ export class McpAppsService extends TypedEventEmitter { } async disconnectServer(serverName: string): Promise { + // Let an in-flight lazy discovery land its connection first so it is + // closed here instead of lingering as a stray reconnect after teardown. + const pendingDiscovery = this.pendingDiscoveries.get(serverName); + if (pendingDiscovery) { + await pendingDiscovery.catch(() => undefined); + } + + this.discoveredServers.delete(serverName); + this.discoveryFailedAt.delete(serverName); + const conn = this.connections.get(serverName); if (!conn) return; @@ -643,8 +679,6 @@ export class McpAppsService extends TypedEventEmitter { }); } this.connections.delete(serverName); - this.discoveredServers.delete(serverName); - this.discoveryFailedAt.delete(serverName); // Clean up associations and cached resources for this server const urisToEvict = new Set(); @@ -667,7 +701,12 @@ export class McpAppsService extends TypedEventEmitter { } async cleanup(): Promise { - const serverNames = [...this.connections.keys()]; + // Include servers whose lazy discovery is still connecting — they have no + // entry in `connections` yet but will land one that must be closed. + const serverNames = new Set([ + ...this.connections.keys(), + ...this.pendingDiscoveries.keys(), + ]); for (const name of serverNames) { await this.disconnectServer(name); }