Skip to content

Commit 288e518

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 288e518

17 files changed

Lines changed: 2320 additions & 16 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: 221 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,218 @@ 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+
tool_call_example = json.dumps(
63+
{
64+
"type": "tool_call",
65+
"tool": "server.tool",
66+
"arguments": {"example": "value"},
67+
},
68+
separators=(",", ":"),
69+
).replace("{\"example\":\"value\"}", "{...}")
70+
final_response_example = json.dumps(
71+
{"type": "final", "content": "..."},
72+
separators=(",", ":"),
73+
)
74+
75+
tool_catalog_text = json.dumps(tools, indent=2, sort_keys=True)
76+
structured_system = (
77+
f"{system}\n\n"
78+
f"Available MCP tools (JSON schema):\n{tool_catalog_text}\n\n"
79+
"Always inspect the available tools first and use them before responding "
80+
"whenever they can help answer the user's request.\n"
81+
"When you need a tool, respond with ONLY a JSON object exactly in this shape:\n"
82+
f"{tool_call_example}\n"
83+
"Do not include a final answer in the same message as a tool call.\n"
84+
"When you are finished, respond with ONLY a JSON object exactly in this shape:\n"
85+
f"{final_response_example}\n"
86+
"Do not wrap the JSON in markdown fences."
87+
)
88+
89+
conversation_history = [user]
90+
remaining_turns = max_tool_turns
91+
current_img_path = img_path
92+
93+
while True:
94+
current_user = "\n\n".join(conversation_history)
95+
response_text, finish_reason = await self.chat_completion(
96+
model=model,
97+
system=structured_system,
98+
user=current_user,
99+
temperature=temperature,
100+
img_path=current_img_path,
101+
)
102+
current_img_path = None
103+
104+
parsed_response = self._parse_tool_or_final_response(response_text)
105+
if parsed_response is None:
106+
return response_text, finish_reason
107+
108+
response_type = parsed_response.get("type", "final")
109+
if response_type == "final":
110+
return str(parsed_response.get("content", "")), finish_reason
111+
112+
if response_type != "tool_call":
113+
return response_text, finish_reason
114+
115+
if remaining_turns <= 0:
116+
self._logger.warning("MCP tool orchestration exceeded the configured turn budget")
117+
return response_text, finish_reason
118+
119+
tool_name = str(parsed_response.get("tool", "")).strip()
120+
arguments = parsed_response.get("arguments") or {}
121+
if not tool_name:
122+
self._logger.warning("MCP tool orchestration returned an empty tool name; aborting tool loop")
123+
return response_text, finish_reason
124+
if not isinstance(arguments, dict):
125+
self._logger.warning("MCP tool orchestration arguments must be a JSON object; aborting tool loop")
126+
return response_text, finish_reason
127+
128+
if tool_name not in allowed_tool_names:
129+
self._logger.warning("MCP tool '%s' was not in the advertised tool catalog; skipping", tool_name)
130+
tool_result = f"Tool not available: {tool_name}"
131+
else:
132+
try:
133+
tool_result = tool_executor(tool_name, arguments)
134+
if inspect.isawaitable(tool_result):
135+
tool_result = await tool_result
136+
except Exception as exc: # noqa: BLE001
137+
self._logger.warning("MCP tool '%s' raised an exception: %s", tool_name, exc)
138+
tool_result = f"Tool error: {exc}"
139+
140+
tool_result_text = self._normalize_tool_result_text(
141+
tool_result,
142+
max_tool_output_chars=max_tool_output_chars,
143+
tool_name=tool_name,
144+
)
145+
conversation_history.append(f"Previous assistant tool request:\n{response_text}")
146+
conversation_history.append(f"Tool result for {tool_name}:\n{tool_result_text}")
147+
remaining_turns -= 1
148+
149+
@classmethod
150+
def _normalize_tool_result_text(
151+
cls,
152+
tool_result: Any,
153+
max_tool_output_chars: int,
154+
tool_name: str = "<unknown>",
155+
) -> str:
156+
if isinstance(tool_result, str):
157+
result_text = tool_result
158+
else:
159+
result_text = json.dumps(tool_result, indent=2, sort_keys=True, default=str)
160+
161+
if len(result_text) > max_tool_output_chars:
162+
cls._logger.warning(
163+
"Tool output for '%s' exceeded per-tool max_tool_output_chars (%s > %s); truncating output",
164+
tool_name,
165+
len(result_text),
166+
max_tool_output_chars,
167+
)
168+
if max_tool_output_chars <= 0:
169+
return ""
170+
suffix = "\n[tool output truncated]"
171+
if max_tool_output_chars <= len(suffix):
172+
return suffix[:max_tool_output_chars]
173+
truncated_prefix_len = max(0, max_tool_output_chars - len(suffix))
174+
return result_text[:truncated_prefix_len] + suffix
175+
return result_text
176+
177+
@staticmethod
178+
def _parse_tool_or_final_response(response_text: str) -> Optional[dict[str, Any]]:
179+
candidate = response_text.strip()
180+
if not candidate:
181+
return None
182+
183+
for json_candidate in BaseAiHandler._iter_json_object_candidates(candidate):
184+
try:
185+
parsed = json.loads(json_candidate)
186+
except json.JSONDecodeError:
187+
continue
188+
189+
if isinstance(parsed, dict):
190+
response_type = parsed.get("type")
191+
if response_type in {"tool_call", "final"}:
192+
return parsed
193+
194+
return None
195+
196+
@staticmethod
197+
def _iter_json_object_candidates(text: str) -> list[str]:
198+
candidates: list[str] = []
199+
depth = 0
200+
start_index: Optional[int] = None
201+
in_string = False
202+
is_escaped = False
203+
204+
for index, char in enumerate(text):
205+
if in_string:
206+
if is_escaped:
207+
is_escaped = False
208+
elif char == "\\":
209+
is_escaped = True
210+
elif char == '"':
211+
in_string = False
212+
continue
213+
214+
if char == '"':
215+
in_string = True
216+
continue
217+
218+
if char == "{":
219+
if depth == 0:
220+
start_index = index
221+
depth += 1
222+
continue
223+
224+
if char == "}" and depth > 0:
225+
depth -= 1
226+
if depth == 0 and start_index is not None:
227+
candidates.append(text[start_index : index + 1])
228+
start_index = None
229+
230+
return candidates
231+
232+
@staticmethod
233+
def _extract_allowed_tool_names(tools: list[dict[str, Any]]) -> set[str]:
234+
allowed: set[str] = set()
235+
for tool in tools:
236+
if not isinstance(tool, dict):
237+
continue
238+
239+
function_info = tool.get("function")
240+
if isinstance(function_info, dict):
241+
function_name = function_info.get("name")
242+
if isinstance(function_name, str) and function_name.strip():
243+
allowed.add(function_name.strip())
244+
245+
simple_name = tool.get("name")
246+
if isinstance(simple_name, str) and simple_name.strip():
247+
allowed.add(simple_name.strip())
248+
249+
return allowed

0 commit comments

Comments
 (0)