Skip to content
This repository was archived by the owner on Aug 6, 2026. It is now read-only.

Commit 315b4a2

Browse files
charlesviencharlesvien
committed
scope MCP UI resource caches by server name
1 parent 9e6e0ca commit 315b4a2

2 files changed

Lines changed: 223 additions & 14 deletions

File tree

‎packages/core/src/mcp-apps/mcp-apps.test.ts‎

Lines changed: 182 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -116,3 +116,185 @@ describe("McpAppsService config resolver", () => {
116116
expect(createConnection).toHaveBeenCalledTimes(2);
117117
});
118118
});
119+
120+
const UI_MIME_TYPE = "text/html;profile=mcp-app";
121+
122+
describe("McpAppsService resource cache isolation", () => {
123+
let service: McpAppsService;
124+
125+
beforeEach(() => {
126+
service = makeService();
127+
});
128+
129+
function stubPerServerReads(): void {
130+
vi.spyOn(internals(service), "getOrCreateConnection").mockImplementation(
131+
async (serverName: string) => ({
132+
name: serverName,
133+
client: {
134+
readResource: async () => ({
135+
contents: [
136+
{
137+
text: `<html>${serverName}</html>`,
138+
mimeType: UI_MIME_TYPE,
139+
},
140+
],
141+
}),
142+
},
143+
}),
144+
);
145+
}
146+
147+
it("does not let one server's resource satisfy another's fetch for the same URI", async () => {
148+
stubPerServerReads();
149+
const uri = "ui://posthog/survey-list.html";
150+
151+
const trusted = await service.getUiResourceByUri("posthog", uri);
152+
const malicious = await service.getUiResourceByUri("evil", uri);
153+
154+
expect(trusted?.html).toBe("<html>posthog</html>");
155+
expect(trusted?.serverName).toBe("posthog");
156+
expect(malicious?.html).toBe("<html>evil</html>");
157+
expect(malicious?.serverName).toBe("evil");
158+
});
159+
160+
it("serves a cache hit only to the server that populated it", async () => {
161+
const getConn = vi
162+
.spyOn(internals(service), "getOrCreateConnection")
163+
.mockImplementation(async (serverName: string) => ({
164+
name: serverName,
165+
client: {
166+
readResource: async () => ({
167+
contents: [
168+
{ text: `<html>${serverName}</html>`, mimeType: UI_MIME_TYPE },
169+
],
170+
}),
171+
},
172+
}));
173+
const uri = "ui://posthog/survey-list.html";
174+
175+
await service.getUiResourceByUri("posthog", uri);
176+
await service.getUiResourceByUri("posthog", uri);
177+
const other = await service.getUiResourceByUri("evil", uri);
178+
179+
expect(getConn).toHaveBeenCalledTimes(2);
180+
expect(other?.serverName).toBe("evil");
181+
});
182+
183+
it("does not share an in-flight fetch across servers for the same URI", async () => {
184+
const reads: string[] = [];
185+
let release: () => void = () => {};
186+
const gate = new Promise<void>((resolve) => {
187+
release = resolve;
188+
});
189+
vi.spyOn(internals(service), "getOrCreateConnection").mockImplementation(
190+
async (serverName: string) => ({
191+
name: serverName,
192+
client: {
193+
readResource: async () => {
194+
reads.push(serverName);
195+
await gate;
196+
return {
197+
contents: [
198+
{ text: `<html>${serverName}</html>`, mimeType: UI_MIME_TYPE },
199+
],
200+
};
201+
},
202+
},
203+
}),
204+
);
205+
const uri = "ui://shared/app.html";
206+
207+
const first = service.getUiResourceByUri("posthog", uri);
208+
const other = service.getUiResourceByUri("evil", uri);
209+
const joined = service.getUiResourceByUri("posthog", uri);
210+
release();
211+
const [r1, r2, r3] = await Promise.all([first, other, joined]);
212+
213+
expect(reads.filter((s) => s === "posthog")).toHaveLength(1);
214+
expect(reads.filter((s) => s === "evil")).toHaveLength(1);
215+
expect(r1?.serverName).toBe("posthog");
216+
expect(r3?.serverName).toBe("posthog");
217+
expect(r2?.serverName).toBe("evil");
218+
});
219+
});
220+
221+
describe("McpAppsService.proxyToolCall authorization", () => {
222+
let service: McpAppsService;
223+
const callTool = vi.fn(async () => ({ ok: true }));
224+
225+
function discoverTools(tools: unknown[]): Promise<void> {
226+
vi.spyOn(internals(service), "getOrCreateConnection").mockResolvedValue({
227+
name: "posthog",
228+
client: {
229+
listTools: async () => ({ tools }),
230+
listResources: async () => ({ resources: [] }),
231+
callTool,
232+
},
233+
});
234+
service.setServerConfigs([config("posthog")]);
235+
return service.handleDiscovery(["posthog"]);
236+
}
237+
238+
beforeEach(() => {
239+
service = makeService();
240+
callTool.mockClear();
241+
});
242+
243+
it("denies a tool that declares no UI metadata", async () => {
244+
await discoverTools([{ name: "exec" }]);
245+
246+
await expect(service.proxyToolCall("posthog", "exec")).rejects.toThrow(
247+
'Tool "exec" is not exposed to apps',
248+
);
249+
expect(callTool).not.toHaveBeenCalled();
250+
});
251+
252+
it("denies a tool that was never discovered", async () => {
253+
await discoverTools([]);
254+
255+
await expect(
256+
service.proxyToolCall("posthog", "delete_all"),
257+
).rejects.toThrow('Tool "delete_all" is not exposed to apps');
258+
expect(callTool).not.toHaveBeenCalled();
259+
});
260+
261+
it("allows a tool that opts in with ui.visibility app", async () => {
262+
await discoverTools([
263+
{ name: "search", _meta: { ui: { visibility: ["app"] } } },
264+
]);
265+
266+
await expect(service.proxyToolCall("posthog", "search")).resolves.toEqual({
267+
ok: true,
268+
});
269+
expect(callTool).toHaveBeenCalledTimes(1);
270+
});
271+
272+
it("allows a tool that carries a UI association", async () => {
273+
await discoverTools([
274+
{
275+
name: "surveys",
276+
_meta: { ui: { resourceUri: "ui://posthog/s.html" } },
277+
},
278+
]);
279+
280+
await expect(service.proxyToolCall("posthog", "surveys")).resolves.toEqual({
281+
ok: true,
282+
});
283+
});
284+
285+
it("still rejects a model-only tool", async () => {
286+
await discoverTools([
287+
{
288+
name: "surveys",
289+
_meta: {
290+
ui: { resourceUri: "ui://posthog/s.html", visibility: ["model"] },
291+
},
292+
},
293+
]);
294+
295+
await expect(service.proxyToolCall("posthog", "surveys")).rejects.toThrow(
296+
"not accessible to apps",
297+
);
298+
expect(callTool).not.toHaveBeenCalled();
299+
});
300+
});

‎packages/core/src/mcp-apps/mcp-apps.ts‎

Lines changed: 41 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,7 @@ export class McpAppsService extends TypedEventEmitter<McpAppsServiceEvents> {
6464
private connections = new Map<string, ServerConnection>();
6565
private resourceCache = new Map<string, McpUiResource>();
6666
private toolAssociations = new Map<string, McpToolUiAssociation>();
67+
private appVisibleTools = new Set<string>();
6768
private toolDefinitions = new Map<string, Tool>();
6869
private serverConfigs = new Map<string, McpServerConnectionConfig>();
6970
private configResolver?: (serverName: string) => Promise<void>;
@@ -83,6 +84,11 @@ export class McpAppsService extends TypedEventEmitter<McpAppsServiceEvents> {
8384
this.log = rootLogger.scope("mcp-apps-service");
8485
}
8586

87+
// Resource URIs are server-chosen and collide across servers.
88+
private cacheKey(serverName: string, resourceUri: string): string {
89+
return `${serverName}\u0000${resourceUri}`;
90+
}
91+
8692
/**
8793
* Store server configs for lazy connections later.
8894
* No connections are created at this point.
@@ -180,9 +186,14 @@ export class McpAppsService extends TypedEventEmitter<McpAppsServiceEvents> {
180186
}
181187

182188
const uiMeta = (tool as McpToolUiMeta)._meta?.ui;
189+
const toolKey = `mcp__${serverName}__${tool.name}`;
190+
191+
if (uiMeta?.visibility?.includes("app")) {
192+
this.appVisibleTools.add(toolKey);
193+
}
194+
183195
if (!uiMeta?.resourceUri) continue;
184196

185-
const toolKey = `mcp__${serverName}__${tool.name}`;
186197
this.toolAssociations.set(toolKey, {
187198
toolKey,
188199
serverName,
@@ -198,7 +209,10 @@ export class McpAppsService extends TypedEventEmitter<McpAppsServiceEvents> {
198209
for (const resource of resourcesList.resources) {
199210
const meta = resource as McpResourceUiMeta;
200211
if (meta._meta?.ui) {
201-
this.resourceMetaCache.set(resource.uri, meta);
212+
this.resourceMetaCache.set(
213+
this.cacheKey(serverName, resource.uri),
214+
meta,
215+
);
202216
}
203217
}
204218
}
@@ -335,7 +349,8 @@ export class McpAppsService extends TypedEventEmitter<McpAppsServiceEvents> {
335349
serverName: string,
336350
resourceUri: string,
337351
): Promise<McpUiResource | null> {
338-
const cached = this.resourceCache.get(resourceUri);
352+
const key = this.cacheKey(serverName, resourceUri);
353+
const cached = this.resourceCache.get(key);
339354
if (cached) {
340355
this.log.debug("fetchUiResourceByUri: cache hit", {
341356
serverName,
@@ -344,7 +359,7 @@ export class McpAppsService extends TypedEventEmitter<McpAppsServiceEvents> {
344359
return cached;
345360
}
346361

347-
const pendingFetch = this.pendingFetches.get(resourceUri);
362+
const pendingFetch = this.pendingFetches.get(key);
348363
if (pendingFetch) {
349364
this.log.debug("fetchUiResourceByUri: joining pending fetch", {
350365
serverName,
@@ -358,11 +373,11 @@ export class McpAppsService extends TypedEventEmitter<McpAppsServiceEvents> {
358373
resourceUri,
359374
});
360375
const fetchPromise = this.doFetchUiResource(serverName, resourceUri);
361-
this.pendingFetches.set(resourceUri, fetchPromise);
376+
this.pendingFetches.set(key, fetchPromise);
362377
try {
363378
return await fetchPromise;
364379
} finally {
365-
this.pendingFetches.delete(resourceUri);
380+
this.pendingFetches.delete(key);
366381
}
367382
}
368383

@@ -409,7 +424,9 @@ export class McpAppsService extends TypedEventEmitter<McpAppsServiceEvents> {
409424
return null;
410425
}
411426

412-
const resourceMeta = this.resourceMetaCache.get(resourceUri);
427+
const resourceMeta = this.resourceMetaCache.get(
428+
this.cacheKey(serverName, resourceUri),
429+
);
413430

414431
const resource: McpUiResource = {
415432
uri: resourceUri,
@@ -421,7 +438,7 @@ export class McpAppsService extends TypedEventEmitter<McpAppsServiceEvents> {
421438
serverName,
422439
};
423440

424-
this.resourceCache.set(resourceUri, resource);
441+
this.resourceCache.set(this.cacheKey(serverName, resourceUri), resource);
425442
this.log.info("Lazily fetched and cached UI resource", {
426443
serverName,
427444
uri: resourceUri,
@@ -456,6 +473,10 @@ export class McpAppsService extends TypedEventEmitter<McpAppsServiceEvents> {
456473
);
457474
}
458475

476+
if (!association && !this.appVisibleTools.has(toolKey)) {
477+
throw new Error(`Tool "${toolName}" is not exposed to apps`);
478+
}
479+
459480
const conn = await this.getOrCreateConnection(serverName);
460481
const result = await conn.client.callTool({
461482
name: toolName,
@@ -540,6 +561,7 @@ export class McpAppsService extends TypedEventEmitter<McpAppsServiceEvents> {
540561
this.resourceCache.clear();
541562
this.resourceMetaCache.clear();
542563
this.toolAssociations.clear();
564+
this.appVisibleTools.clear();
543565
this.toolDefinitions.clear();
544566
this.pendingConnections.clear();
545567
this.pendingFetches.clear();
@@ -578,13 +600,17 @@ export class McpAppsService extends TypedEventEmitter<McpAppsServiceEvents> {
578600
}
579601
}
580602

581-
// Only evict cached resources not referenced by remaining associations
582-
const stillReferenced = new Set(
583-
[...this.toolAssociations.values()].map((a) => a.resourceUri),
584-
);
585603
for (const uri of urisToEvict) {
586-
if (!stillReferenced.has(uri)) {
587-
this.resourceCache.delete(uri);
604+
const key = this.cacheKey(serverName, uri);
605+
this.resourceCache.delete(key);
606+
this.resourceMetaCache.delete(key);
607+
this.pendingFetches.delete(key);
608+
}
609+
610+
const toolKeyPrefix = `mcp__${serverName}__`;
611+
for (const toolKey of this.appVisibleTools) {
612+
if (toolKey.startsWith(toolKeyPrefix)) {
613+
this.appVisibleTools.delete(toolKey);
588614
}
589615
}
590616
}
@@ -597,6 +623,7 @@ export class McpAppsService extends TypedEventEmitter<McpAppsServiceEvents> {
597623
this.resourceCache.clear();
598624
this.resourceMetaCache.clear();
599625
this.toolAssociations.clear();
626+
this.appVisibleTools.clear();
600627
this.toolDefinitions.clear();
601628
this.serverConfigs.clear();
602629
this.pendingConnections.clear();

0 commit comments

Comments
 (0)