feat: Add support for MCP Resources (#13178)

Co-authored-by: Jack Wotherspoon <jackwoth@google.com>
This commit is contained in:
Alex Gavrilescu
2025-12-09 03:43:12 +01:00
committed by GitHub
parent 720b31cb8b
commit 560550f5df
20 changed files with 1146 additions and 80 deletions

View File

@@ -16,8 +16,8 @@ An MCP server enables the Gemini CLI to:
through standardized schema definitions.
- **Execute tools:** Call specific tools with defined arguments and receive
structured responses.
- **Access resources:** Read data from specific resources (though the Gemini CLI
primarily focuses on tool execution).
- **Access resources:** Read data from specific resources that the server
exposes (files, API payloads, reports, etc.).
With an MCP server, you can extend the Gemini CLI's capabilities to perform
actions beyond its built-in features, such as interacting with databases, APIs,
@@ -40,6 +40,7 @@ The discovery process is orchestrated by `discoverMcpTools()`, which:
4. **Sanitizes and validates** tool schemas for compatibility with the Gemini
API
5. **Registers tools** in the global tool registry with conflict resolution
6. **Fetches and registers resources** if the server exposes any
### Execution layer (`mcp-tool.ts`)
@@ -59,6 +60,32 @@ The Gemini CLI supports three MCP transport types:
- **SSE Transport:** Connects to Server-Sent Events endpoints
- **Streamable HTTP Transport:** Uses HTTP streaming for communication
## Working with MCP resources
Some MCP servers expose contextual “resources” in addition to the tools and
prompts. Gemini CLI discovers these automatically and gives you the possibility
to reference them in the chat.
### Discovery and listing
- When discovery runs, the CLI fetches each servers `resources/list` results.
- The `/mcp` command displays a Resources section alongside Tools and Prompts
for every connected server.
This returns a concise, plain-text list of URIs plus metadata.
### Referencing resources in a conversation
You can use the same `@` syntax already known for referencing local files:
```
@server://resource/path
```
Resource URIs appear in the completion menu together with filesystem paths. When
you submit the message, the CLI calls `resources/read` and injects the content
in the conversation.
## How to set up your MCP server
The Gemini CLI uses the `mcpServers` configuration in your `settings.json` file

View File

@@ -63,6 +63,7 @@ describe('mcpCommand', () => {
getPromptRegistry: ReturnType<typeof vi.fn>;
getGeminiClient: ReturnType<typeof vi.fn>;
getMcpClientManager: ReturnType<typeof vi.fn>;
getResourceRegistry: ReturnType<typeof vi.fn>;
};
beforeEach(() => {
@@ -93,6 +94,9 @@ describe('mcpCommand', () => {
getBlockedMcpServers: vi.fn(),
getMcpServers: vi.fn(),
})),
getResourceRegistry: vi.fn().mockReturnValue({
getAllResources: vi.fn().mockReturnValue([]),
}),
};
mockContext = createMockCommandContext({
@@ -141,6 +145,10 @@ describe('mcpCommand', () => {
};
mockConfig.getMcpServers = vi.fn().mockReturnValue(mockMcpServers);
mockConfig.getMcpClientManager = vi.fn().mockReturnValue({
getMcpServers: vi.fn().mockReturnValue(mockMcpServers),
getBlockedMcpServers: vi.fn().mockReturnValue([]),
});
});
it('should display configured MCP servers with status indicators and their tools', async () => {
@@ -169,6 +177,30 @@ describe('mcpCommand', () => {
getAllTools: vi.fn().mockReturnValue(allTools),
});
const resourcesByServer: Record<
string,
Array<{ name: string; uri: string }>
> = {
server1: [
{
name: 'Server1 Resource',
uri: 'file:///server1/resource1.txt',
},
],
server2: [],
server3: [],
};
mockConfig.getResourceRegistry = vi.fn().mockReturnValue({
getAllResources: vi.fn().mockReturnValue(
Object.entries(resourcesByServer).flatMap(([serverName, resources]) =>
resources.map((entry) => ({
serverName,
...entry,
})),
),
),
});
await mcpCommand.action!(mockContext, '');
expect(mockContext.ui.addItem).toHaveBeenCalledWith(
@@ -180,6 +212,12 @@ describe('mcpCommand', () => {
description: tool.description,
schema: tool.schema,
})),
resources: expect.arrayContaining([
expect.objectContaining({
serverName: 'server1',
uri: 'file:///server1/resource1.txt',
}),
]),
}),
expect.any(Number),
);

View File

@@ -12,6 +12,7 @@ import type {
import { CommandKind } from './types.js';
import type {
DiscoveredMCPPrompt,
DiscoveredMCPResource,
MessageActionReturn,
} from '@google/gemini-cli-core';
import {
@@ -230,6 +231,13 @@ const listAction = async (
serverNames.includes(prompt.serverName as string),
) as DiscoveredMCPPrompt[];
const resourceRegistry = config.getResourceRegistry();
const mcpResources = resourceRegistry
.getAllResources()
.filter((entry) =>
serverNames.includes(entry.serverName),
) as DiscoveredMCPResource[];
const authStatus: HistoryItemMcpStatus['authStatus'] = {};
const tokenStorage = new MCPOAuthTokenStorage();
for (const serverName of serverNames) {
@@ -265,6 +273,13 @@ const listAction = async (
name: prompt.name,
description: prompt.description,
})),
resources: mcpResources.map((resource) => ({
serverName: resource.serverName,
name: resource.name,
uri: resource.uri,
mimeType: resource.mimeType,
description: resource.description,
})),
authStatus,
blockedServers: blockedMcpServers,
discoveryInProgress,

View File

@@ -36,6 +36,7 @@ describe('McpStatus', () => {
},
],
prompts: [],
resources: [],
blockedServers: [],
serverStatus: () => MCPServerStatus.CONNECTED,
authStatus: {},
@@ -147,6 +148,24 @@ describe('McpStatus', () => {
unmount();
});
it('renders correctly with resources', () => {
const { lastFrame, unmount } = render(
<McpStatus
{...baseProps}
resources={[
{
serverName: 'server-1',
name: 'resource-1',
uri: 'file:///tmp/resource-1.txt',
description: 'A test resource',
},
]}
/>,
);
expect(lastFrame()).toMatchSnapshot();
unmount();
});
it('renders correctly with a blocked server', () => {
const { lastFrame, unmount } = render(
<McpStatus

View File

@@ -12,6 +12,7 @@ import { theme } from '../../semantic-colors.js';
import type {
HistoryItemMcpStatus,
JsonMcpPrompt,
JsonMcpResource,
JsonMcpTool,
} from '../../types.js';
@@ -19,6 +20,7 @@ interface McpStatusProps {
servers: Record<string, MCPServerConfig>;
tools: JsonMcpTool[];
prompts: JsonMcpPrompt[];
resources: JsonMcpResource[];
blockedServers: Array<{ name: string; extensionName: string }>;
serverStatus: (serverName: string) => MCPServerStatus;
authStatus: HistoryItemMcpStatus['authStatus'];
@@ -32,6 +34,7 @@ export const McpStatus: React.FC<McpStatusProps> = ({
servers,
tools,
prompts,
resources,
blockedServers,
serverStatus,
authStatus,
@@ -83,9 +86,14 @@ export const McpStatus: React.FC<McpStatusProps> = ({
const serverPrompts = prompts.filter(
(prompt) => prompt.serverName === serverName,
);
const serverResources = resources.filter(
(resource) => resource.serverName === serverName,
);
const originalStatus = serverStatus(serverName);
const hasCachedItems =
serverTools.length > 0 || serverPrompts.length > 0;
serverTools.length > 0 ||
serverPrompts.length > 0 ||
serverResources.length > 0;
const status =
originalStatus === MCPServerStatus.DISCONNECTED && hasCachedItems
? MCPServerStatus.CONNECTED
@@ -121,6 +129,7 @@ export const McpStatus: React.FC<McpStatusProps> = ({
const toolCount = serverTools.length;
const promptCount = serverPrompts.length;
const resourceCount = serverResources.length;
const parts = [];
if (toolCount > 0) {
parts.push(`${toolCount} ${toolCount === 1 ? 'tool' : 'tools'}`);
@@ -130,6 +139,11 @@ export const McpStatus: React.FC<McpStatusProps> = ({
`${promptCount} ${promptCount === 1 ? 'prompt' : 'prompts'}`,
);
}
if (resourceCount > 0) {
parts.push(
`${resourceCount} ${resourceCount === 1 ? 'resource' : 'resources'}`,
);
}
const serverAuthStatus = authStatus[serverName];
let authStatusNode: React.ReactNode = null;
@@ -233,6 +247,34 @@ export const McpStatus: React.FC<McpStatusProps> = ({
))}
</Box>
)}
{serverResources.length > 0 && (
<Box flexDirection="column" marginLeft={2}>
<Text color={theme.text.primary}>Resources:</Text>
{serverResources.map((resource, index) => {
const label = resource.name || resource.uri || 'resource';
return (
<Box
key={`${resource.serverName}-resource-${index}`}
flexDirection="column"
>
<Text>
- <Text color={theme.text.primary}>{label}</Text>
{resource.uri ? ` (${resource.uri})` : ''}
{resource.mimeType ? ` [${resource.mimeType}]` : ''}
</Text>
{showDescriptions && resource.description && (
<Box marginLeft={2}>
<Text color={theme.text.secondary}>
{resource.description.trim()}
</Text>
</Box>
)}
</Box>
);
})}
</Box>
)}
</Box>
);
})}

View File

@@ -116,6 +116,20 @@ A test server
"
`;
exports[`McpStatus > renders correctly with resources 1`] = `
"Configured MCP servers:
🟢 server-1 - Ready (1 tool, 1 resource)
A test server
Tools:
- tool-1
A test tool
Resources:
- resource-1 (file:///tmp/resource-1.txt)
A test resource
"
`;
exports[`McpStatus > renders correctly with schema enabled 1`] = `
"Configured MCP servers:

View File

@@ -7,7 +7,7 @@
import type { Mock } from 'vitest';
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { handleAtCommand } from './atCommandProcessor.js';
import type { Config } from '@google/gemini-cli-core';
import type { Config, DiscoveredMCPResource } from '@google/gemini-cli-core';
import {
FileDiscoveryService,
GlobTool,
@@ -86,6 +86,13 @@ describe('handleAtCommand', () => {
}),
getUsageStatisticsEnabled: () => false,
getEnableExtensionReloading: () => false,
getResourceRegistry: () => ({
findResourceByUri: () => undefined,
getAllResources: () => [],
}),
getMcpClientManager: () => ({
getClient: () => undefined,
}),
} as unknown as Config;
const registry = new ToolRegistry(mockConfig);
@@ -1241,4 +1248,98 @@ describe('handleAtCommand', () => {
);
expect(userTurnCalls).toHaveLength(0);
});
describe('MCP resource attachments', () => {
it('attaches MCP resource content when @serverName:uri matches registry', async () => {
const serverName = 'server-1';
const resourceUri = 'resource://server-1/logs';
const prefixedUri = `${serverName}:${resourceUri}`;
const resource = {
serverName,
uri: resourceUri,
name: 'logs',
discoveredAt: Date.now(),
} as DiscoveredMCPResource;
vi.spyOn(mockConfig, 'getResourceRegistry').mockReturnValue({
findResourceByUri: (identifier: string) =>
identifier === prefixedUri ? resource : undefined,
getAllResources: () => [],
} as never);
const readResource = vi.fn().mockResolvedValue({
contents: [{ text: 'mcp resource body' }],
});
vi.spyOn(mockConfig, 'getMcpClientManager').mockReturnValue({
getClient: () => ({ readResource }),
} as never);
const result = await handleAtCommand({
query: `@${prefixedUri}`,
config: mockConfig,
addItem: mockAddItem,
onDebugMessage: mockOnDebugMessage,
messageId: 42,
signal: abortController.signal,
});
expect(readResource).toHaveBeenCalledWith(resourceUri);
const processedParts = Array.isArray(result.processedQuery)
? result.processedQuery
: [];
const containsResourceText = processedParts.some((part) => {
const text = typeof part === 'string' ? part : part?.text;
return typeof text === 'string' && text.includes('mcp resource body');
});
expect(containsResourceText).toBe(true);
expect(mockAddItem).toHaveBeenCalledWith(
expect.objectContaining({ type: 'tool_group' }),
expect.any(Number),
);
});
it('returns an error if MCP client is unavailable', async () => {
const serverName = 'server-1';
const resourceUri = 'resource://server-1/logs';
const prefixedUri = `${serverName}:${resourceUri}`;
vi.spyOn(mockConfig, 'getResourceRegistry').mockReturnValue({
findResourceByUri: (identifier: string) =>
identifier === prefixedUri
? ({
serverName,
uri: resourceUri,
discoveredAt: Date.now(),
} as DiscoveredMCPResource)
: undefined,
getAllResources: () => [],
} as never);
vi.spyOn(mockConfig, 'getMcpClientManager').mockReturnValue({
getClient: () => undefined,
} as never);
const result = await handleAtCommand({
query: `@${prefixedUri}`,
config: mockConfig,
addItem: mockAddItem,
onDebugMessage: mockOnDebugMessage,
messageId: 42,
signal: abortController.signal,
});
expect(result.shouldProceed).toBe(false);
expect(mockAddItem).toHaveBeenCalledWith(
expect.objectContaining({
type: 'tool_group',
tools: expect.arrayContaining([
expect.objectContaining({
resultDisplay: expect.stringContaining(
"MCP client for server 'server-1' is not available or not connected.",
),
}),
]),
}),
expect.any(Number),
);
});
});
});

View File

@@ -7,7 +7,11 @@
import * as fs from 'node:fs/promises';
import * as path from 'node:path';
import type { PartListUnion, PartUnion } from '@google/genai';
import type { AnyToolInvocation, Config } from '@google/gemini-cli-core';
import type {
AnyToolInvocation,
Config,
DiscoveredMCPResource,
} from '@google/gemini-cli-core';
import {
debugLogger,
getErrorMessage,
@@ -15,6 +19,7 @@ import {
unescapePath,
ReadManyFilesTool,
} from '@google/gemini-cli-core';
import { Buffer } from 'node:buffer';
import type { HistoryItem, IndividualToolCallDisplay } from '../types.js';
import { ToolCallStatus } from '../types.js';
import type { UseHistoryManagerReturn } from './useHistoryManager.js';
@@ -113,13 +118,14 @@ function parseAllAtCommands(query: string): AtCommandPart[] {
}
/**
* Processes user input potentially containing one or more '@<path>' commands.
* If found, it attempts to read the specified files/directories using the
* 'read_many_files' tool. The user query is modified to include resolved paths,
* and the content of the files is appended in a structured block.
* Processes user input containing one or more '@<path>' commands.
* - Workspace paths are read via the 'read_many_files' tool.
* - MCP resource URIs are read via each server's `resources/read`.
* The user query is updated with inline content blocks so the LLM receives the
* referenced context directly.
*
* @returns An object indicating whether the main hook should proceed with an
* LLM call and the processed query parts (including file content).
* LLM call and the processed query parts (including file/resource content).
*/
export async function handleAtCommand({
query,
@@ -129,6 +135,9 @@ export async function handleAtCommand({
messageId: userMessageTimestamp,
signal,
}: HandleAtCommandParams): Promise<HandleAtCommandResult> {
const resourceRegistry = config.getResourceRegistry();
const mcpClientManager = config.getMcpClientManager();
const commandParts = parseAllAtCommands(query);
const atPathCommandParts = commandParts.filter(
(part) => part.type === 'atPath',
@@ -144,8 +153,9 @@ export async function handleAtCommand({
const respectFileIgnore = config.getFileFilteringOptions();
const pathSpecsToRead: string[] = [];
const resourceAttachments: DiscoveredMCPResource[] = [];
const atPathToResolvedSpecMap = new Map<string, string>();
const contentLabelsForDisplay: string[] = [];
const fileLabelsForDisplay: string[] = [];
const absoluteToRelativePathMap = new Map<string, string>();
const ignoredByReason: Record<string, string[]> = {
git: [],
@@ -191,7 +201,13 @@ export async function handleAtCommand({
return { processedQuery: null, shouldProceed: false };
}
// Check if path should be ignored based on filtering options
// Check if this is an MCP resource reference (serverName:uri format)
const resourceMatch = resourceRegistry.findResourceByUri(pathName);
if (resourceMatch) {
resourceAttachments.push(resourceMatch);
atPathToResolvedSpecMap.set(originalAtPath, pathName);
continue;
}
const workspaceContext = config.getWorkspaceContext();
if (!workspaceContext.isPathWithinWorkspace(pathName)) {
@@ -324,7 +340,7 @@ export async function handleAtCommand({
pathSpecsToRead.push(currentPathSpec);
atPathToResolvedSpecMap.set(originalAtPath, currentPathSpec);
const displayPath = path.isAbsolute(pathName) ? relativePath : pathName;
contentLabelsForDisplay.push(displayPath);
fileLabelsForDisplay.push(displayPath);
break;
}
}
@@ -397,7 +413,7 @@ export async function handleAtCommand({
}
// Fallback for lone "@" or completely invalid @-commands resulting in empty initialQueryText
if (pathSpecsToRead.length === 0) {
if (pathSpecsToRead.length === 0 && resourceAttachments.length === 0) {
onDebugMessage('No valid file paths found in @ commands to read.');
if (initialQueryText === '@' && query.trim() === '@') {
// If the only thing was a lone @, pass original query (which might have spaces)
@@ -413,7 +429,86 @@ export async function handleAtCommand({
};
}
const processedQueryParts: PartUnion[] = [{ text: initialQueryText }];
const processedQueryParts: PartListUnion = [{ text: initialQueryText }];
const resourcePromises = resourceAttachments.map(async (resource) => {
const uri = resource.uri!;
const client = mcpClientManager?.getClient(resource.serverName);
try {
if (!client) {
throw new Error(
`MCP client for server '${resource.serverName}' is not available or not connected.`,
);
}
const response = await client.readResource(uri);
const parts = convertResourceContentsToParts(response);
return {
success: true,
parts,
uri,
display: {
callId: `mcp-resource-${resource.serverName}-${uri}`,
name: `resources/read (${resource.serverName})`,
description: uri,
status: ToolCallStatus.Success,
resultDisplay: `Successfully read resource ${uri}`,
confirmationDetails: undefined,
} as IndividualToolCallDisplay,
};
} catch (error) {
return {
success: false,
parts: [],
uri,
display: {
callId: `mcp-resource-${resource.serverName}-${uri}`,
name: `resources/read (${resource.serverName})`,
description: uri,
status: ToolCallStatus.Error,
resultDisplay: `Error reading resource ${uri}: ${getErrorMessage(error)}`,
confirmationDetails: undefined,
} as IndividualToolCallDisplay,
};
}
});
const resourceResults = await Promise.all(resourcePromises);
const resourceReadDisplays: IndividualToolCallDisplay[] = [];
let resourceErrorOccurred = false;
for (const result of resourceResults) {
resourceReadDisplays.push(result.display);
if (result.success) {
processedQueryParts.push({ text: `\nContent from @${result.uri}:\n` });
processedQueryParts.push(...result.parts);
} else {
resourceErrorOccurred = true;
}
}
if (resourceErrorOccurred) {
addItem(
{ type: 'tool_group', tools: resourceReadDisplays } as Omit<
HistoryItem,
'id'
>,
userMessageTimestamp,
);
return { processedQuery: null, shouldProceed: false };
}
if (pathSpecsToRead.length === 0) {
if (resourceReadDisplays.length > 0) {
addItem(
{ type: 'tool_group', tools: resourceReadDisplays } as Omit<
HistoryItem,
'id'
>,
userMessageTimestamp,
);
}
return { processedQuery: processedQueryParts, shouldProceed: true };
}
const toolArgs = {
include: pathSpecsToRead,
@@ -423,20 +518,20 @@ export async function handleAtCommand({
},
// Use configuration setting
};
let toolCallDisplay: IndividualToolCallDisplay;
let readManyFilesDisplay: IndividualToolCallDisplay | undefined;
let invocation: AnyToolInvocation | undefined = undefined;
try {
invocation = readManyFilesTool.build(toolArgs);
const result = await invocation.execute(signal);
toolCallDisplay = {
readManyFilesDisplay = {
callId: `client-read-${userMessageTimestamp}`,
name: readManyFilesTool.displayName,
description: invocation.getDescription(),
status: ToolCallStatus.Success,
resultDisplay:
result.returnDisplay ||
`Successfully read: ${contentLabelsForDisplay.join(', ')}`,
`Successfully read: ${fileLabelsForDisplay.join(', ')}`,
confirmationDetails: undefined,
};
@@ -486,32 +581,67 @@ export async function handleAtCommand({
);
}
addItem(
{ type: 'tool_group', tools: [toolCallDisplay] } as Omit<
HistoryItem,
'id'
>,
userMessageTimestamp,
);
if (resourceReadDisplays.length > 0 || readManyFilesDisplay) {
addItem(
{
type: 'tool_group',
tools: [
...resourceReadDisplays,
...(readManyFilesDisplay ? [readManyFilesDisplay] : []),
],
} as Omit<HistoryItem, 'id'>,
userMessageTimestamp,
);
}
return { processedQuery: processedQueryParts, shouldProceed: true };
} catch (error: unknown) {
toolCallDisplay = {
readManyFilesDisplay = {
callId: `client-read-${userMessageTimestamp}`,
name: readManyFilesTool.displayName,
description:
invocation?.getDescription() ??
'Error attempting to execute tool to read files',
status: ToolCallStatus.Error,
resultDisplay: `Error reading files (${contentLabelsForDisplay.join(', ')}): ${getErrorMessage(error)}`,
resultDisplay: `Error reading files (${fileLabelsForDisplay.join(', ')}): ${getErrorMessage(error)}`,
confirmationDetails: undefined,
};
addItem(
{ type: 'tool_group', tools: [toolCallDisplay] } as Omit<
HistoryItem,
'id'
>,
{
type: 'tool_group',
tools: [...resourceReadDisplays, readManyFilesDisplay],
} as Omit<HistoryItem, 'id'>,
userMessageTimestamp,
);
return { processedQuery: null, shouldProceed: false };
}
}
function convertResourceContentsToParts(response: {
contents?: Array<{
text?: string;
blob?: string;
mimeType?: string;
resource?: {
text?: string;
blob?: string;
mimeType?: string;
};
}>;
}): PartUnion[] {
const parts: PartUnion[] = [];
for (const content of response.contents ?? []) {
const candidate = content.resource ?? content;
if (candidate.text) {
parts.push({ text: candidate.text });
continue;
}
if (candidate.blob) {
const sizeBytes = Buffer.from(candidate.blob, 'base64').length;
const mimeType = candidate.mimeType ?? 'application/octet-stream';
parts.push({
text: `[Binary resource content ${mimeType}, ${sizeBytes} bytes]`,
});
}
}
return parts;
}

View File

@@ -49,6 +49,9 @@ describe('useAtCompletion', () => {
})),
getEnableRecursiveFileSearch: () => true,
getFileFilteringDisableFuzzySearch: () => false,
getResourceRegistry: vi.fn().mockReturnValue({
getAllResources: () => [],
}),
} as unknown as Config;
vi.clearAllMocks();
});
@@ -174,6 +177,34 @@ describe('useAtCompletion', () => {
});
});
describe('MCP resource suggestions', () => {
it('should include MCP resources in the suggestion list using fuzzy matching', async () => {
mockConfig.getResourceRegistry = vi.fn().mockReturnValue({
getAllResources: () => [
{
serverName: 'server-1',
uri: 'file:///tmp/server-1/logs.txt',
name: 'logs',
discoveredAt: Date.now(),
},
],
});
const { result } = renderHook(() =>
useTestHarnessForAtCompletion(true, 'logs', mockConfig, '/tmp'),
);
await waitFor(() => {
expect(
result.current.suggestions.some(
(suggestion) =>
suggestion.value === 'server-1:file:///tmp/server-1/logs.txt',
),
).toBe(true);
});
});
});
describe('UI State and Loading Behavior', () => {
it('should be in a loading state during initial file system crawl', async () => {
testRootDir = await createTmpDir({});

View File

@@ -9,6 +9,7 @@ import type { Config, FileSearch } from '@google/gemini-cli-core';
import { FileSearchFactory, escapePath } from '@google/gemini-cli-core';
import type { Suggestion } from '../components/SuggestionsDisplay.js';
import { MAX_SUGGESTIONS_TO_SHOW } from '../components/SuggestionsDisplay.js';
import { AsyncFzf } from 'fzf';
export enum AtCompletionStatus {
IDLE = 'idle',
@@ -97,6 +98,61 @@ export interface UseAtCompletionProps {
setIsLoadingSuggestions: (isLoading: boolean) => void;
}
interface ResourceSuggestionCandidate {
searchKey: string;
suggestion: Suggestion;
}
function buildResourceCandidates(
config?: Config,
): ResourceSuggestionCandidate[] {
const registry = config?.getResourceRegistry?.();
if (!registry) {
return [];
}
const resources = registry.getAllResources().map((resource) => {
// Use serverName:uri format to disambiguate resources from different MCP servers
const prefixedUri = `${resource.serverName}:${resource.uri}`;
return {
// Include prefixedUri in searchKey so users can search by the displayed format
searchKey: `${prefixedUri} ${resource.name ?? ''}`.toLowerCase(),
suggestion: {
label: prefixedUri,
value: prefixedUri,
},
} satisfies ResourceSuggestionCandidate;
});
return resources;
}
async function searchResourceCandidates(
pattern: string,
candidates: ResourceSuggestionCandidate[],
): Promise<Suggestion[]> {
if (candidates.length === 0) {
return [];
}
const normalizedPattern = pattern.toLowerCase();
if (!normalizedPattern) {
return candidates
.slice(0, MAX_SUGGESTIONS_TO_SHOW)
.map((candidate) => candidate.suggestion);
}
const fzf = new AsyncFzf(candidates, {
selector: (candidate: ResourceSuggestionCandidate) => candidate.searchKey,
});
const results = await fzf.find(normalizedPattern, {
limit: MAX_SUGGESTIONS_TO_SHOW * 3,
});
return results.map(
(result: { item: ResourceSuggestionCandidate }) => result.item.suggestion,
);
}
export function useAtCompletion(props: UseAtCompletionProps): void {
const {
enabled,
@@ -210,11 +266,28 @@ export function useAtCompletion(props: UseAtCompletionProps): void {
return;
}
const suggestions = results.map((p) => ({
const fileSuggestions = results.map((p) => ({
label: p,
value: escapePath(p),
}));
dispatch({ type: 'SEARCH_SUCCESS', payload: suggestions });
const resourceCandidates = buildResourceCandidates(config);
const resourceSuggestions = (
await searchResourceCandidates(
state.pattern ?? '',
resourceCandidates,
)
).map((suggestion) => ({
...suggestion,
label: suggestion.label.replace(/^@/, ''),
value: suggestion.value.replace(/^@/, ''),
}));
const combinedSuggestions = [
...fileSuggestions,
...resourceSuggestions,
];
dispatch({ type: 'SEARCH_SUCCESS', payload: combinedSuggestions });
} catch (error) {
if (!(error instanceof Error && error.name === 'AbortError')) {
dispatch({ type: 'ERROR' });

View File

@@ -224,11 +224,20 @@ export interface JsonMcpPrompt {
description?: string;
}
export interface JsonMcpResource {
serverName: string;
name?: string;
uri?: string;
mimeType?: string;
description?: string;
}
export type HistoryItemMcpStatus = HistoryItemBase & {
type: 'mcp_status';
servers: Record<string, MCPServerConfig>;
tools: JsonMcpTool[];
prompts: JsonMcpPrompt[];
resources: JsonMcpResource[];
authStatus: Record<
string,
'authenticated' | 'expired' | 'unauthenticated' | 'not-configured'

View File

@@ -11,7 +11,11 @@ export type HighlightToken = {
type: 'default' | 'command' | 'file';
};
const HIGHLIGHT_REGEX = /(^\/[a-zA-Z0-9_-]+|@(?:\\ |[a-zA-Z0-9_./-])+)/g;
// Matches slash commands (e.g., /help) and @ references (files or MCP resource URIs).
// The @ pattern uses a negated character class to support URIs like `@file:///example.txt`
// which contain colons. It matches any character except delimiters: comma, whitespace,
// semicolon, common punctuation, and brackets.
const HIGHLIGHT_REGEX = /(^\/[a-zA-Z0-9_-]+|@(?:\\ |[^,\s;!?()[\]{}])+)/g;
export function parseInputForHighlighting(
text: string,

View File

@@ -17,6 +17,7 @@ import {
createContentGeneratorConfig,
} from '../core/contentGenerator.js';
import { PromptRegistry } from '../prompts/prompt-registry.js';
import { ResourceRegistry } from '../resources/resource-registry.js';
import { ToolRegistry } from '../tools/tool-registry.js';
import { LSTool } from '../tools/ls.js';
import { ReadFileTool } from '../tools/read-file.js';
@@ -331,6 +332,7 @@ export class Config {
private allowedMcpServers: string[];
private blockedMcpServers: string[];
private promptRegistry!: PromptRegistry;
private resourceRegistry!: ResourceRegistry;
private agentRegistry!: AgentRegistry;
private sessionId: string;
private fileSystemService: FileSystemService;
@@ -656,6 +658,7 @@ export class Config {
await this.getGitService();
}
this.promptRegistry = new PromptRegistry();
this.resourceRegistry = new ResourceRegistry();
this.agentRegistry = new AgentRegistry(this);
await this.agentRegistry.initialize();
@@ -921,6 +924,10 @@ export class Config {
return this.promptRegistry;
}
getResourceRegistry(): ResourceRegistry {
return this.resourceRegistry;
}
getDebugMode(): boolean {
return this.debugMode;
}

View File

@@ -103,6 +103,7 @@ export * from './tools/tools.js';
export * from './tools/tool-error.js';
export * from './tools/tool-registry.js';
export * from './tools/tool-names.js';
export * from './resources/resource-registry.js';
// Export prompt logic
export * from './prompts/mcp-prompts.js';

View File

@@ -0,0 +1,76 @@
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, expect, it, beforeEach } from 'vitest';
import type { Resource } from '@modelcontextprotocol/sdk/types.js';
import { ResourceRegistry } from './resource-registry.js';
describe('ResourceRegistry', () => {
let registry: ResourceRegistry;
beforeEach(() => {
registry = new ResourceRegistry();
});
const createResource = (overrides: Partial<Resource> = {}): Resource => ({
uri: 'file:///tmp/foo.txt',
name: 'foo',
description: 'example resource',
mimeType: 'text/plain',
...overrides,
});
it('stores resources per server', () => {
registry.setResourcesForServer('a', [createResource()]);
registry.setResourcesForServer('b', [createResource({ uri: 'foo' })]);
expect(
registry.getAllResources().filter((res) => res.serverName === 'a'),
).toHaveLength(1);
expect(
registry.getAllResources().filter((res) => res.serverName === 'b'),
).toHaveLength(1);
});
it('clears resources for server before adding new ones', () => {
registry.setResourcesForServer('a', [
createResource(),
createResource({ uri: 'bar' }),
]);
registry.setResourcesForServer('a', [createResource({ uri: 'baz' })]);
const resources = registry
.getAllResources()
.filter((res) => res.serverName === 'a');
expect(resources).toHaveLength(1);
expect(resources[0].uri).toBe('baz');
});
it('finds resources by serverName:uri identifier', () => {
registry.setResourcesForServer('a', [createResource()]);
registry.setResourcesForServer('b', [
createResource({ uri: 'file:///tmp/bar.txt' }),
]);
expect(
registry.findResourceByUri('b:file:///tmp/bar.txt')?.serverName,
).toBe('b');
expect(
registry.findResourceByUri('a:file:///tmp/foo.txt')?.serverName,
).toBe('a');
expect(registry.findResourceByUri('a:file:///tmp/bar.txt')).toBeUndefined();
expect(registry.findResourceByUri('nonexistent')).toBeUndefined();
});
it('clears resources for a server', () => {
registry.setResourcesForServer('a', [createResource()]);
registry.removeResourcesByServer('a');
expect(
registry.getAllResources().filter((res) => res.serverName === 'a'),
).toHaveLength(0);
});
});

View File

@@ -0,0 +1,72 @@
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { Resource } from '@modelcontextprotocol/sdk/types.js';
const resourceKey = (serverName: string, uri: string): string =>
`${serverName}::${uri}`;
export interface MCPResource extends Resource {
serverName: string;
discoveredAt: number;
}
export type DiscoveredMCPResource = MCPResource;
/**
* Tracks resources discovered from MCP servers so other
* components can query or include them in conversations.
*/
export class ResourceRegistry {
private resources: Map<string, MCPResource> = new Map();
/**
* Replace the resources for a specific server.
*/
setResourcesForServer(serverName: string, resources: Resource[]): void {
this.removeResourcesByServer(serverName);
const discoveredAt = Date.now();
for (const resource of resources) {
if (!resource.uri) {
continue;
}
this.resources.set(resourceKey(serverName, resource.uri), {
serverName,
discoveredAt,
...resource,
});
}
}
getAllResources(): MCPResource[] {
return Array.from(this.resources.values());
}
/**
* Find a resource by its identifier.
* Format: serverName:uri (e.g., "myserver:file:///data.txt")
*/
findResourceByUri(identifier: string): MCPResource | undefined {
const colonIndex = identifier.indexOf(':');
if (colonIndex <= 0) {
return undefined;
}
const serverName = identifier.substring(0, colonIndex);
const uri = identifier.substring(colonIndex + 1);
return this.resources.get(resourceKey(serverName, uri));
}
removeResourcesByServer(serverName: string): void {
for (const key of Array.from(this.resources.keys())) {
if (key.startsWith(`${serverName}::`)) {
this.resources.delete(key);
}
}
}
clear(): void {
this.resources.clear();
}
}

View File

@@ -29,6 +29,7 @@ vi.mock('./mcp-client.js', async () => {
describe('McpClientManager', () => {
let mockedMcpClient: MockedObject<McpClient>;
let mockConfig: MockedObject<Config>;
let toolRegistry: ToolRegistry;
beforeEach(() => {
mockedMcpClient = vi.mockObject({
@@ -43,6 +44,7 @@ describe('McpClientManager', () => {
isTrustedFolder: vi.fn().mockReturnValue(true),
getMcpServers: vi.fn().mockReturnValue({}),
getPromptRegistry: () => {},
getResourceRegistry: () => {},
getDebugMode: () => false,
getWorkspaceContext: () => {},
getAllowedMcpServers: vi.fn().mockReturnValue([]),
@@ -52,6 +54,7 @@ describe('McpClientManager', () => {
isInitialized: vi.fn(),
}),
} as unknown as Config);
toolRegistry = {} as ToolRegistry;
});
afterEach(() => {
@@ -62,7 +65,7 @@ describe('McpClientManager', () => {
mockConfig.getMcpServers.mockReturnValue({
'test-server': {},
});
const manager = new McpClientManager({} as ToolRegistry, mockConfig);
const manager = new McpClientManager(toolRegistry, mockConfig);
await manager.startConfiguredMcpServers();
expect(mockedMcpClient.connect).toHaveBeenCalledOnce();
expect(mockedMcpClient.discover).toHaveBeenCalledOnce();
@@ -73,7 +76,7 @@ describe('McpClientManager', () => {
'test-server': {},
});
mockConfig.isTrustedFolder.mockReturnValue(false);
const manager = new McpClientManager({} as ToolRegistry, mockConfig);
const manager = new McpClientManager(toolRegistry, mockConfig);
await manager.startConfiguredMcpServers();
expect(mockedMcpClient.connect).not.toHaveBeenCalled();
expect(mockedMcpClient.discover).not.toHaveBeenCalled();
@@ -84,7 +87,7 @@ describe('McpClientManager', () => {
'test-server': {},
});
mockConfig.getBlockedMcpServers.mockReturnValue(['test-server']);
const manager = new McpClientManager({} as ToolRegistry, mockConfig);
const manager = new McpClientManager(toolRegistry, mockConfig);
await manager.startConfiguredMcpServers();
expect(mockedMcpClient.connect).not.toHaveBeenCalled();
expect(mockedMcpClient.discover).not.toHaveBeenCalled();
@@ -96,14 +99,14 @@ describe('McpClientManager', () => {
'another-server': {},
});
mockConfig.getAllowedMcpServers.mockReturnValue(['another-server']);
const manager = new McpClientManager({} as ToolRegistry, mockConfig);
const manager = new McpClientManager(toolRegistry, mockConfig);
await manager.startConfiguredMcpServers();
expect(mockedMcpClient.connect).toHaveBeenCalledOnce();
expect(mockedMcpClient.discover).toHaveBeenCalledOnce();
});
it('should start servers from extensions', async () => {
const manager = new McpClientManager({} as ToolRegistry, mockConfig);
const manager = new McpClientManager(toolRegistry, mockConfig);
await manager.startExtension({
name: 'test-extension',
mcpServers: {
@@ -120,7 +123,7 @@ describe('McpClientManager', () => {
});
it('should not start servers from disabled extensions', async () => {
const manager = new McpClientManager({} as ToolRegistry, mockConfig);
const manager = new McpClientManager(toolRegistry, mockConfig);
await manager.startExtension({
name: 'test-extension',
mcpServers: {
@@ -141,7 +144,7 @@ describe('McpClientManager', () => {
'test-server': {},
});
mockConfig.getBlockedMcpServers.mockReturnValue(['test-server']);
const manager = new McpClientManager({} as ToolRegistry, mockConfig);
const manager = new McpClientManager(toolRegistry, mockConfig);
await manager.startConfiguredMcpServers();
expect(manager.getBlockedMcpServers()).toEqual([
{ name: 'test-server', extensionName: '' },
@@ -154,7 +157,7 @@ describe('McpClientManager', () => {
'test-server': {},
});
mockedMcpClient.getServerConfig.mockReturnValue({});
const manager = new McpClientManager({} as ToolRegistry, mockConfig);
const manager = new McpClientManager(toolRegistry, mockConfig);
await manager.startConfiguredMcpServers();
expect(mockedMcpClient.connect).toHaveBeenCalledTimes(1);
@@ -173,7 +176,7 @@ describe('McpClientManager', () => {
'test-server': {},
});
mockedMcpClient.getServerConfig.mockReturnValue({});
const manager = new McpClientManager({} as ToolRegistry, mockConfig);
const manager = new McpClientManager(toolRegistry, mockConfig);
await manager.startConfiguredMcpServers();
expect(mockedMcpClient.connect).toHaveBeenCalledTimes(1);
@@ -187,7 +190,7 @@ describe('McpClientManager', () => {
});
it('should throw an error if the server does not exist', async () => {
const manager = new McpClientManager({} as ToolRegistry, mockConfig);
const manager = new McpClientManager(toolRegistry, mockConfig);
await expect(manager.restartServer('non-existent')).rejects.toThrow(
'No MCP server registered with the name "non-existent"',
);

View File

@@ -52,6 +52,10 @@ export class McpClientManager {
return this.blockedMcpServers;
}
getClient(serverName: string): McpClient | undefined {
return this.clients.get(serverName);
}
/**
* For all the MCP servers associated with this extension:
*
@@ -174,6 +178,7 @@ export class McpClientManager {
config,
this.toolRegistry,
this.cliConfig.getPromptRegistry(),
this.cliConfig.getResourceRegistry(),
this.cliConfig.getWorkspaceContext(),
this.cliConfig,
this.cliConfig.getDebugMode(),

View File

@@ -30,6 +30,7 @@ import {
populateMcpServerCommand,
} from './mcp-client.js';
import type { ToolRegistry } from './tool-registry.js';
import type { ResourceRegistry } from '../resources/resource-registry.js';
import * as fs from 'node:fs';
import * as os from 'node:os';
import * as path from 'node:path';
@@ -77,6 +78,7 @@ describe('mcp-client', () => {
getStatus: vi.fn(),
registerCapabilities: vi.fn(),
setRequestHandler: vi.fn(),
setNotificationHandler: vi.fn(),
getServerCapabilities: vi.fn().mockReturnValue({ tools: {} }),
listTools: vi.fn().mockResolvedValue({
tools: [
@@ -105,13 +107,22 @@ describe('mcp-client', () => {
sortTools: vi.fn(),
getMessageBus: vi.fn().mockReturnValue(undefined),
} as unknown as ToolRegistry;
const promptRegistry = {
registerPrompt: vi.fn(),
removePromptsByServer: vi.fn(),
} as unknown as PromptRegistry;
const resourceRegistry = {
setResourcesForServer: vi.fn(),
removeResourcesByServer: vi.fn(),
} as unknown as ResourceRegistry;
const client = new McpClient(
'test-server',
{
command: 'test-command',
},
mockedToolRegistry,
{} as PromptRegistry,
promptRegistry,
resourceRegistry,
workspaceContext,
{} as Config,
false,
@@ -135,6 +146,7 @@ describe('mcp-client', () => {
getStatus: vi.fn(),
registerCapabilities: vi.fn(),
setRequestHandler: vi.fn(),
setNotificationHandler: vi.fn(),
getServerCapabilities: vi.fn().mockReturnValue({ tools: {} }),
listTools: vi.fn().mockResolvedValue({
@@ -175,13 +187,22 @@ describe('mcp-client', () => {
sortTools: vi.fn(),
getMessageBus: vi.fn().mockReturnValue(undefined),
} as unknown as ToolRegistry;
const promptRegistry = {
registerPrompt: vi.fn(),
removePromptsByServer: vi.fn(),
} as unknown as PromptRegistry;
const resourceRegistry = {
setResourcesForServer: vi.fn(),
removeResourcesByServer: vi.fn(),
} as unknown as ResourceRegistry;
const client = new McpClient(
'test-server',
{
command: 'test-command',
},
mockedToolRegistry,
{} as PromptRegistry,
promptRegistry,
resourceRegistry,
workspaceContext,
{} as Config,
false,
@@ -201,6 +222,7 @@ describe('mcp-client', () => {
getStatus: vi.fn(),
registerCapabilities: vi.fn(),
setRequestHandler: vi.fn(),
setNotificationHandler: vi.fn(),
getServerCapabilities: vi.fn().mockReturnValue({ prompts: {} }),
listTools: vi.fn().mockResolvedValue({ tools: [] }),
listPrompts: vi.fn().mockRejectedValue(new Error('Test error')),
@@ -216,20 +238,29 @@ describe('mcp-client', () => {
registerTool: vi.fn(),
getMessageBus: vi.fn().mockReturnValue(undefined),
} as unknown as ToolRegistry;
const promptRegistry = {
registerPrompt: vi.fn(),
removePromptsByServer: vi.fn(),
} as unknown as PromptRegistry;
const resourceRegistry = {
setResourcesForServer: vi.fn(),
removeResourcesByServer: vi.fn(),
} as unknown as ResourceRegistry;
const client = new McpClient(
'test-server',
{
command: 'test-command',
},
mockedToolRegistry,
{} as PromptRegistry,
promptRegistry,
resourceRegistry,
workspaceContext,
{} as Config,
false,
);
await client.connect();
await expect(client.discover({} as Config)).rejects.toThrow(
'No prompts or tools found on the server.',
'No prompts, tools, or resources found on the server.',
);
expect(coreEvents.emitFeedback).toHaveBeenCalledWith(
'error',
@@ -246,6 +277,7 @@ describe('mcp-client', () => {
getStatus: vi.fn(),
registerCapabilities: vi.fn(),
setRequestHandler: vi.fn(),
setNotificationHandler: vi.fn(),
getServerCapabilities: vi.fn().mockReturnValue({ prompts: {} }),
listPrompts: vi.fn().mockResolvedValue({ prompts: [] }),
request: vi.fn().mockResolvedValue({}),
@@ -261,20 +293,29 @@ describe('mcp-client', () => {
sortTools: vi.fn(),
getMessageBus: vi.fn().mockReturnValue(undefined),
} as unknown as ToolRegistry;
const promptRegistry = {
registerPrompt: vi.fn(),
removePromptsByServer: vi.fn(),
} as unknown as PromptRegistry;
const resourceRegistry = {
setResourcesForServer: vi.fn(),
removeResourcesByServer: vi.fn(),
} as unknown as ResourceRegistry;
const client = new McpClient(
'test-server',
{
command: 'test-command',
},
mockedToolRegistry,
{} as PromptRegistry,
promptRegistry,
resourceRegistry,
workspaceContext,
{} as Config,
false,
);
await client.connect();
await expect(client.discover({} as Config)).rejects.toThrow(
'No prompts or tools found on the server.',
'No prompts, tools, or resources found on the server.',
);
});
@@ -286,6 +327,7 @@ describe('mcp-client', () => {
getStatus: vi.fn(),
registerCapabilities: vi.fn(),
setRequestHandler: vi.fn(),
setNotificationHandler: vi.fn(),
getServerCapabilities: vi.fn().mockReturnValue({ tools: {} }),
listTools: vi.fn().mockResolvedValue({
tools: [
@@ -310,13 +352,22 @@ describe('mcp-client', () => {
sortTools: vi.fn(),
getMessageBus: vi.fn().mockReturnValue(undefined),
} as unknown as ToolRegistry;
const promptRegistry = {
registerPrompt: vi.fn(),
removePromptsByServer: vi.fn(),
} as unknown as PromptRegistry;
const resourceRegistry = {
setResourcesForServer: vi.fn(),
removeResourcesByServer: vi.fn(),
} as unknown as ResourceRegistry;
const client = new McpClient(
'test-server',
{
command: 'test-command',
},
mockedToolRegistry,
{} as PromptRegistry,
promptRegistry,
resourceRegistry,
workspaceContext,
{} as Config,
false,
@@ -334,6 +385,7 @@ describe('mcp-client', () => {
getStatus: vi.fn(),
registerCapabilities: vi.fn(),
setRequestHandler: vi.fn(),
setNotificationHandler: vi.fn(),
getServerCapabilities: vi.fn().mockReturnValue({ tools: {} }),
listTools: vi.fn().mockResolvedValue({
tools: [
@@ -373,13 +425,22 @@ describe('mcp-client', () => {
sortTools: vi.fn(),
getMessageBus: vi.fn().mockReturnValue(undefined),
} as unknown as ToolRegistry;
const promptRegistry = {
registerPrompt: vi.fn(),
removePromptsByServer: vi.fn(),
} as unknown as PromptRegistry;
const resourceRegistry = {
setResourcesForServer: vi.fn(),
removeResourcesByServer: vi.fn(),
} as unknown as ResourceRegistry;
const client = new McpClient(
'test-server',
{
command: 'test-command',
},
mockedToolRegistry,
{} as PromptRegistry,
promptRegistry,
resourceRegistry,
workspaceContext,
{} as Config,
false,
@@ -405,6 +466,165 @@ describe('mcp-client', () => {
});
});
it('should discover resources when a server only exposes resources', async () => {
const mockedClient = {
connect: vi.fn(),
discover: vi.fn(),
disconnect: vi.fn(),
getStatus: vi.fn(),
registerCapabilities: vi.fn(),
setRequestHandler: vi.fn(),
setNotificationHandler: vi.fn(),
getServerCapabilities: vi.fn().mockReturnValue({ resources: {} }),
request: vi.fn().mockImplementation(({ method }) => {
if (method === 'resources/list') {
return Promise.resolve({
resources: [
{
uri: 'file:///tmp/resource.txt',
name: 'resource',
description: 'Test Resource',
mimeType: 'text/plain',
},
],
});
}
return Promise.resolve({ prompts: [] });
}),
} as unknown as ClientLib.Client;
vi.mocked(ClientLib.Client).mockReturnValue(mockedClient);
vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue(
{} as SdkClientStdioLib.StdioClientTransport,
);
const mockedToolRegistry = {
registerTool: vi.fn(),
sortTools: vi.fn(),
getMessageBus: vi.fn().mockReturnValue(undefined),
} as unknown as ToolRegistry;
const promptRegistry = {
registerPrompt: vi.fn(),
removePromptsByServer: vi.fn(),
} as unknown as PromptRegistry;
const resourceRegistry = {
setResourcesForServer: vi.fn(),
removeResourcesByServer: vi.fn(),
} as unknown as ResourceRegistry;
const client = new McpClient(
'test-server',
{
command: 'test-command',
},
mockedToolRegistry,
promptRegistry,
resourceRegistry,
workspaceContext,
{} as Config,
false,
);
await client.connect();
await client.discover({} as Config);
expect(resourceRegistry.setResourcesForServer).toHaveBeenCalledWith(
'test-server',
[
expect.objectContaining({
uri: 'file:///tmp/resource.txt',
name: 'resource',
}),
],
);
});
it('refreshes registry when resource list change notification is received', async () => {
let listCallCount = 0;
let resourceListHandler:
| ((notification: unknown) => Promise<void> | void)
| undefined;
const mockedClient = {
connect: vi.fn(),
discover: vi.fn(),
disconnect: vi.fn(),
getStatus: vi.fn(),
registerCapabilities: vi.fn(),
setRequestHandler: vi.fn(),
setNotificationHandler: vi.fn((_, handler) => {
resourceListHandler = handler;
}),
getServerCapabilities: vi
.fn()
.mockReturnValue({ resources: { listChanged: true } }),
request: vi.fn().mockImplementation(({ method }) => {
if (method === 'resources/list') {
listCallCount += 1;
if (listCallCount === 1) {
return Promise.resolve({
resources: [
{
uri: 'file:///tmp/one.txt',
},
],
});
}
return Promise.resolve({
resources: [
{
uri: 'file:///tmp/two.txt',
},
],
});
}
return Promise.resolve({ prompts: [] });
}),
} as unknown as ClientLib.Client;
vi.mocked(ClientLib.Client).mockReturnValue(mockedClient);
vi.spyOn(SdkClientStdioLib, 'StdioClientTransport').mockReturnValue(
{} as SdkClientStdioLib.StdioClientTransport,
);
const mockedToolRegistry = {
registerTool: vi.fn(),
sortTools: vi.fn(),
getMessageBus: vi.fn().mockReturnValue(undefined),
} as unknown as ToolRegistry;
const promptRegistry = {
registerPrompt: vi.fn(),
removePromptsByServer: vi.fn(),
} as unknown as PromptRegistry;
const resourceRegistry = {
setResourcesForServer: vi.fn(),
removeResourcesByServer: vi.fn(),
} as unknown as ResourceRegistry;
const client = new McpClient(
'test-server',
{
command: 'test-command',
},
mockedToolRegistry,
promptRegistry,
resourceRegistry,
workspaceContext,
{} as Config,
false,
);
await client.connect();
await client.discover({} as Config);
expect(mockedClient.setNotificationHandler).toHaveBeenCalledOnce();
expect(resourceListHandler).toBeDefined();
await resourceListHandler?.({
method: 'notifications/resources/list_changed',
});
expect(resourceRegistry.setResourcesForServer).toHaveBeenLastCalledWith(
'test-server',
[expect.objectContaining({ uri: 'file:///tmp/two.txt' })],
);
expect(coreEvents.emitFeedback).toHaveBeenCalledWith(
'info',
'Resources updated for server: test-server',
);
});
it('should remove tools and prompts on disconnect', async () => {
const mockedClient = {
connect: vi.fn(),
@@ -412,6 +632,7 @@ describe('mcp-client', () => {
getStatus: vi.fn(),
registerCapabilities: vi.fn(),
setRequestHandler: vi.fn(),
setNotificationHandler: vi.fn(),
getServerCapabilities: vi
.fn()
.mockReturnValue({ tools: {}, prompts: {} }),
@@ -447,6 +668,10 @@ describe('mcp-client', () => {
unregisterPrompt: vi.fn(),
removePromptsByServer: vi.fn(),
} as unknown as PromptRegistry;
const resourceRegistry = {
setResourcesForServer: vi.fn(),
removeResourcesByServer: vi.fn(),
} as unknown as ResourceRegistry;
const client = new McpClient(
'test-server',
{
@@ -454,6 +679,7 @@ describe('mcp-client', () => {
},
mockedToolRegistry,
mockedPromptRegistry,
resourceRegistry,
workspaceContext,
{} as Config,
false,
@@ -469,6 +695,7 @@ describe('mcp-client', () => {
expect(mockedClient.close).toHaveBeenCalledOnce();
expect(mockedToolRegistry.removeMcpToolsByServer).toHaveBeenCalledOnce();
expect(mockedPromptRegistry.removePromptsByServer).toHaveBeenCalledOnce();
expect(resourceRegistry.removeResourcesByServer).toHaveBeenCalledOnce();
});
});
@@ -501,6 +728,7 @@ describe('mcp-client', () => {
{ command: 'test-command' },
{} as ToolRegistry,
{} as PromptRegistry,
{} as ResourceRegistry,
workspaceContext,
{} as Config,
false,
@@ -536,6 +764,7 @@ describe('mcp-client', () => {
{ command: 'test-command' },
{} as ToolRegistry,
{} as PromptRegistry,
{} as ResourceRegistry,
workspaceContext,
{} as Config,
false,
@@ -590,6 +819,7 @@ describe('mcp-client', () => {
{ command: 'test-command' },
mockedToolRegistry,
{} as PromptRegistry,
{} as ResourceRegistry,
workspaceContext,
{} as Config,
false,
@@ -659,6 +889,7 @@ describe('mcp-client', () => {
{ command: 'test-command' },
mockedToolRegistry,
{} as PromptRegistry,
{} as ResourceRegistry,
workspaceContext,
{} as Config,
false,
@@ -728,6 +959,7 @@ describe('mcp-client', () => {
{ command: 'cmd-a' },
mockedToolRegistry,
{} as PromptRegistry,
{} as ResourceRegistry,
workspaceContext,
{} as Config,
false,
@@ -739,6 +971,7 @@ describe('mcp-client', () => {
{ command: 'cmd-b' },
mockedToolRegistry,
{} as PromptRegistry,
{} as ResourceRegistry,
workspaceContext,
{} as Config,
false,
@@ -820,6 +1053,7 @@ describe('mcp-client', () => {
{ command: 'test-command', timeout: 100 },
mockedToolRegistry,
{} as PromptRegistry,
{} as ResourceRegistry,
workspaceContext,
{} as Config,
false,
@@ -883,6 +1117,7 @@ describe('mcp-client', () => {
{ command: 'test-command' },
mockedToolRegistry,
{} as PromptRegistry,
{} as ResourceRegistry,
workspaceContext,
{} as Config,
false,

View File

@@ -20,9 +20,14 @@ import type { Transport } from '@modelcontextprotocol/sdk/shared/transport.js';
import type {
GetPromptResult,
Prompt,
ReadResourceResult,
Resource,
} from '@modelcontextprotocol/sdk/types.js';
import {
ListResourcesResultSchema,
ListRootsRequestSchema,
ReadResourceResultSchema,
ResourceListChangedNotificationSchema,
ToolListChangedNotificationSchema,
type Tool as McpTool,
} from '@modelcontextprotocol/sdk/types.js';
@@ -54,6 +59,7 @@ import type { ToolRegistry } from './tool-registry.js';
import { debugLogger } from '../utils/debugLogger.js';
import type { MessageBus } from '../confirmation-bus/message-bus.js';
import { coreEvents } from '../utils/events.js';
import type { ResourceRegistry } from '../resources/resource-registry.js';
export const MCP_DEFAULT_TIMEOUT_MSEC = 10 * 60 * 1000; // default to 10 minutes
@@ -98,14 +104,17 @@ export class McpClient {
private client: Client | undefined;
private transport: Transport | undefined;
private status: MCPServerStatus = MCPServerStatus.DISCONNECTED;
private isRefreshing: boolean = false;
private pendingRefresh: boolean = false;
private isRefreshingTools: boolean = false;
private pendingToolRefresh: boolean = false;
private isRefreshingResources: boolean = false;
private pendingResourceRefresh: boolean = false;
constructor(
private readonly serverName: string,
private readonly serverConfig: MCPServerConfig,
private readonly toolRegistry: ToolRegistry,
private readonly promptRegistry: PromptRegistry,
private readonly resourceRegistry: ResourceRegistry,
private readonly workspaceContext: WorkspaceContext,
private readonly cliConfig: Config,
private readonly debugMode: boolean,
@@ -130,24 +139,8 @@ export class McpClient {
this.workspaceContext,
);
// setup dynamic tool listener
const capabilities = this.client.getServerCapabilities();
this.registerNotificationHandlers();
if (capabilities?.tools?.listChanged) {
debugLogger.log(
`Server '${this.serverName}' supports tool updates. Listening for changes...`,
);
this.client.setNotificationHandler(
ToolListChangedNotificationSchema,
async () => {
debugLogger.log(
`🔔 Received tool update notification from '${this.serverName}'`,
);
await this.refreshTools();
},
);
}
const originalOnError = this.client.onerror;
this.client.onerror = (error) => {
if (this.status !== MCPServerStatus.CONNECTED) {
@@ -176,9 +169,11 @@ export class McpClient {
const prompts = await this.discoverPrompts();
const tools = await this.discoverTools(cliConfig);
const resources = await this.discoverResources();
this.updateResourceRegistry(resources);
if (prompts.length === 0 && tools.length === 0) {
throw new Error('No prompts or tools found on the server.');
if (prompts.length === 0 && tools.length === 0 && resources.length === 0) {
throw new Error('No prompts, tools, or resources found on the server.');
}
for (const tool of tools) {
@@ -196,6 +191,7 @@ export class McpClient {
}
this.toolRegistry.removeMcpToolsByServer(this.serverName);
this.promptRegistry.removePromptsByServer(this.serverName);
this.resourceRegistry.removeResourcesByServer(this.serverName);
this.updateStatus(MCPServerStatus.DISCONNECTING);
const client = this.client;
this.client = undefined;
@@ -250,6 +246,128 @@ export class McpClient {
return discoverPrompts(this.serverName, this.client!, this.promptRegistry);
}
private async discoverResources(): Promise<Resource[]> {
this.assertConnected();
return discoverResources(this.serverName, this.client!);
}
private updateResourceRegistry(resources: Resource[]): void {
this.resourceRegistry.setResourcesForServer(this.serverName, resources);
}
async readResource(uri: string): Promise<ReadResourceResult> {
this.assertConnected();
return this.client!.request(
{
method: 'resources/read',
params: { uri },
},
ReadResourceResultSchema,
);
}
/**
* Registers notification handlers for dynamic updates from the MCP server.
* This includes handlers for tool list changes and resource list changes.
*/
private registerNotificationHandlers(): void {
if (!this.client) {
return;
}
const capabilities = this.client.getServerCapabilities();
if (capabilities?.tools?.listChanged) {
debugLogger.log(
`Server '${this.serverName}' supports tool updates. Listening for changes...`,
);
this.client.setNotificationHandler(
ToolListChangedNotificationSchema,
async () => {
debugLogger.log(
`🔔 Received tool update notification from '${this.serverName}'`,
);
await this.refreshTools();
},
);
}
if (capabilities?.resources?.listChanged) {
debugLogger.log(
`Server '${this.serverName}' supports resource updates. Listening for changes...`,
);
this.client.setNotificationHandler(
ResourceListChangedNotificationSchema,
async () => {
debugLogger.log(
`🔔 Received resource update notification from '${this.serverName}'`,
);
await this.refreshResources();
},
);
}
}
/**
* Refreshes the resources for this server by re-querying the MCP `resources/list` endpoint.
*
* This method implements a **Coalescing Pattern** to handle rapid bursts of notifications
* (e.g., during server startup or bulk updates) without overwhelming the server or
* creating race conditions in the ResourceRegistry.
*/
private async refreshResources(): Promise<void> {
if (this.isRefreshingResources) {
debugLogger.log(
`Resource refresh for '${this.serverName}' is already in progress. Pending update.`,
);
this.pendingResourceRefresh = true;
return;
}
this.isRefreshingResources = true;
try {
do {
this.pendingResourceRefresh = false;
if (this.status !== MCPServerStatus.CONNECTED || !this.client) break;
const timeoutMs = this.serverConfig.timeout ?? MCP_DEFAULT_TIMEOUT_MSEC;
const abortController = new AbortController();
const timeoutId = setTimeout(() => abortController.abort(), timeoutMs);
let newResources;
try {
newResources = await this.discoverResources();
} catch (err) {
debugLogger.error(
`Resource discovery failed during refresh: ${getErrorMessage(err)}`,
);
clearTimeout(timeoutId);
break;
}
this.updateResourceRegistry(newResources);
clearTimeout(timeoutId);
coreEvents.emitFeedback(
'info',
`Resources updated for server: ${this.serverName}`,
);
} while (this.pendingResourceRefresh);
} catch (error) {
debugLogger.error(
`Critical error in resource refresh loop for ${this.serverName}: ${getErrorMessage(error)}`,
);
} finally {
this.isRefreshingResources = false;
this.pendingResourceRefresh = false;
}
}
getServerConfig(): MCPServerConfig {
return this.serverConfig;
}
@@ -266,19 +384,19 @@ export class McpClient {
* creating race conditions in the global ToolRegistry.
*/
private async refreshTools(): Promise<void> {
if (this.isRefreshing) {
if (this.isRefreshingTools) {
debugLogger.log(
`Tool refresh for '${this.serverName}' is already in progress. Pending update.`,
);
this.pendingRefresh = true;
this.pendingToolRefresh = true;
return;
}
this.isRefreshing = true;
this.isRefreshingTools = true;
try {
do {
this.pendingRefresh = false;
this.pendingToolRefresh = false;
if (this.status !== MCPServerStatus.CONNECTED || !this.client) break;
@@ -316,14 +434,14 @@ export class McpClient {
'info',
`Tools updated for server: ${this.serverName}`,
);
} while (this.pendingRefresh);
} while (this.pendingToolRefresh);
} catch (error) {
debugLogger.error(
`Critical error in refresh loop for ${this.serverName}: ${getErrorMessage(error)}`,
);
} finally {
this.isRefreshing = false;
this.pendingRefresh = false;
this.isRefreshingTools = false;
this.pendingToolRefresh = false;
}
}
}
@@ -944,6 +1062,52 @@ export async function discoverPrompts(
}
}
export async function discoverResources(
mcpServerName: string,
mcpClient: Client,
): Promise<Resource[]> {
if (mcpClient.getServerCapabilities()?.resources == null) {
return [];
}
const resources = await listResources(mcpServerName, mcpClient);
return resources;
}
async function listResources(
mcpServerName: string,
mcpClient: Client,
): Promise<Resource[]> {
const resources: Resource[] = [];
let cursor: string | undefined;
try {
do {
const response = await mcpClient.request(
{
method: 'resources/list',
params: cursor ? { cursor } : {},
},
ListResourcesResultSchema,
);
resources.push(...(response.resources ?? []));
cursor = response.nextCursor ?? undefined;
} while (cursor);
} catch (error) {
if (error instanceof Error && error.message?.includes('Method not found')) {
return [];
}
coreEvents.emitFeedback(
'error',
`Error discovering resources from ${mcpServerName}: ${getErrorMessage(
error,
)}`,
error,
);
throw error;
}
return resources;
}
/**
* Invokes a prompt on a connected MCP client.
*