Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 46 additions & 3 deletions agentrun/tool/__tool_async_template.py
Original file line number Diff line number Diff line change
Expand Up @@ -261,6 +261,40 @@ def _parse_protocol_spec_mcp_url(self) -> Tuple[str, str, Dict[str, str]]:

return url, session_affinity, spec_headers

def _infer_protocol_spec_mcp_session_affinity(self) -> Optional[str]:
"""从 protocol_spec 推断 MCP session_affinity / Infer MCP session_affinity from protocol_spec

用于 MCP_REMOTE + proxy_enabled=true 且 mcp_config.session_affinity
为空的场景。proxy 模式仍使用数据面 URL,不使用 protocol_spec 中的
上游 URL 和 headers。
Used when MCP_REMOTE proxy is enabled but mcp_config.session_affinity
is empty. Proxy mode still uses data endpoint URL, not upstream URL
or headers from protocol_spec.

Returns:
Optional[str]: MCP_STREAMABLE、MCP_SSE 或 None
"""
if not self.protocol_spec:
return None

try:
spec = json.loads(self.protocol_spec)
except (json.JSONDecodeError, TypeError):
return None

mcp_servers = spec.get("mcpServers")
if not mcp_servers or not isinstance(mcp_servers, dict):
return None

first_server = next(iter(mcp_servers.values()), None)
if not first_server or not isinstance(first_server, dict):
return None

transport_type = first_server.get("transportType", "sse")
if transport_type == "streamable-http":
return "MCP_STREAMABLE"
return "MCP_SSE"

def _get_mcp_endpoint(
self, config: Optional[Config] = None
) -> Optional[Tuple[str, str, Dict[str, str]]]:
Expand Down Expand Up @@ -290,9 +324,18 @@ def _get_mcp_endpoint(
if not data_endpoint or not effective_name:
return None

session_affinity = pydash.get(
self, "mcp_config.session_affinity", "MCP_SSE"
)
session_affinity = pydash.get(self, "mcp_config.session_affinity")
if not session_affinity:
is_mcp_remote_with_proxy = (
self.create_method == "MCP_REMOTE"
and pydash.get(self, "mcp_config.proxy_enabled", False)
)
if is_mcp_remote_with_proxy:
session_affinity = (
self._infer_protocol_spec_mcp_session_affinity()
)
if not session_affinity:
session_affinity = "MCP_SSE"

if session_affinity == "MCP_STREAMABLE":
return (
Expand Down
49 changes: 46 additions & 3 deletions agentrun/tool/tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -286,6 +286,40 @@ def _parse_protocol_spec_mcp_url(self) -> Tuple[str, str, Dict[str, str]]:

return url, session_affinity, spec_headers

def _infer_protocol_spec_mcp_session_affinity(self) -> Optional[str]:
"""从 protocol_spec 推断 MCP session_affinity / Infer MCP session_affinity from protocol_spec

用于 MCP_REMOTE + proxy_enabled=true 且 mcp_config.session_affinity
为空的场景。proxy 模式仍使用数据面 URL,不使用 protocol_spec 中的
上游 URL 和 headers。
Used when MCP_REMOTE proxy is enabled but mcp_config.session_affinity
is empty. Proxy mode still uses data endpoint URL, not upstream URL
or headers from protocol_spec.

Returns:
Optional[str]: MCP_STREAMABLE、MCP_SSE 或 None
"""
if not self.protocol_spec:
return None

try:
spec = json.loads(self.protocol_spec)
except (json.JSONDecodeError, TypeError):
return None

mcp_servers = spec.get("mcpServers")
if not mcp_servers or not isinstance(mcp_servers, dict):
return None

first_server = next(iter(mcp_servers.values()), None)
if not first_server or not isinstance(first_server, dict):
return None

transport_type = first_server.get("transportType", "sse")
if transport_type == "streamable-http":
return "MCP_STREAMABLE"
return "MCP_SSE"

def _get_mcp_endpoint(
self, config: Optional[Config] = None
) -> Optional[Tuple[str, str, Dict[str, str]]]:
Expand Down Expand Up @@ -315,9 +349,18 @@ def _get_mcp_endpoint(
if not data_endpoint or not effective_name:
return None

session_affinity = pydash.get(
self, "mcp_config.session_affinity", "MCP_SSE"
)
session_affinity = pydash.get(self, "mcp_config.session_affinity")
if not session_affinity:
is_mcp_remote_with_proxy = (
self.create_method == "MCP_REMOTE"
and pydash.get(self, "mcp_config.proxy_enabled", False)
)
if is_mcp_remote_with_proxy:
session_affinity = (
self._infer_protocol_spec_mcp_session_affinity()
)
if not session_affinity:
session_affinity = "MCP_SSE"

if session_affinity == "MCP_STREAMABLE":
return (
Expand Down
36 changes: 36 additions & 0 deletions tests/unittests/tool/test_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -1161,6 +1161,42 @@ def test_get_mcp_endpoint_mcp_remote_with_proxy_uses_data_endpoint(self):
{},
)

def test_get_mcp_endpoint_mcp_remote_with_proxy_infers_streamable(self):
"""测试 MCP_REMOTE proxy 模式按 protocol_spec 推断 streamable。"""
tool = Tool(
tool_name="my-tool",
tool_type="MCP",
create_method="MCP_REMOTE",
data_endpoint="https://example.com",
mcp_config=McpConfig(proxy_enabled=True),
protocol_spec='{"mcpServers":{"s1":{"transportType":"streamable-http","url":"https://external-mcp.com/mcp"}}}',
)
result = tool._get_mcp_endpoint()
assert result == (
"https://example.com/tools/my-tool/mcp",
"MCP_STREAMABLE",
{},
)

def test_get_mcp_endpoint_mcp_remote_with_proxy_empty_affinity_infers_streamable(
self,
):
"""测试空 session_affinity 也按 protocol_spec 推断 streamable。"""
tool = Tool(
tool_name="my-tool",
tool_type="MCP",
create_method="MCP_REMOTE",
data_endpoint="https://example.com",
mcp_config=McpConfig(session_affinity="", proxy_enabled=True),
protocol_spec='{"mcpServers":{"s1":{"transportType":"streamable-http","url":"https://external-mcp.com/mcp"}}}',
)
result = tool._get_mcp_endpoint()
assert result == (
"https://example.com/tools/my-tool/mcp",
"MCP_STREAMABLE",
{},
)

def test_get_mcp_endpoint_mcp_bundle_uses_data_endpoint(self):
"""测试 MCP_BUNDLE 类型使用 data_endpoint 拼接"""
tool = Tool(
Expand Down