From 2062b705ccad6048a99771435d781a0850268c1e Mon Sep 17 00:00:00 2001 From: 117503445 Date: Sat, 30 May 2026 02:24:26 +0800 Subject: [PATCH] =?UTF-8?q?feat(tool):=20=E6=8E=A8=E6=96=AD=20MCP=20sessio?= =?UTF-8?q?n=5Faffinity=20=E4=BB=A5=E6=94=AF=E6=8C=81=20MCP=5FREMOTE=20pro?= =?UTF-8?q?xy=20=E6=A8=A1=E5=BC=8F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 黑曜 --- agentrun/tool/__tool_async_template.py | 49 ++++++++++++++++++++++++-- agentrun/tool/tool.py | 49 ++++++++++++++++++++++++-- tests/unittests/tool/test_tool.py | 36 +++++++++++++++++++ 3 files changed, 128 insertions(+), 6 deletions(-) diff --git a/agentrun/tool/__tool_async_template.py b/agentrun/tool/__tool_async_template.py index 3833c48..fc57c5d 100644 --- a/agentrun/tool/__tool_async_template.py +++ b/agentrun/tool/__tool_async_template.py @@ -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]]]: @@ -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 ( diff --git a/agentrun/tool/tool.py b/agentrun/tool/tool.py index 2414585..d0b7dcf 100644 --- a/agentrun/tool/tool.py +++ b/agentrun/tool/tool.py @@ -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]]]: @@ -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 ( diff --git a/tests/unittests/tool/test_tool.py b/tests/unittests/tool/test_tool.py index adcaa72..5ffc354 100644 --- a/tests/unittests/tool/test_tool.py +++ b/tests/unittests/tool/test_tool.py @@ -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(