Merge pull request #111 from 117503445/fix/mcp-remote-proxy-streamable-endpoint

feat(tool): 推断 MCP session_affinity 以支持 MCP_REMOTE proxy 模式

Author: OhYee

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 (

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 (

tests/unittests/tool/test_tool.py

@@ -1161,6 +1161,77 @@ def test_get_mcp_endpoint_mcp_remote_with_proxy_uses_data_endpoint(self):
             {},
         )
 
+    @pytest.mark.parametrize(
+        "protocol_spec",
+        [
+            None,
+            "invalid json",
+            "{}",
+            '{"mcpServers":{}}',
+            '{"mcpServers":[]}',
+            '{"mcpServers":{"s1":null}}',
+            '{"mcpServers":{"s1":[]}}',
+        ],
+    )
+    def test_infer_protocol_spec_mcp_session_affinity_invalid_spec(
+        self,
+        protocol_spec,
+    ):
+        """测试无效 protocol_spec 无法推断 session_affinity。"""
+        tool = Tool(tool_name="my-tool", protocol_spec=protocol_spec)
+        assert tool._infer_protocol_spec_mcp_session_affinity() is None
+
+    @pytest.mark.parametrize(
+        "protocol_spec",
+        [
+            '{"mcpServers":{"s1":{"url":"https://external-mcp.com/sse"}}}',
+            '{"mcpServers":{"s1":{"transportType":"sse","url":"https://external-mcp.com/sse"}}}',
+        ],
+    )
+    def test_infer_protocol_spec_mcp_session_affinity_sse(
+        self,
+        protocol_spec,
+    ):
+        """测试 protocol_spec 缺省或显式 sse 时推断 MCP_SSE。"""
+        tool = Tool(tool_name="my-tool", protocol_spec=protocol_spec)
+        assert tool._infer_protocol_spec_mcp_session_affinity() == "MCP_SSE"
+
+    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(