Skip to content

Commit 7a7991a

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 7a7991a

17 files changed

Lines changed: 2258 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: 209 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,8 @@
1+
import inspect
2+
import json
3+
import logging
14
from abc import ABC, abstractmethod
5+
from typing import Any, Awaitable, Callable, Optional
26

37

48
class BaseAiHandler(ABC):
@@ -10,6 +14,8 @@ class BaseAiHandler(ABC):
1014
def __init__(self):
1115
pass
1216

17+
_logger = logging.getLogger(__name__)
18+
1319
@property
1420
@abstractmethod
1521
def deployment_id(self):
@@ -26,3 +32,206 @@ async def chat_completion(self, model: str, system: str, user: str, temperature:
2632
temperature (float): the temperature to use for the chat completion
2733
"""
2834
pass
35+
36+
async def chat_completion_with_tools(
37+
self,
38+
model: str,
39+
system: str,
40+
user: str,
41+
tools: Optional[list[dict[str, Any]]] = None,
42+
tool_executor: Optional[Callable[[str, dict[str, Any]], Any | Awaitable[Any]]] = None,
43+
temperature: float = 0.2,
44+
img_path: str = None,
45+
max_tool_turns: int = 4,
46+
max_tool_output_chars: int = 12000,
47+
):
48+
"""
49+
Run a structured tool-calling loop on top of plain chat completion.
50+
51+
The model is instructed to emit JSON tool requests in the form:
52+
{"type": "tool_call", "tool": "server.tool", "arguments": {...}}
53+
and to finish with:
54+
{"type": "final", "content": "..."}
55+
56+
max_tool_output_chars is applied per tool call, not across all tool calls.
57+
"""
58+
if not tools or tool_executor is None:
59+
return await self.chat_completion(model, system, user, temperature=temperature, img_path=img_path)
60+
61+
allowed_tool_names = self._extract_allowed_tool_names(tools)
62+
63+
tool_catalog_text = json.dumps(tools, indent=2, sort_keys=True)
64+
structured_system = (
65+
f"{system}\n\n"
66+
f"Available MCP tools (JSON schema):\n{tool_catalog_text}\n\n"
67+
"Always inspect the available tools first and use them before responding "
68+
"whenever they can help answer the user's request.\n"
69+
"When you need a tool, respond with ONLY a JSON object exactly in this shape:\n"
70+
'{"type":"tool_call","tool":"server.tool","arguments":{...}}\n'
71+
"Do not include a final answer in the same message as a tool call.\n"
72+
"When you are finished, respond with ONLY a JSON object exactly in this shape:\n"
73+
'{"type":"final","content":"..."}\n'
74+
"Do not wrap the JSON in markdown fences."
75+
)
76+
77+
conversation_history = [user]
78+
remaining_turns = max_tool_turns
79+
current_img_path = img_path
80+
81+
while True:
82+
current_user = "\n\n".join(conversation_history)
83+
response_text, finish_reason = await self.chat_completion(
84+
model=model,
85+
system=structured_system,
86+
user=current_user,
87+
temperature=temperature,
88+
img_path=current_img_path,
89+
)
90+
current_img_path = None
91+
92+
parsed_response = self._parse_tool_or_final_response(response_text)
93+
if parsed_response is None:
94+
return response_text, finish_reason
95+
96+
response_type = parsed_response.get("type", "final")
97+
if response_type == "final":
98+
return str(parsed_response.get("content", "")), finish_reason
99+
100+
if response_type != "tool_call":
101+
return response_text, finish_reason
102+
103+
if remaining_turns <= 0:
104+
self._logger.warning("MCP tool orchestration exceeded the configured turn budget")
105+
return response_text, finish_reason
106+
107+
tool_name = str(parsed_response.get("tool", "")).strip()
108+
arguments = parsed_response.get("arguments") or {}
109+
if not tool_name:
110+
self._logger.warning("MCP tool orchestration returned an empty tool name; aborting tool loop")
111+
return response_text, finish_reason
112+
if not isinstance(arguments, dict):
113+
self._logger.warning("MCP tool orchestration arguments must be a JSON object; aborting tool loop")
114+
return response_text, finish_reason
115+
116+
if tool_name not in allowed_tool_names:
117+
self._logger.warning("MCP tool '%s' was not in the advertised tool catalog; skipping", tool_name)
118+
tool_result = f"Tool not available: {tool_name}"
119+
else:
120+
try:
121+
tool_result = tool_executor(tool_name, arguments)
122+
if inspect.isawaitable(tool_result):
123+
tool_result = await tool_result
124+
except Exception as exc: # noqa: BLE001
125+
self._logger.warning("MCP tool '%s' raised an exception: %s", tool_name, exc)
126+
tool_result = f"Tool error: {exc}"
127+
128+
tool_result_text = self._normalize_tool_result_text(
129+
tool_result,
130+
max_tool_output_chars=max_tool_output_chars,
131+
tool_name=tool_name,
132+
)
133+
conversation_history.append(f"Previous assistant tool request:\n{response_text}")
134+
conversation_history.append(f"Tool result for {tool_name}:\n{tool_result_text}")
135+
remaining_turns -= 1
136+
137+
@classmethod
138+
def _normalize_tool_result_text(
139+
cls,
140+
tool_result: Any,
141+
max_tool_output_chars: int,
142+
tool_name: str = "<unknown>",
143+
) -> str:
144+
if isinstance(tool_result, str):
145+
result_text = tool_result
146+
else:
147+
result_text = json.dumps(tool_result, indent=2, sort_keys=True, default=str)
148+
149+
if len(result_text) > max_tool_output_chars:
150+
cls._logger.warning(
151+
"Tool output for '%s' exceeded per-tool max_tool_output_chars (%s > %s); truncating output",
152+
tool_name,
153+
len(result_text),
154+
max_tool_output_chars,
155+
)
156+
if max_tool_output_chars <= 0:
157+
return ""
158+
suffix = "\n[tool output truncated]"
159+
if max_tool_output_chars <= len(suffix):
160+
return suffix[:max_tool_output_chars]
161+
truncated_prefix_len = max(0, max_tool_output_chars - len(suffix))
162+
return result_text[:truncated_prefix_len] + suffix
163+
return result_text
164+
165+
@staticmethod
166+
def _parse_tool_or_final_response(response_text: str) -> Optional[dict[str, Any]]:
167+
candidate = response_text.strip()
168+
if not candidate:
169+
return None
170+
171+
for json_candidate in BaseAiHandler._iter_json_object_candidates(candidate):
172+
try:
173+
parsed = json.loads(json_candidate)
174+
except json.JSONDecodeError:
175+
continue
176+
177+
if isinstance(parsed, dict):
178+
response_type = parsed.get("type")
179+
if response_type in {"tool_call", "final"}:
180+
return parsed
181+
182+
return None
183+
184+
@staticmethod
185+
def _iter_json_object_candidates(text: str) -> list[str]:
186+
candidates: list[str] = []
187+
depth = 0
188+
start_index: Optional[int] = None
189+
in_string = False
190+
is_escaped = False
191+
192+
for index, char in enumerate(text):
193+
if in_string:
194+
if is_escaped:
195+
is_escaped = False
196+
elif char == "\\":
197+
is_escaped = True
198+
elif char == '"':
199+
in_string = False
200+
continue
201+
202+
if char == '"':
203+
in_string = True
204+
continue
205+
206+
if char == "{":
207+
if depth == 0:
208+
start_index = index
209+
depth += 1
210+
continue
211+
212+
if char == "}" and depth > 0:
213+
depth -= 1
214+
if depth == 0 and start_index is not None:
215+
candidates.append(text[start_index : index + 1])
216+
start_index = None
217+
218+
return candidates
219+
220+
@staticmethod
221+
def _extract_allowed_tool_names(tools: list[dict[str, Any]]) -> set[str]:
222+
allowed: set[str] = set()
223+
for tool in tools:
224+
if not isinstance(tool, dict):
225+
continue
226+
227+
function_info = tool.get("function")
228+
if isinstance(function_info, dict):
229+
function_name = function_info.get("name")
230+
if isinstance(function_name, str) and function_name.strip():
231+
allowed.add(function_name.strip())
232+
233+
simple_name = tool.get("name")
234+
if isinstance(simple_name, str) and simple_name.strip():
235+
allowed.add(simple_name.strip())
236+
237+
return allowed

0 commit comments

Comments
 (0)