diff --git a/docs/tools/mcp-server.md b/docs/tools/mcp-server.md index 5d24407f57..a489f833c2 100644 --- a/docs/tools/mcp-server.md +++ b/docs/tools/mcp-server.md @@ -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 server’s `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 diff --git a/packages/cli/src/ui/commands/mcpCommand.test.ts b/packages/cli/src/ui/commands/mcpCommand.test.ts index 91eea4acd7..45bc935fe6 100644 --- a/packages/cli/src/ui/commands/mcpCommand.test.ts +++ b/packages/cli/src/ui/commands/mcpCommand.test.ts @@ -63,6 +63,7 @@ describe('mcpCommand', () => { getPromptRegistry: ReturnType; getGeminiClient: ReturnType; getMcpClientManager: ReturnType; + getResourceRegistry: ReturnType; }; 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), ); diff --git a/packages/cli/src/ui/commands/mcpCommand.ts b/packages/cli/src/ui/commands/mcpCommand.ts index 254a2b1cb5..45c4cb2577 100644 --- a/packages/cli/src/ui/commands/mcpCommand.ts +++ b/packages/cli/src/ui/commands/mcpCommand.ts @@ -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, diff --git a/packages/cli/src/ui/components/views/McpStatus.test.tsx b/packages/cli/src/ui/components/views/McpStatus.test.tsx index 0fb37510b5..86dbc7bb36 100644 --- a/packages/cli/src/ui/components/views/McpStatus.test.tsx +++ b/packages/cli/src/ui/components/views/McpStatus.test.tsx @@ -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( + , + ); + expect(lastFrame()).toMatchSnapshot(); + unmount(); + }); + it('renders correctly with a blocked server', () => { const { lastFrame, unmount } = render( ; 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 = ({ servers, tools, prompts, + resources, blockedServers, serverStatus, authStatus, @@ -83,9 +86,14 @@ export const McpStatus: React.FC = ({ 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 = ({ 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 = ({ `${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 = ({ ))} )} + + {serverResources.length > 0 && ( + + Resources: + {serverResources.map((resource, index) => { + const label = resource.name || resource.uri || 'resource'; + return ( + + + - {label} + {resource.uri ? ` (${resource.uri})` : ''} + {resource.mimeType ? ` [${resource.mimeType}]` : ''} + + {showDescriptions && resource.description && ( + + + {resource.description.trim()} + + + )} + + ); + })} + + )} ); })} diff --git a/packages/cli/src/ui/components/views/__snapshots__/McpStatus.test.tsx.snap b/packages/cli/src/ui/components/views/__snapshots__/McpStatus.test.tsx.snap index bf04cfb381..e7b7af36f2 100644 --- a/packages/cli/src/ui/components/views/__snapshots__/McpStatus.test.tsx.snap +++ b/packages/cli/src/ui/components/views/__snapshots__/McpStatus.test.tsx.snap @@ -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: diff --git a/packages/cli/src/ui/hooks/atCommandProcessor.test.ts b/packages/cli/src/ui/hooks/atCommandProcessor.test.ts index a2f6c35f66..68eaa423de 100644 --- a/packages/cli/src/ui/hooks/atCommandProcessor.test.ts +++ b/packages/cli/src/ui/hooks/atCommandProcessor.test.ts @@ -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), + ); + }); + }); }); diff --git a/packages/cli/src/ui/hooks/atCommandProcessor.ts b/packages/cli/src/ui/hooks/atCommandProcessor.ts index 6d3e9071f8..6f620f3f4d 100644 --- a/packages/cli/src/ui/hooks/atCommandProcessor.ts +++ b/packages/cli/src/ui/hooks/atCommandProcessor.ts @@ -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 '@' 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 '@' 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 { + 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(); - const contentLabelsForDisplay: string[] = []; + const fileLabelsForDisplay: string[] = []; const absoluteToRelativePathMap = new Map(); const ignoredByReason: Record = { 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, + 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, 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; +} diff --git a/packages/cli/src/ui/hooks/useAtCompletion.test.ts b/packages/cli/src/ui/hooks/useAtCompletion.test.ts index 74c92b9159..c4e7c98d19 100644 --- a/packages/cli/src/ui/hooks/useAtCompletion.test.ts +++ b/packages/cli/src/ui/hooks/useAtCompletion.test.ts @@ -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({}); diff --git a/packages/cli/src/ui/hooks/useAtCompletion.ts b/packages/cli/src/ui/hooks/useAtCompletion.ts index 1be9a42e06..c361089947 100644 --- a/packages/cli/src/ui/hooks/useAtCompletion.ts +++ b/packages/cli/src/ui/hooks/useAtCompletion.ts @@ -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 { + 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' }); diff --git a/packages/cli/src/ui/types.ts b/packages/cli/src/ui/types.ts index 84ffc40694..32bd291e5d 100644 --- a/packages/cli/src/ui/types.ts +++ b/packages/cli/src/ui/types.ts @@ -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; tools: JsonMcpTool[]; prompts: JsonMcpPrompt[]; + resources: JsonMcpResource[]; authStatus: Record< string, 'authenticated' | 'expired' | 'unauthenticated' | 'not-configured' diff --git a/packages/cli/src/ui/utils/highlight.ts b/packages/cli/src/ui/utils/highlight.ts index 1e48b36f13..19866b330a 100644 --- a/packages/cli/src/ui/utils/highlight.ts +++ b/packages/cli/src/ui/utils/highlight.ts @@ -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, diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 2d26e29966..b9a78ca7b9 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -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; } diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index c06c0a8e84..e667185ca5 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -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'; diff --git a/packages/core/src/resources/resource-registry.test.ts b/packages/core/src/resources/resource-registry.test.ts new file mode 100644 index 0000000000..29b34d316d --- /dev/null +++ b/packages/core/src/resources/resource-registry.test.ts @@ -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 => ({ + 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); + }); +}); diff --git a/packages/core/src/resources/resource-registry.ts b/packages/core/src/resources/resource-registry.ts new file mode 100644 index 0000000000..1c2c754504 --- /dev/null +++ b/packages/core/src/resources/resource-registry.ts @@ -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 = 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(); + } +} diff --git a/packages/core/src/tools/mcp-client-manager.test.ts b/packages/core/src/tools/mcp-client-manager.test.ts index 5bb3db7bc5..0ff201e8e6 100644 --- a/packages/core/src/tools/mcp-client-manager.test.ts +++ b/packages/core/src/tools/mcp-client-manager.test.ts @@ -29,6 +29,7 @@ vi.mock('./mcp-client.js', async () => { describe('McpClientManager', () => { let mockedMcpClient: MockedObject; let mockConfig: MockedObject; + 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"', ); diff --git a/packages/core/src/tools/mcp-client-manager.ts b/packages/core/src/tools/mcp-client-manager.ts index 23a34f45d9..fbc3b3e423 100644 --- a/packages/core/src/tools/mcp-client-manager.ts +++ b/packages/core/src/tools/mcp-client-manager.ts @@ -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(), diff --git a/packages/core/src/tools/mcp-client.test.ts b/packages/core/src/tools/mcp-client.test.ts index 2075d13c3c..192025689d 100644 --- a/packages/core/src/tools/mcp-client.test.ts +++ b/packages/core/src/tools/mcp-client.test.ts @@ -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) + | 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, diff --git a/packages/core/src/tools/mcp-client.ts b/packages/core/src/tools/mcp-client.ts index f0ef143cf1..86f0043a7d 100644 --- a/packages/core/src/tools/mcp-client.ts +++ b/packages/core/src/tools/mcp-client.ts @@ -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 { + this.assertConnected(); + return discoverResources(this.serverName, this.client!); + } + + private updateResourceRegistry(resources: Resource[]): void { + this.resourceRegistry.setResourcesForServer(this.serverName, resources); + } + + async readResource(uri: string): Promise { + 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 { + 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 { - 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 { + if (mcpClient.getServerCapabilities()?.resources == null) { + return []; + } + + const resources = await listResources(mcpServerName, mcpClient); + return resources; +} + +async function listResources( + mcpServerName: string, + mcpClient: Client, +): Promise { + 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. *