diff --git a/src/everything/__tests__/resources.test.ts b/src/everything/__tests__/resources.test.ts index a22b904175..edb57387e2 100644 --- a/src/everything/__tests__/resources.test.ts +++ b/src/everything/__tests__/resources.test.ts @@ -262,6 +262,59 @@ describe('Session Resources', () => { expect(handlerResult.contents[0].text).toBe('Test content here'); expect(handlerResult.contents[0].mimeType).toBe('text/plain'); }); + + it('should not remove a resource another session registered under the same URI', () => { + const removedFromA: string[] = []; + const removedFromB: string[] = []; + + const sessionA = { + registerResource: vi.fn((_name, uri) => { + return { remove: () => removedFromA.push(String(uri)) }; + }), + } as unknown as McpServer; + + const sessionB = { + registerResource: vi.fn((_name, uri) => { + return { remove: () => removedFromB.push(String(uri)) }; + }), + } as unknown as McpServer; + + const resource = { + uri: 'demo://resource/session/shared.gz', + name: 'shared.gz', + mimeType: 'application/gzip', + }; + + registerSessionResource(sessionA, resource, 'blob', 'from-a'); + registerSessionResource(sessionB, resource, 'blob', 'from-b'); + + // Each session must keep serving its own resource. + expect(removedFromA).toEqual([]); + expect(removedFromB).toEqual([]); + }); + + it('should still remove its own previous registration for the same URI', () => { + const removed: string[] = []; + + const server = { + registerResource: vi.fn((_name, uri) => { + return { remove: () => removed.push(String(uri)) }; + }), + } as unknown as McpServer; + + const resource = { + uri: 'demo://resource/session/retry.gz', + name: 'retry.gz', + mimeType: 'application/gzip', + }; + + // A tool called twice with the same output name must not fail with + // "Resource already registered"; the earlier registration is replaced. + registerSessionResource(server, resource, 'blob', 'first'); + registerSessionResource(server, resource, 'blob', 'second'); + + expect(removed).toEqual(['demo://resource/session/retry.gz']); + }); }); }); diff --git a/src/everything/resources/session.ts b/src/everything/resources/session.ts index 10e0db33c1..5059a900dc 100644 --- a/src/everything/resources/session.ts +++ b/src/everything/resources/session.ts @@ -2,11 +2,30 @@ import { McpServer, RegisteredResource } from "@modelcontextprotocol/sdk/server/ import { Resource, ResourceLink } from "@modelcontextprotocol/sdk/types.js"; /** - * Tracks registered session resources by URI to allow updating/removing on re-registration. - * This prevents "Resource already registered" errors when a tool creates a resource - * with the same URI multiple times during a session. + * Tracks registered session resources per server, by URI, to allow updating/removing on + * re-registration. This prevents "Resource already registered" errors when a tool creates a + * resource with the same URI multiple times during a session. + * + * The registry is scoped to the `McpServer` that owns the resource rather than being module + * level: every session gets its own server (see `server/index.ts`), so a module-level map lets + * one session's re-registration remove the resource another session registered under the same + * URI. A `WeakMap` also lets a closed session's registry be collected. + */ +const registeredResources = new WeakMap>(); + +/** + * Returns the resource registry owned by `server`, creating it on first use. */ -const registeredResources = new Map(); +const sessionResourcesFor = (server: McpServer): Map => { + let resources = registeredResources.get(server); + + if (!resources) { + resources = new Map(); + registeredResources.set(server, resources); + } + + return resources; +}; /** * Generates a session-scoped resource URI string based on the provided resource name. @@ -54,11 +73,12 @@ export const registerSessionResource = ( blob: payload, }; - // Check if a resource with this URI is already registered and remove it - const existingResource = registeredResources.get(uri); + // Check if this server already registered a resource with this URI and remove that one + const sessionResources = sessionResourcesFor(server); + const existingResource = sessionResources.get(uri); if (existingResource) { existingResource.remove(); - registeredResources.delete(uri); + sessionResources.delete(uri); } // Register file resource @@ -74,7 +94,7 @@ export const registerSessionResource = ( ); // Track the registered resource for potential future removal - registeredResources.set(uri, registeredResource); + sessionResources.set(uri, registeredResource); return { type: "resource_link", ...resource }; };