Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 44 additions & 0 deletions src/everything/__tests__/resources.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -221,6 +221,50 @@ describe('Session Resources', () => {
);
});

it('should not evict a resource with the same URI from another server', () => {
const firstRegistration = { remove: vi.fn() };
const secondRegistration = { remove: vi.fn() };
const firstServer = {
registerResource: vi.fn(() => firstRegistration),
} as unknown as McpServer;
const secondServer = {
registerResource: vi.fn(() => secondRegistration),
} as unknown as McpServer;
const resource = {
uri: 'demo://resource/session/shared-file',
name: 'shared-file',
mimeType: 'text/plain',
};

registerSessionResource(firstServer, resource, 'text', 'first session');
registerSessionResource(secondServer, resource, 'text', 'second session');

expect(firstRegistration.remove).not.toHaveBeenCalled();
expect(secondRegistration.remove).not.toHaveBeenCalled();
});

it('should replace a resource with the same URI on the same server', () => {
const firstRegistration = { remove: vi.fn() };
const secondRegistration = { remove: vi.fn() };
const server = {
registerResource: vi
.fn()
.mockReturnValueOnce(firstRegistration)
.mockReturnValueOnce(secondRegistration),
} as unknown as McpServer;
const resource = {
uri: 'demo://resource/session/replaced-file',
name: 'replaced-file',
mimeType: 'text/plain',
};

registerSessionResource(server, resource, 'text', 'first payload');
registerSessionResource(server, resource, 'text', 'second payload');

expect(firstRegistration.remove).toHaveBeenCalledOnce();
expect(secondRegistration.remove).not.toHaveBeenCalled();
});

it('should register blob resource correctly', () => {
const mockServer = {
registerResource: vi.fn(),
Expand Down
18 changes: 13 additions & 5 deletions src/everything/resources/session.ts
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,10 @@ import { Resource, ResourceLink } from "@modelcontextprotocol/sdk/types.js";
* This prevents "Resource already registered" errors when a tool creates a resource
* with the same URI multiple times during a session.
*/
const registeredResources = new Map<string, RegisteredResource>();
const registeredResources = new WeakMap<
McpServer,
Map<string, RegisteredResource>
>();

/**
* Generates a session-scoped resource URI string based on the provided resource name.
Expand Down Expand Up @@ -54,11 +57,16 @@ export const registerSessionResource = (
blob: payload,
};

// Check if a resource with this URI is already registered and remove it
const existingResource = registeredResources.get(uri);
const serverResources =
registeredResources.get(server) ?? new Map<string, RegisteredResource>();
registeredResources.set(server, serverResources);

// Check if a resource with this URI is already registered on this server and remove it
const resourceKey = uri.toString();
const existingResource = serverResources.get(resourceKey);
if (existingResource) {
existingResource.remove();
registeredResources.delete(uri);
serverResources.delete(resourceKey);
}

// Register file resource
Expand All @@ -74,7 +82,7 @@ export const registerSessionResource = (
);

// Track the registered resource for potential future removal
registeredResources.set(uri, registeredResource);
serverResources.set(resourceKey, registeredResource);

return { type: "resource_link", ...resource };
};