Flexible MCPClient Initialization #705
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @suluyana, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
I've introduced significant enhancements to the "MCPClient" and "ToolManager" classes, primarily focusing on improving their lifecycle management and configuration flexibility. This involves enabling both classes to function as asynchronous context managers, which simplifies resource handling. Furthermore, the "MCPClient" now supports dynamic updates to its server configurations, allowing for more adaptable tool management. The "LLMAgent" has been updated to integrate with these changes, facilitating a more controlled injection of the "MCPClient".
Highlights
- Asynchronous Context Management: Both "MCPClient" and "ToolManager" now implement "aenter" and "aexit", allowing them to be used with "async with" statements for streamlined resource acquisition and release.
- Dynamic MCP Configuration: The "MCPClient" gains a new "add_mcp_config" method, enabling the addition of new MCP server configurations to an already initialized client, enhancing runtime adaptability.
- LLMAgent Integration: The "LLMAgent"'s initialization now accepts an "mcp_client" instance, which is then passed to the "ToolManager", providing more explicit control over the "MCPClient"'s lifecycle.
- Flexible MCPClient Initialization: The "MCPClient"'s constructor has been refactored to allow its "config" parameter to be optional and to handle "mcp_config" more flexibly, improving its usability in various scenarios.
- Comprehensive Unit Tests: A new dedicated test file, "tests/tools/test_mcp_client.py", has been added to thoroughly validate the new functionalities of "MCPClient", including its context manager behavior and dynamic configuration updates.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments or fill out our survey to provide feedback.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Code Review
This pull request refactors the agent to allow passing a pre-initialized MCPClient, which is a good improvement for resource management. It also makes MCPClient an async context manager and adds a comprehensive set of unit tests for it.
However, the review identified several critical issues in the implementation. The initialization logic for MCPClient is fragile and can lead to KeyError exceptions. The logic in ToolManager.connect for handling a passed-in client is incorrect and will cause a crash. Additionally, there's a resource leak in ToolManager.cleanup where it fails to clean up MCPClient instances it creates. These issues need to be addressed to ensure the stability and correctness of the changes.
ms_agent/tools/mcp_client.py
Outdated
| self.mcp_config: Dict[str, Dict[str, Any]] = dict() | ||
| if config is not None: | ||
| self.mcp_config = Config.convert_mcp_servers_to_json(config) | ||
| self._exclude_functions = {} | ||
| if mcp_config is not None: | ||
| self.mcp_config.update(mcp_config) |
There was a problem hiding this comment.
The initialization logic for self.mcp_config is fragile and can lead to runtime errors.
- If
MCPClientis initialized without any arguments,self.mcp_configwill be an empty dictionary. This will cause aKeyErrorin methods likeconnect()that expect themcpServerskey. - The use of
self.mcp_config.update(mcp_config)will incorrectly overwrite server configurations fromconfigifmcp_configis also provided, instead of merging them.
I suggest a more robust implementation that correctly initializes and merges the configurations.
| self.mcp_config: Dict[str, Dict[str, Any]] = dict() | |
| if config is not None: | |
| self.mcp_config = Config.convert_mcp_servers_to_json(config) | |
| self._exclude_functions = {} | |
| if mcp_config is not None: | |
| self.mcp_config.update(mcp_config) | |
| self.mcp_config: Dict[str, Dict[str, Any]] = {'mcpServers': {}} | |
| if config is not None: | |
| config_from_file = Config.convert_mcp_servers_to_json(config) | |
| self.mcp_config['mcpServers'].update(config_from_file.get('mcpServers', {})) | |
| self._exclude_functions = {} | |
| if mcp_config is not None: | |
| self.mcp_config['mcpServers'].update(mcp_config.get('mcpServers', {})) |
ms_agent/tools/tool_manager.py
Outdated
| self.servers = await self.mcp_client.add_mcp_config(self.mcp_config) | ||
| self.mcp_config = self.servers.mcp_confg |
There was a problem hiding this comment.
There are a couple of issues in this block that will cause a crash.
- The
add_mcp_configmethod returnsNone, so assigning its result toself.serverswill makeself.serversNone. You should probably assignself.mcp_clienttoself.servers. - The next line attempts to access
self.servers.mcp_confg, which will raise anAttributeErrorsinceself.serversisNone. There is also a typo inmcp_confg.
| self.servers = await self.mcp_client.add_mcp_config(self.mcp_config) | |
| self.mcp_config = self.servers.mcp_confg | |
| self.servers = self.mcp_client | |
| await self.servers.add_mcp_config(self.mcp_config) |
ms_agent/tools/mcp_client.py
Outdated
| new_mcp_config = mcp_config['mcpServers'] | ||
| servers = self.mcp_config["mcpServers"] |
There was a problem hiding this comment.
The direct dictionary access on these lines can raise a KeyError if the mcpServers key is missing from mcp_config or self.mcp_config. This can happen if the client was initialized without a valid configuration. To make this method more robust, you should use safe access methods like .get() or .setdefault().
| new_mcp_config = mcp_config['mcpServers'] | |
| servers = self.mcp_config["mcpServers"] | |
| new_mcp_config = mcp_config.get('mcpServers', {}) | |
| servers = self.mcp_config.setdefault("mcpServers", {}) |
| for tool in self.extra_tools: | ||
| await tool.connect() | ||
| await self.reindex_tool() | ||
|
|
||
| async def cleanup(self): | ||
| await self.servers.cleanup() | ||
| self.servers = None |
There was a problem hiding this comment.
Replacing await self.servers.cleanup() with self.servers = None introduces a resource leak. If the ToolManager creates its own MCPClient instance (in the else branch of the connect method), it is responsible for cleaning it up. The current implementation fails to do so.
The cleanup should be conditional. You could track if the client is managed internally. For example:
In __init__:
self._managed_client = mcp_client is NoneIn cleanup:
if self._managed_client and self.servers:
await self.servers.cleanup()
self.servers = None
suluyana
changed the title
Change Summary
Related issue number
Checklist
pre-commit installandpre-commit run --all-filesbefore git commit, and passed lint check.