Skip to content

Commit 39b32f8

Browse files
SodawyxSodawyx
committed
refactor(tool): update MCP endpoint handling and add protocol spec parsing
This change updates the `_get_mcp_endpoint` method to return both the endpoint URL and session affinity as a tuple, and introduces a new method `_parse_protocol_spec_mcp_url` to handle parsing of MCP URLs from protocol specifications when using `MCP_REMOTE` without proxy enabled. This improves consistency across synchronous and asynchronous implementations of the tool class. Co-developed-by: Aone Copilot <noreply@alibaba-inc.com> Signed-off-by: Sodawyx <sodawyx@126.com>
1 parent 7124dd3 commit 39b32f8

File tree

3 files changed

+381
-54
lines changed

3 files changed

+381
-54
lines changed

‎agentrun/tool/__tool_async_template.py‎

Lines changed: 88 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,10 @@
55
"""
66

77
import io
8+
import json
89
import os
910
import shutil
10-
from typing import Any, Dict, List, Optional
11+
from typing import Any, Dict, List, Optional, Tuple
1112
from urllib.parse import urlparse
1213
import zipfile
1314

@@ -196,16 +197,85 @@ def _get_tool_type(self) -> Optional[ToolType]:
196197
return None
197198
return None
198199

200+
def _parse_protocol_spec_mcp_url(self) -> Tuple[str, str]:
201+
"""从 protocol_spec 解析 MCP 服务器 URL 和 session_affinity / Parse MCP server URL and session_affinity from protocol_spec
202+
203+
用于 MCP_REMOTE + proxy_enabled=false 场景，从 protocol_spec JSON 中提取
204+
第一个 mcpServers entry 的 url 和 transportType。
205+
Used for MCP_REMOTE + proxy_enabled=false scenario, extracts url and
206+
transportType from the first mcpServers entry in protocol_spec JSON.
207+
208+
Returns:
209+
Tuple[str, str]: (mcp_url, session_affinity)
210+
211+
Raises:
212+
ValueError: protocol_spec 为空、格式不合法或缺少必要字段时抛出
213+
"""
214+
if not self.protocol_spec:
215+
raise ValueError(
216+
"protocol_spec is required for MCP_REMOTE tool with proxy"
217+
" disabled, but it is empty for tool"
218+
f" '{self.tool_name or self.name}'"
219+
)
220+
221+
try:
222+
spec = json.loads(self.protocol_spec)
223+
except (json.JSONDecodeError, TypeError) as exc:
224+
raise ValueError(
225+
"Failed to parse protocol_spec for tool"
226+
f" '{self.tool_name or self.name}': {exc}"
227+
) from exc
228+
229+
mcp_servers = spec.get("mcpServers")
230+
if not mcp_servers or not isinstance(mcp_servers, dict):
231+
raise ValueError(
232+
"mcpServers not found or invalid in protocol_spec for tool"
233+
f" '{self.tool_name or self.name}'"
234+
)
235+
236+
first_server = next(iter(mcp_servers.values()), None)
237+
if not first_server or not isinstance(first_server, dict):
238+
raise ValueError(
239+
"No MCP server entry found in protocol_spec for tool"
240+
f" '{self.tool_name or self.name}'"
241+
)
242+
243+
url = first_server.get("url")
244+
if not url:
245+
raise ValueError(
246+
"url not found in MCP server entry of protocol_spec for tool"
247+
f" '{self.tool_name or self.name}'"
248+
)
249+
250+
transport_type = first_server.get("transportType", "sse")
251+
if transport_type == "streamable-http":
252+
session_affinity = "MCP_STREAMABLE"
253+
else:
254+
session_affinity = "MCP_SSE"
255+
256+
return url, session_affinity
257+
199258
def _get_mcp_endpoint(
200259
self, config: Optional[Config] = None
201-
) -> Optional[str]:
202-
"""获取 MCP 数据链路 URL / Get MCP data endpoint URL
260+
) -> Optional[Tuple[str, str]]:
261+
"""获取 MCP 数据链路 URL 和 session_affinity / Get MCP data endpoint URL and session_affinity
262+
263+
MCP_REMOTE + proxy_enabled=false 时从 protocol_spec 解析 URL 和 session_affinity。
264+
其他场景使用 data_endpoint 拼接，session_affinity 从 mcp_config 获取。
265+
For MCP_REMOTE with proxy disabled, parses URL and session_affinity from protocol_spec.
266+
Otherwise constructs URL from data_endpoint and gets session_affinity from mcp_config.
203267
204-
根据 session_affinity 决定使用 /mcp 还是 /sse 路径。
205-
如果 self.data_endpoint 为空，则从 Config 中获取。
206-
Determines /mcp or /sse path based on session_affinity.
207-
Falls back to Config if self.data_endpoint is not set.
268+
Returns:
269+
Optional[Tuple[str, str]]: (endpoint_url, session_affinity) 或 None
208270
"""
271+
is_mcp_remote_without_proxy = (
272+
self.create_method == "MCP_REMOTE"
273+
and not pydash.get(self, "mcp_config.proxy_enabled", False)
274+
)
275+
276+
if is_mcp_remote_without_proxy:
277+
return self._parse_protocol_spec_mcp_url()
278+
209279
effective_name = self.tool_name or self.name
210280
data_endpoint = self.data_endpoint
211281
if not data_endpoint:
@@ -219,8 +289,11 @@ def _get_mcp_endpoint(
219289
)
220290

221291
if session_affinity == "MCP_STREAMABLE":
222-
return f"{data_endpoint}/tools/{effective_name}/mcp"
223-
return f"{data_endpoint}/tools/{effective_name}/sse"
292+
return (
293+
f"{data_endpoint}/tools/{effective_name}/mcp",
294+
session_affinity,
295+
)
296+
return f"{data_endpoint}/tools/{effective_name}/sse", session_affinity
224297

225298
async def list_tools_async(
226299
self, config: Optional[Config] = None
@@ -240,16 +313,14 @@ async def list_tools_async(
240313
if tool_type == ToolType.MCP:
241314
from .api.mcp import ToolMCPSession
242315

243-
mcp_endpoint = self._get_mcp_endpoint(config)
244-
if not mcp_endpoint:
316+
endpoint_result = self._get_mcp_endpoint(config)
317+
if not endpoint_result:
245318
logger.warning(
246319
"MCP endpoint not available for tool %s", self.name
247320
)
248321
return []
249322

250-
session_affinity = pydash.get(
251-
self, "mcp_config.session_affinity", "MCP_SSE"
252-
)
323+
mcp_endpoint, session_affinity = endpoint_result
253324

254325
# MCP_REMOTE + proxy_enabled=false 时直连外部服务，不走 RAM 鉴权
255326
# Only skip RAM auth for MCP_REMOTE with proxy disabled (direct external connection)
@@ -309,15 +380,13 @@ async def call_tool_async(
309380
if tool_type == ToolType.MCP:
310381
from .api.mcp import ToolMCPSession
311382

312-
mcp_endpoint = self._get_mcp_endpoint(config)
313-
if not mcp_endpoint:
383+
endpoint_result = self._get_mcp_endpoint(config)
384+
if not endpoint_result:
314385
raise ValueError(
315386
f"MCP endpoint not available for tool {self.name}"
316387
)
317388

318-
session_affinity = pydash.get(
319-
self, "mcp_config.session_affinity", "MCP_SSE"
320-
)
389+
mcp_endpoint, session_affinity = endpoint_result
321390

322391
# MCP_REMOTE + proxy_enabled=false 时直连外部服务，不走 RAM 鉴权
323392
# Only skip RAM auth for MCP_REMOTE with proxy disabled (direct external connection)

‎agentrun/tool/tool.py‎

Lines changed: 94 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -15,9 +15,10 @@
1515
"""
1616

1717
import io
18+
import json
1819
import os
1920
import shutil
20-
from typing import Any, Dict, List, Optional
21+
from typing import Any, Dict, List, Optional, Tuple
2122
from urllib.parse import urlparse
2223
import zipfile
2324

@@ -221,16 +222,85 @@ def _get_tool_type(self) -> Optional[ToolType]:
221222
return None
222223
return None
223224

225+
def _parse_protocol_spec_mcp_url(self) -> Tuple[str, str]:
226+
"""从 protocol_spec 解析 MCP 服务器 URL 和 session_affinity / Parse MCP server URL and session_affinity from protocol_spec
227+
228+
用于 MCP_REMOTE + proxy_enabled=false 场景，从 protocol_spec JSON 中提取
229+
第一个 mcpServers entry 的 url 和 transportType。
230+
Used for MCP_REMOTE + proxy_enabled=false scenario, extracts url and
231+
transportType from the first mcpServers entry in protocol_spec JSON.
232+
233+
Returns:
234+
Tuple[str, str]: (mcp_url, session_affinity)
235+
236+
Raises:
237+
ValueError: protocol_spec 为空、格式不合法或缺少必要字段时抛出
238+
"""
239+
if not self.protocol_spec:
240+
raise ValueError(
241+
"protocol_spec is required for MCP_REMOTE tool with proxy"
242+
" disabled, but it is empty for tool"
243+
f" '{self.tool_name or self.name}'"
244+
)
245+
246+
try:
247+
spec = json.loads(self.protocol_spec)
248+
except (json.JSONDecodeError, TypeError) as exc:
249+
raise ValueError(
250+
"Failed to parse protocol_spec for tool"
251+
f" '{self.tool_name or self.name}': {exc}"
252+
) from exc
253+
254+
mcp_servers = spec.get("mcpServers")
255+
if not mcp_servers or not isinstance(mcp_servers, dict):
256+
raise ValueError(
257+
"mcpServers not found or invalid in protocol_spec for tool"
258+
f" '{self.tool_name or self.name}'"
259+
)
260+
261+
first_server = next(iter(mcp_servers.values()), None)
262+
if not first_server or not isinstance(first_server, dict):
263+
raise ValueError(
264+
"No MCP server entry found in protocol_spec for tool"
265+
f" '{self.tool_name or self.name}'"
266+
)
267+
268+
url = first_server.get("url")
269+
if not url:
270+
raise ValueError(
271+
"url not found in MCP server entry of protocol_spec for tool"
272+
f" '{self.tool_name or self.name}'"
273+
)
274+
275+
transport_type = first_server.get("transportType", "sse")
276+
if transport_type == "streamable-http":
277+
session_affinity = "MCP_STREAMABLE"
278+
else:
279+
session_affinity = "MCP_SSE"
280+
281+
return url, session_affinity
282+
224283
def _get_mcp_endpoint(
225284
self, config: Optional[Config] = None
226-
) -> Optional[str]:
227-
"""获取 MCP 数据链路 URL / Get MCP data endpoint URL
285+
) -> Optional[Tuple[str, str]]:
286+
"""获取 MCP 数据链路 URL 和 session_affinity / Get MCP data endpoint URL and session_affinity
228287
229-
根据 session_affinity 决定使用 /mcp 还是 /sse 路径。
230-
如果 self.data_endpoint 为空，则从 Config 中获取。
231-
Determines /mcp or /sse path based on session_affinity.
232-
Falls back to Config if self.data_endpoint is not set.
288+
MCP_REMOTE + proxy_enabled=false 时从 protocol_spec 解析 URL 和 session_affinity。
289+
其他场景使用 data_endpoint 拼接，session_affinity 从 mcp_config 获取。
290+
For MCP_REMOTE with proxy disabled, parses URL and session_affinity from protocol_spec.
291+
Otherwise constructs URL from data_endpoint and gets session_affinity from mcp_config.
292+
293+
Returns:
294+
Optional[Tuple[str, str]]: (endpoint_url, session_affinity) 或 None
233295
"""
296+
is_mcp_remote_without_proxy = (
297+
self.create_method == "MCP_REMOTE"
298+
and not pydash.get(self, "mcp_config.proxy_enabled", False)
299+
)
300+
301+
if is_mcp_remote_without_proxy:
302+
return self._parse_protocol_spec_mcp_url()
303+
234304
effective_name = self.tool_name or self.name
235305
data_endpoint = self.data_endpoint
236306
if not data_endpoint:
@@ -244,8 +314,11 @@ def _get_mcp_endpoint(
244314
)
245315

246316
if session_affinity == "MCP_STREAMABLE":
247-
return f"{data_endpoint}/tools/{effective_name}/mcp"
248-
return f"{data_endpoint}/tools/{effective_name}/sse"
317+
return (
318+
f"{data_endpoint}/tools/{effective_name}/mcp",
319+
session_affinity,
320+
)
321+
return f"{data_endpoint}/tools/{effective_name}/sse", session_affinity
249322

250323
async def list_tools_async(
251324
self, config: Optional[Config] = None
@@ -265,16 +338,14 @@ async def list_tools_async(
265338
if tool_type == ToolType.MCP:
266339
from .api.mcp import ToolMCPSession
267340

268-
mcp_endpoint = self._get_mcp_endpoint(config)
269-
if not mcp_endpoint:
341+
endpoint_result = self._get_mcp_endpoint(config)
342+
if not endpoint_result:
270343
logger.warning(
271344
"MCP endpoint not available for tool %s", self.name
272345
)
273346
return []
274347

275-
session_affinity = pydash.get(
276-
self, "mcp_config.session_affinity", "MCP_SSE"
277-
)
348+
mcp_endpoint, session_affinity = endpoint_result
278349

279350
# MCP_REMOTE + proxy_enabled=false 时直连外部服务，不走 RAM 鉴权
280351
# Only skip RAM auth for MCP_REMOTE with proxy disabled (direct external connection)
@@ -327,16 +398,14 @@ def list_tools(self, config: Optional[Config] = None) -> List[ToolInfo]:
327398
if tool_type == ToolType.MCP:
328399
from .api.mcp import ToolMCPSession
329400

330-
mcp_endpoint = self._get_mcp_endpoint(config)
331-
if not mcp_endpoint:
401+
endpoint_result = self._get_mcp_endpoint(config)
402+
if not endpoint_result:
332403
logger.warning(
333404
"MCP endpoint not available for tool %s", self.name
334405
)
335406
return []
336407

337-
session_affinity = pydash.get(
338-
self, "mcp_config.session_affinity", "MCP_SSE"
339-
)
408+
mcp_endpoint, session_affinity = endpoint_result
340409

341410
# MCP_REMOTE + proxy_enabled=false 时直连外部服务，不走 RAM 鉴权
342411
# Only skip RAM auth for MCP_REMOTE with proxy disabled (direct external connection)
@@ -396,15 +465,13 @@ async def call_tool_async(
396465
if tool_type == ToolType.MCP:
397466
from .api.mcp import ToolMCPSession
398467

399-
mcp_endpoint = self._get_mcp_endpoint(config)
400-
if not mcp_endpoint:
468+
endpoint_result = self._get_mcp_endpoint(config)
469+
if not endpoint_result:
401470
raise ValueError(
402471
f"MCP endpoint not available for tool {self.name}"
403472
)
404473

405-
session_affinity = pydash.get(
406-
self, "mcp_config.session_affinity", "MCP_SSE"
407-
)
474+
mcp_endpoint, session_affinity = endpoint_result
408475

409476
# MCP_REMOTE + proxy_enabled=false 时直连外部服务，不走 RAM 鉴权
410477
# Only skip RAM auth for MCP_REMOTE with proxy disabled (direct external connection)
@@ -469,15 +536,13 @@ def call_tool(
469536
if tool_type == ToolType.MCP:
470537
from .api.mcp import ToolMCPSession
471538

472-
mcp_endpoint = self._get_mcp_endpoint(config)
473-
if not mcp_endpoint:
539+
endpoint_result = self._get_mcp_endpoint(config)
540+
if not endpoint_result:
474541
raise ValueError(
475542
f"MCP endpoint not available for tool {self.name}"
476543
)
477544

478-
session_affinity = pydash.get(
479-
self, "mcp_config.session_affinity", "MCP_SSE"
480-
)
545+
mcp_endpoint, session_affinity = endpoint_result
481546

482547
# MCP_REMOTE + proxy_enabled=false 时直连外部服务，不走 RAM 鉴权
483548
# Only skip RAM auth for MCP_REMOTE with proxy disabled (direct external connection)

0 commit comments

Comments
 (0)