mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-02-01 22:48:03 +00:00
feat: Add support for MCP Resources (#13178)
Co-authored-by: Jack Wotherspoon <jackwoth@google.com>
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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),
|
||||
);
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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>
|
||||
);
|
||||
})}
|
||||
|
||||
@@ -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:
|
||||
|
||||
|
||||
@@ -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),
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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({});
|
||||
|
||||
@@ -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' });
|
||||
|
||||
@@ -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'
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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';
|
||||
|
||||
76
packages/core/src/resources/resource-registry.test.ts
Normal file
76
packages/core/src/resources/resource-registry.test.ts
Normal 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);
|
||||
});
|
||||
});
|
||||
72
packages/core/src/resources/resource-registry.ts
Normal file
72
packages/core/src/resources/resource-registry.ts
Normal 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();
|
||||
}
|
||||
}
|
||||
@@ -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"',
|
||||
);
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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.
|
||||
*
|
||||
|
||||
Reference in New Issue
Block a user