Skip to content

Commit 8008c02

Browse files
committed
feat: add MCP tool orchestration and server config support
Introduce server-side MCP config loading from /etc/pr-agent/mcp.json or MCP_CONFIG_PATH, including JSONC parsing and VS Code / Claude schema normalization. Add the MCP runtime, HTTP and stdio clients, structured tool-calling orchestration on the base AI handler, and wire /ask, /review, and /improve through the MCP-aware integration helper. Expose MCP runtime status in /config output, document the configuration flow and AWS Knowledge example, and add focused tests for config loading, runtime behavior, tool orchestration, integration, and discovery.
1 parent 0e37fc8 commit 8008c02

17 files changed

Lines changed: 1514 additions & 14 deletions

docs/docs/usage-guide/additional_configurations.md

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,27 @@ To print all the available configurations as a comment on your PR, you can use t
99
/config
1010
```
1111

12+
When MCP is enabled, the `/config` comment also includes a small MCP runtime status block showing whether MCP is enabled and which servers are configured and connected.
13+
14+
## MCP runtime configuration
15+
16+
PR-Agent can load MCP servers from a server-side JSON or JSONC file. By default, it reads `/etc/pr-agent/mcp.json`, and you can override that path with `MCP_CONFIG_PATH` or the `[mcp].config_path` setting.
17+
18+
The file may use either the `servers` key, which matches the VS Code MCP schema, or `mcpServers`, which matches the Claude Desktop schema.
19+
20+
For example, an AWS Knowledge MCP server can be configured like this:
21+
22+
```json
23+
{
24+
"servers": {
25+
"AWS Knowledge": {
26+
"url": "https://knowledge-mcp.global.api.aws",
27+
"type": "http"
28+
}
29+
}
30+
}
31+
```
32+
1233
![possible_config1](https://codium.ai/images/pr_agent/possible_config1.png){width=512}
1334

1435
To view the **actual** configurations used for a specific tool, after all the user settings are applied, you can add for each tool a `--config.output_relevant_configurations=true` suffix.

docs/docs/usage-guide/automations_and_usage.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,8 @@ For example, if you want to edit the `review` tool configurations, you can run:
7575

7676
Any configuration value in [configuration file](https://github.com/the-pr-agent/pr-agent/blob/main/pr_agent/settings/configuration.toml) file can be similarly edited. Comment `/config` to see the list of available configurations.
7777

78+
If you want PR-Agent to use MCP tools, mount a server-side MCP config file at `/etc/pr-agent/mcp.json` or point `MCP_CONFIG_PATH` at another JSON/JSONC file. The `/config` comment will show the active MCP runtime status when MCP is enabled.
79+
7880
## PR-Agent Automatic Feedback
7981

8082
### Disabling all automatic feedback

pr_agent/algo/ai_handlers/base_ai_handler.py

Lines changed: 149 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,7 @@
1+
import inspect
2+
import json
13
from abc import ABC, abstractmethod
4+
from typing import Any, Awaitable, Callable, Optional
25

36

47
class BaseAiHandler(ABC):
@@ -26,3 +29,149 @@ async def chat_completion(self, model: str, system: str, user: str, temperature:
2629
temperature (float): the temperature to use for the chat completion
2730
"""
2831
pass
32+
33+
async def chat_completion_with_tools(
34+
self,
35+
model: str,
36+
system: str,
37+
user: str,
38+
tools: Optional[list[dict[str, Any]]] = None,
39+
tool_executor: Optional[Callable[[str, dict[str, Any]], Any | Awaitable[Any]]] = None,
40+
temperature: float = 0.2,
41+
img_path: str = None,
42+
max_tool_turns: int = 4,
43+
max_tool_output_chars: int = 12000,
44+
):
45+
"""
46+
Run a structured tool-calling loop on top of plain chat completion.
47+
48+
The model is instructed to emit JSON tool requests in the form:
49+
{"type": "tool_call", "tool": "server.tool", "arguments": {...}}
50+
and to finish with:
51+
{"type": "final", "content": "..."}
52+
"""
53+
if not tools or tool_executor is None:
54+
return await self.chat_completion(model, system, user, temperature=temperature, img_path=img_path)
55+
56+
tool_catalog_text = json.dumps(tools, indent=2, sort_keys=True)
57+
structured_system = (
58+
f"{system}\n\n"
59+
f"Available MCP tools (JSON schema):\n{tool_catalog_text}\n\n"
60+
"When you need a tool, respond with ONLY a JSON object exactly in this shape:\n"
61+
'{"type":"tool_call","tool":"server.tool","arguments":{...}}\n'
62+
"Do not include a final answer in the same message as a tool call.\n"
63+
"When you are finished, respond with ONLY a JSON object exactly in this shape:\n"
64+
'{"type":"final","content":"..."}\n'
65+
"Do not wrap the JSON in markdown fences."
66+
)
67+
68+
conversation_history = [user]
69+
remaining_turns = max_tool_turns
70+
current_img_path = img_path
71+
72+
while True:
73+
current_user = "\n\n".join(conversation_history)
74+
response_text, finish_reason = await self.chat_completion(
75+
model=model,
76+
system=structured_system,
77+
user=current_user,
78+
temperature=temperature,
79+
img_path=current_img_path,
80+
)
81+
current_img_path = None
82+
83+
parsed_response = self._parse_tool_or_final_response(response_text)
84+
if parsed_response is None:
85+
return response_text, finish_reason
86+
87+
response_type = parsed_response.get("type", "final")
88+
if response_type == "final":
89+
return str(parsed_response.get("content", "")), finish_reason
90+
91+
if response_type != "tool_call":
92+
return response_text, finish_reason
93+
94+
if remaining_turns <= 0:
95+
raise ValueError("MCP tool orchestration exceeded the configured turn budget")
96+
97+
tool_name = str(parsed_response.get("tool", "")).strip()
98+
arguments = parsed_response.get("arguments") or {}
99+
if not tool_name:
100+
raise ValueError("MCP tool orchestration returned an empty tool name")
101+
if not isinstance(arguments, dict):
102+
raise ValueError("MCP tool orchestration arguments must be a JSON object")
103+
104+
tool_result = tool_executor(tool_name, arguments)
105+
if inspect.isawaitable(tool_result):
106+
tool_result = await tool_result
107+
108+
tool_result_text = self._normalize_tool_result_text(tool_result, max_tool_output_chars)
109+
conversation_history.append(f"Previous assistant tool request:\n{response_text}")
110+
conversation_history.append(f"Tool result for {tool_name}:\n{tool_result_text}")
111+
remaining_turns -= 1
112+
113+
@staticmethod
114+
def _normalize_tool_result_text(tool_result: Any, max_tool_output_chars: int) -> str:
115+
if isinstance(tool_result, str):
116+
result_text = tool_result
117+
else:
118+
result_text = json.dumps(tool_result, indent=2, sort_keys=True, default=str)
119+
120+
if len(result_text) > max_tool_output_chars:
121+
return result_text[: max_tool_output_chars - 20] + "\n[tool output truncated]"
122+
return result_text
123+
124+
@staticmethod
125+
def _parse_tool_or_final_response(response_text: str) -> Optional[dict[str, Any]]:
126+
candidate = response_text.strip()
127+
if not candidate:
128+
return None
129+
130+
for json_candidate in BaseAiHandler._iter_json_object_candidates(candidate):
131+
try:
132+
parsed = json.loads(json_candidate)
133+
except json.JSONDecodeError:
134+
continue
135+
136+
if isinstance(parsed, dict):
137+
response_type = parsed.get("type")
138+
if response_type in {"tool_call", "final"}:
139+
return parsed
140+
141+
return None
142+
143+
@staticmethod
144+
def _iter_json_object_candidates(text: str) -> list[str]:
145+
candidates: list[str] = []
146+
depth = 0
147+
start_index: Optional[int] = None
148+
in_string = False
149+
is_escaped = False
150+
151+
for index, char in enumerate(text):
152+
if in_string:
153+
if is_escaped:
154+
is_escaped = False
155+
elif char == "\\":
156+
is_escaped = True
157+
elif char == '"':
158+
in_string = False
159+
continue
160+
161+
if char == '"':
162+
in_string = True
163+
continue
164+
165+
if char == "{":
166+
if depth == 0:
167+
start_index = index
168+
depth += 1
169+
continue
170+
171+
if char == "}" and depth > 0:
172+
depth -= 1
173+
if depth == 0 and start_index is not None:
174+
candidates.append(text[start_index : index + 1])
175+
start_index = None
176+
177+
return candidates

pr_agent/config_loader.py

Lines changed: 135 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,15 @@
1+
import json
2+
import os
13
from os.path import abspath, dirname, join
24
from pathlib import Path
3-
from typing import Optional
5+
from typing import Any, Optional
46

57
from dynaconf import Dynaconf
68
from starlette_context import context
79

810
PR_AGENT_TOML_KEY = 'pr-agent'
11+
MCP_CONFIG_ENV_VAR = "MCP_CONFIG_PATH"
12+
DEFAULT_MCP_CONFIG_PATH = "/etc/pr-agent/mcp.json"
913

1014
current_dir = dirname(abspath(__file__))
1115

@@ -60,6 +64,134 @@ def get_settings(use_context=False):
6064
return global_settings
6165

6266

67+
def _get_logger():
68+
try:
69+
from pr_agent.log import get_logger
70+
return get_logger()
71+
except Exception:
72+
class DummyLogger:
73+
def debug(self, *args, **kwargs): pass
74+
def info(self, *args, **kwargs): pass
75+
def warning(self, *args, **kwargs): pass
76+
def error(self, *args, **kwargs): pass
77+
return DummyLogger()
78+
79+
80+
def _strip_json_comments(content: str) -> str:
81+
"""Strip line and block comments from JSONC-style config while preserving newlines."""
82+
stripped = []
83+
in_string = False
84+
in_line_comment = False
85+
in_block_comment = False
86+
is_escaped = False
87+
index = 0
88+
89+
while index < len(content):
90+
char = content[index]
91+
next_char = content[index + 1] if index + 1 < len(content) else ""
92+
93+
if in_line_comment:
94+
if char == "\n":
95+
in_line_comment = False
96+
stripped.append(char)
97+
index += 1
98+
continue
99+
100+
if in_block_comment:
101+
if char == "*" and next_char == "/":
102+
in_block_comment = False
103+
index += 2
104+
continue
105+
if char == "\n":
106+
stripped.append(char)
107+
index += 1
108+
continue
109+
110+
if in_string:
111+
stripped.append(char)
112+
if is_escaped:
113+
is_escaped = False
114+
elif char == "\\":
115+
is_escaped = True
116+
elif char == '"':
117+
in_string = False
118+
index += 1
119+
continue
120+
121+
if char == '"':
122+
in_string = True
123+
stripped.append(char)
124+
index += 1
125+
continue
126+
127+
if char == "/" and next_char == "/":
128+
in_line_comment = True
129+
index += 2
130+
continue
131+
132+
if char == "/" and next_char == "*":
133+
in_block_comment = True
134+
index += 2
135+
continue
136+
137+
stripped.append(char)
138+
index += 1
139+
140+
return "".join(stripped)
141+
142+
143+
def _resolve_mcp_config_path() -> Path:
144+
env_path = os.getenv(MCP_CONFIG_ENV_VAR)
145+
if env_path:
146+
return Path(env_path).expanduser()
147+
configured_path = get_settings().get("MCP.CONFIG_PATH", DEFAULT_MCP_CONFIG_PATH)
148+
return Path(str(configured_path)).expanduser()
149+
150+
151+
def _normalize_mcp_servers(config_data: dict[str, Any]) -> dict[str, Any]:
152+
servers = config_data.get("servers")
153+
if servers is None:
154+
servers = config_data.get("mcpServers")
155+
if servers is None:
156+
raise ValueError("MCP config must define either 'servers' or 'mcpServers'")
157+
if not isinstance(servers, dict):
158+
raise ValueError("MCP server definitions must be a JSON object")
159+
return servers
160+
161+
162+
def load_mcp_server_config(config_path: Path) -> dict[str, Any]:
163+
if not config_path.is_file():
164+
raise FileNotFoundError(f"MCP config file not found: {config_path}")
165+
config_text = config_path.read_text(encoding="utf-8")
166+
try:
167+
config_data = json.loads(_strip_json_comments(config_text))
168+
except json.JSONDecodeError as exc:
169+
raise ValueError(f"Invalid MCP config JSON in {config_path}: {exc}") from exc
170+
if not isinstance(config_data, dict):
171+
raise ValueError("MCP config root must be a JSON object")
172+
servers = _normalize_mcp_servers(config_data)
173+
return {"servers": servers}
174+
175+
176+
def apply_mcp_server_config():
177+
logger = _get_logger()
178+
config_path = _resolve_mcp_config_path()
179+
if not config_path.exists():
180+
logger.debug(f"MCP config file not found, skipping load: {config_path}")
181+
return
182+
try:
183+
config_data = load_mcp_server_config(config_path)
184+
settings = get_settings()
185+
settings.set("MCP.SERVERS", config_data["servers"], merge=False)
186+
settings.set("MCP.SERVER_CONFIG", config_data, merge=False)
187+
settings.set("MCP.ACTIVE_CONFIG_PATH", str(config_path), merge=False)
188+
logger.info(f"Loaded MCP server configuration from {config_path}")
189+
except Exception as exc:
190+
logger.error(f"Failed to load MCP server configuration from {config_path}: {exc}")
191+
if get_settings().get("MCP.FAIL_ON_INVALID_CONFIG", False):
192+
raise
193+
194+
63195
# Add local configuration from pyproject.toml of the project being reviewed
64196
def _find_repository_root() -> Optional[Path]:
65197
"""
@@ -90,6 +222,8 @@ def _find_pyproject() -> Optional[Path]:
90222
if pyproject_path is not None:
91223
get_settings().load_file(pyproject_path, env=f'tool.{PR_AGENT_TOML_KEY}')
92224

225+
apply_mcp_server_config()
226+
93227

94228
def apply_secrets_manager_config():
95229
"""

pr_agent/mcp/__init__.py

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
from pr_agent.mcp.runtime import (
2+
MCPHttpClient,
3+
MCPRuntime,
4+
MCPRuntimeError,
5+
MCPStdioClient,
6+
MCPToolDefinition,
7+
)
8+
9+
__all__ = [
10+
"MCPRuntime",
11+
"MCPRuntimeError",
12+
"MCPToolDefinition",
13+
"MCPStdioClient",
14+
"MCPHttpClient",
15+
]

0 commit comments

Comments
 (0)