Implement MCP tool registry and integration in PRReviewer#2356
Open
benny338813 wants to merge 4 commits intoThe-PR-Agent:mainfrom
Open
Implement MCP tool registry and integration in PRReviewer#2356benny338813 wants to merge 4 commits intoThe-PR-Agent:mainfrom
benny338813 wants to merge 4 commits intoThe-PR-Agent:mainfrom
Conversation
added 4 commits
April 29, 2026 01:39
Contributor
Review Summary by QodoImplement MCP tool registry and LLM tool calling integration
WalkthroughsDescription• Implement tool registry system for dynamic tool registration • Add MCP (Model Context Protocol) integration framework • Refactor tool instantiation to use registry pattern • Support tool calling with LLM integration Diagramflowchart LR
A["Tool Classes"] -->|register| B["ToolRegistry"]
B -->|lookup| C["PRAgent"]
C -->|get_tools| D["MCPHandler"]
D -->|call_tool| E["LLM with Tools"]
E -->|tool_calls| D
File Changes1. pr_agent/tools/registry.py
|
Contributor
Code Review by Qodo
1. Commented-out command2class block
|
![qodo-free-for-open-source-projects[bot]](https://avatars.githubusercontent.com/in/1068000?s=60&v=4){kind=small}
Comment on lines
+28
to
+50
| # command2class = { | ||
| # "auto_review": PRReviewer, | ||
| # "answer": PRReviewer, | ||
| # "review": PRReviewer, | ||
| # "review_pr": PRReviewer, | ||
| # "describe": PRDescription, | ||
| # "describe_pr": PRDescription, | ||
| # "improve": PRCodeSuggestions, | ||
| # "improve_code": PRCodeSuggestions, | ||
| # "ask": PRQuestions, | ||
| # "ask_question": PRQuestions, | ||
| # "ask_line": PR_LineQuestions, | ||
| # "update_changelog": PRUpdateChangelog, | ||
| # "config": PRConfig, | ||
| # "settings": PRConfig, | ||
| # "help": PRHelpMessage, | ||
| # "similar_issue": PRSimilarIssue, | ||
| # "add_docs": PRAddDocs, | ||
| # "generate_labels": PRGenerateLabels, | ||
| # "help_docs": PRHelpDocs, | ||
| # } | ||
|
|
||
| # commands = list(command2class.keys()) |
Contributor
There was a problem hiding this comment.
1. Commented-out command2class block 📘 Rule violation ⚙ Maintainability
A large command2class mapping is left in the code as commented-out lines, which is dead code and adds maintenance noise. This violates the requirement to avoid commented-out code in submitted changes.
Agent Prompt
## Issue description
A large commented-out `command2class` block was added, which is dead code and violates the no-commented-out-code requirement.
## Issue Context
The tool dispatch was replaced by `ToolRegistry`, so this legacy mapping should be removed rather than kept commented.
## Fix Focus Areas
- pr_agent/agent/pr_agent.py[28-50]
ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools
| if not tool_class: | ||
| get_logger().warning(f"Unknown command: {action}") | ||
| return False | ||
|
|
Contributor
There was a problem hiding this comment.
2. Whitespace-only line has trailing spaces 📘 Rule violation ⚙ Maintainability
A whitespace-only line with indentation spaces was added, which will fail pre-commit trailing-whitespace hooks. This violates the repository’s pre-commit hygiene requirements.
Agent Prompt
## Issue description
A blank line containing spaces was introduced, which violates pre-commit trailing-whitespace rules.
## Issue Context
This will typically fail CI/pre-commit and should be cleaned up.
## Fix Focus Areas
- pr_agent/agent/pr_agent.py[110-110]
ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools
Comment on lines
+1
to
+20
| from typing import Dict, Type | ||
| from pr_agent.tools.base import PRTool | ||
|
|
||
| class ToolRegistry: | ||
| _tools: Dict[str, Type[PRTool]] = {} | ||
|
|
||
| @classmethod | ||
| def register(cls, command: str): | ||
| def wrapper(subclass: Type[PRTool]): | ||
| cls._tools[command] = subclass | ||
| return subclass | ||
| return wrapper | ||
|
|
||
| @classmethod | ||
| def get_tool(cls, command: str) -> Optional[Type[PRTool]]: | ||
| return cls._tools.get(command) | ||
|
|
||
| @classmethod | ||
| def get_all_commands(cls) -> List[str]: | ||
| return list(cls._tools.keys()) |
Contributor
There was a problem hiding this comment.
3. toolregistry missing typing imports 📘 Rule violation ≡ Correctness
ToolRegistry uses Optional and List in type hints without importing them, which will raise NameError at import time and fail linting. This introduces a correctness/lint violation.
Agent Prompt
## Issue description
`ToolRegistry` references `Optional` and `List` but does not import them, causing runtime import errors and failing lint checks.
## Issue Context
This is a new file introduced by the PR and should be clean under Ruff.
## Fix Focus Areas
- pr_agent/tools/registry.py[1-20]
ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools
Comment on lines
+1
to
+6
| import json | ||
| from typing import List, Dict, Any, Optional | ||
| from mcp import ClientSession, StdioServerParameters | ||
| from mcp.client.stdio import stdio_client | ||
| import asyncio | ||
|
|
Contributor
There was a problem hiding this comment.
4. mcp_handler has unused imports 📘 Rule violation ⚙ Maintainability
pr_agent/algo/mcp_handler.py imports json and asyncio but does not use them, which will be flagged by Ruff/flake8-style unused-import checks. This violates the repository lint/tooling compliance requirement.
Agent Prompt
## Issue description
The new `mcp_handler.py` file includes unused imports that will fail lint checks.
## Issue Context
Ruff typically enforces unused import removal.
## Fix Focus Areas
- pr_agent/algo/mcp_handler.py[1-6]
ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools
Comment on lines
+212
to
+216
| mcp_config = get_settings().get("mcp", {}) | ||
| if mcp_config: | ||
| from pr_agent.algo.mcp_handler import MCPHandler | ||
| mcp_handler = MCPHandler(mcp_config.command, mcp_config.args) | ||
| async with mcp_handler as handler: |
Contributor
There was a problem hiding this comment.
5. mcp_config accessed without validation 📘 Rule violation ☼ Reliability
mcp_config is fetched via get_settings().get("mcp", {}) and then accessed as
mcp_config.command/mcp_config.args without validating presence/types, which can raise
AttributeError at runtime for missing/incorrect config shapes. This violates the requirement to
use a consistent config access pattern and validate/normalize config values before use.
Agent Prompt
## Issue description
`mcp_config` is retrieved in a way that may yield a plain dict or an object missing required keys, but the code assumes `mcp_config.command` and `mcp_config.args` always exist.
## Issue Context
Dynaconf values can vary by layer/override; missing keys should be handled gracefully.
## Fix Focus Areas
- pr_agent/tools/pr_reviewer.py[212-216]
ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools
Comment on lines
+228
to
+246
| if finish_reason == "tool_calls" and tool_calls: | ||
| # Execute MCP tools | ||
| for tool_call in tool_calls: | ||
| tool_name = tool_call.function.name | ||
| tool_args = json.loads(tool_call.function.arguments) | ||
| get_logger().info(f"Executing tool: {tool_name} with args {tool_args}") | ||
| tool_result = await handler.call_tool(tool_name, tool_args) | ||
|
|
||
| # Add tool result to user prompt | ||
| user_prompt += f"\n\nTool Result ({tool_name}): {json.dumps(tool_result)}" | ||
|
|
||
| # Re-run completion with tool results | ||
| response, finish_reason, _ = await self.ai_handler.chat_completion( | ||
| model=model, | ||
| temperature=get_settings().config.temperature, | ||
| system=system_prompt, | ||
| user=user_prompt, | ||
| tools=None # Don't allow recursive tool calls for now | ||
| ) |
Contributor
There was a problem hiding this comment.
6. Tool call args not guarded 📘 Rule violation ☼ Reliability
Tool execution assumes tool_call.function.arguments is valid JSON and that MCP tool calls always succeed, but exceptions from json.loads(...) or handler.call_tool(...) are not handled. This can crash the review flow instead of failing gracefully, violating robust error handling and defensive input validation requirements.
Agent Prompt
## Issue description
Tool calling can raise exceptions (invalid JSON arguments, MCP server errors) and currently crashes the review path.
## Issue Context
Tool calls are driven by model output (external/untrusted shape) and an external MCP server; both require defensive handling.
## Fix Focus Areas
- pr_agent/tools/pr_reviewer.py[228-246]
ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools
| get_logger().info(f"\nAI response:\n{resp}") | ||
|
|
||
| return resp, finish_reason | ||
| return resp, finish_reason, tool_calls |
Contributor
There was a problem hiding this comment.
7. litellmaihandler now returns 3 📘 Rule violation ≡ Correctness
LiteLLMAIHandler.chat_completion() now returns three values (resp, finish_reason, tool_calls), but multiple call sites still unpack only two values, which will raise `ValueError: too many values to unpack` at runtime. This is a robustness issue that can break normal tool execution flows.
Agent Prompt
## Issue description
`LiteLLMAIHandler.chat_completion()` now returns a 3-tuple, but many callers unpack a 2-tuple, causing runtime failures.
## Issue Context
Either keep backward compatibility (e.g., return 2 values unless tool-calling is enabled / provide a separate method), or update all call sites (and the `BaseAiHandler` interface if needed) to unpack the third return value.
## Fix Focus Areas
- pr_agent/algo/ai_handlers/litellm_ai_handler.py[270-272]
- pr_agent/algo/ai_handlers/litellm_ai_handler.py[420-442]
- pr_agent/tools/pr_add_docs.py[92-93]
ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools
Comment on lines
105
to
123
| action = action.lstrip("/").lower() | ||
| if action not in command2class: | ||
| tool_class = ToolRegistry.get_tool(action) | ||
| if not tool_class: | ||
| get_logger().warning(f"Unknown command: {action}") | ||
| return False | ||
|
|
||
| with get_logger().contextualize(command=action, pr_url=pr_url): | ||
| get_logger().info("PR-Agent request handler started", analytics=True) | ||
| if action == "answer": | ||
| if notify: | ||
| notify() | ||
| await PRReviewer(pr_url, is_answer=True, args=args, ai_handler=self.ai_handler).run() | ||
| elif action == "auto_review": | ||
| await PRReviewer(pr_url, is_auto=True, args=args, ai_handler=self.ai_handler).run() | ||
| elif action in command2class: | ||
| if notify: | ||
| notify() | ||
|
|
||
| await command2class[action](pr_url, ai_handler=self.ai_handler, args=args).run() | ||
|
|
||
| if notify: | ||
| notify() | ||
|
|
||
| if tool_class == PRReviewer: | ||
| is_answer = (action == "answer") | ||
| is_auto = (action == "auto_review") | ||
| await PRReviewer(pr_url, is_answer=is_answer, is_auto=is_auto, args=args, ai_handler=self.ai_handler).run() | ||
| else: | ||
| return False | ||
| await tool_class(pr_url, ai_handler=self.ai_handler, args=args).run() | ||
| return True |
Contributor
There was a problem hiding this comment.
8. Unregistered commands now fail 🐞 Bug ≡ Correctness
PRAgent now dispatches exclusively via ToolRegistry, but only PRReviewer/PRDescription are registered; commands like improve, ask, add_docs, etc. will be treated as unknown and the CLI will reject them because choices=commands is derived from the registry.
Agent Prompt
### Issue description
Command dispatch now depends on `ToolRegistry`, but most tools are not registered, so many commands become unusable and CLI `choices` excludes them.
### Issue Context
Previously, `command2class` provided coverage for many commands. After the migration, each tool must self-register (or a compatibility mapping must remain).
### Fix Focus Areas
- pr_agent/agent/pr_agent.py[14-52]
- pr_agent/agent/pr_agent.py[105-123]
- pr_agent/cli.py[5-56]
- pr_agent/tools/pr_code_suggestions.py[33-60]
- pr_agent/tools/pr_add_docs.py[18-40]
- pr_agent/tools/pr_questions.py[18-60]
- pr_agent/tools/pr_config.py[8-31]
### Expected change
Add `@ToolRegistry.register(...)` decorators for every supported command (e.g., `improve`, `improve_code`, `ask`, `ask_question`, `add_docs`, `update_changelog`, `help_docs`, `generate_labels`, `config/settings`, `help`, `similar_issue`, etc.), or reintroduce a backward-compatible `command2class`/fallback dispatch while migrating.
ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools
Comment on lines
+39
to
64
| def __init__(self, pr_url: str, ai_handler: partial[BaseAiHandler,] = LiteLLMAIHandler, args: list = None): | ||
| super().__init__(pr_url, ai_handler=ai_handler, args=args) | ||
| self.git_provider = get_git_provider_with_context(pr_url) | ||
| self.main_pr_language = get_main_pr_language( | ||
| self.git_provider.get_languages(), self.git_provider.get_files() | ||
| self.ai_handler = ai_handler() | ||
| self.vars = { | ||
| "title": self.git_provider.pr.title, | ||
| "branch": self.git_provider.get_pr_branch(), | ||
| "description": self.git_provider.get_pr_description(), | ||
| "language": get_main_pr_language( | ||
| self.git_provider.get_languages(), self.git_provider.get_files() | ||
| ), | ||
| "commit_messages_str": self.git_provider.get_commit_messages(), | ||
| "extra_instructions": get_settings().pr_description_prompt.extra_instructions, | ||
| 'duplicate_prompt_examples': get_settings().config.get('duplicate_prompt_examples', False), | ||
| "enable_custom_labels": get_settings().config.enable_custom_labels, | ||
| } | ||
| self.token_handler = TokenHandler( | ||
| self.git_provider.pr, | ||
| self.vars, | ||
| get_settings().pr_description_prompt.system, | ||
| get_settings().pr_description_prompt.user | ||
| ) | ||
|
|
||
| self.pr_id = self.git_provider.get_pr_id() | ||
| self.keys_fix = ["filename:", "language:", "changes_summary:", "changes_title:", "description:", "title:"] | ||
|
|
Contributor
There was a problem hiding this comment.
9. Prdescription init crashes 🐞 Bug ≡ Correctness
PRDescription.__init__ now duplicates initialization logic and references self.main_pr_language even though it is no longer set, causing AttributeError during tool construction.
Agent Prompt
### Issue description
`PRDescription.__init__` contains two initialization blocks and uses `self.main_pr_language` without defining it, causing an `AttributeError` and inconsistent object state.
### Issue Context
This appears to be a partial refactor to `PRTool`/registry, leaving old initialization in place.
### Fix Focus Areas
- pr_agent/tools/pr_description.py[39-103]
### Expected change
Remove the duplicated initialization and ensure `self.main_pr_language` (or equivalent) is set exactly once before being used (and align `self.vars` / `TokenHandler` creation to that single source of truth).
ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools
Comment on lines
+228
to
+238
| if finish_reason == "tool_calls" and tool_calls: | ||
| # Execute MCP tools | ||
| for tool_call in tool_calls: | ||
| tool_name = tool_call.function.name | ||
| tool_args = json.loads(tool_call.function.arguments) | ||
| get_logger().info(f"Executing tool: {tool_name} with args {tool_args}") | ||
| tool_result = await handler.call_tool(tool_name, tool_args) | ||
|
|
||
| # Add tool result to user prompt | ||
| user_prompt += f"\n\nTool Result ({tool_name}): {json.dumps(tool_result)}" | ||
|
|
Contributor
There was a problem hiding this comment.
10. Mcp tool_calls parsing crash 🐞 Bug ≡ Correctness
When MCP tools are enabled, PRReviewer assumes each tool_call has .function.name/.arguments, but the LiteLLM response passes tool_calls as dicts, so tool execution will crash before any MCP tool can run.
Agent Prompt
### Issue description
`PRReviewer` parses `tool_calls` using attribute access (`tool_call.function.name`), but the handler returns tool calls as dictionaries from the OpenAI/LiteLLM response.
### Issue Context
Tool calls are typically shaped like:
```json
{"id": "...", "type": "function", "function": {"name": "...", "arguments": "{...}"}}
```
### Fix Focus Areas
- pr_agent/tools/pr_reviewer.py[228-238]
- pr_agent/algo/ai_handlers/litellm_ai_handler.py[458-470]
### Expected change
Update parsing to support dict tool calls (e.g., `tool_call["function"]["name"]` and `tool_call["function"]["arguments"]`), and add defensive handling for missing/invalid JSON arguments so MCP execution can proceed or fail gracefully.
ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.