Deduplicate mcp/ prpc calls

Deduplicates the various `prpc call` subprocess.run calls throughout
mcp/ in favor of a new run_prpc_call helper.

Also makes a few drive-by cleanup changes that were caught by WIP
presubmit checks.

Change-Id: I3e401cb959ba96f2224ea1a7f7f58795330f643a
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/tools/depot_tools/+/6847120
Commit-Queue: Struan Shrimpton <sshrimp@google.com>
Auto-Submit: Brian Sheedy <bsheedy@chromium.org>
Reviewed-by: Struan Shrimpton <sshrimp@google.com>

mcp/buildbucket.py

@@ -3,12 +3,15 @@
 # found in the LICENSE file.
 """Tools for interacting with buildbucket"""
 import json
-import subprocess
 import urllib.parse
 
 from mcp.server import fastmcp
 import telemetry
 
+import common
+
+BUILDBUCKET_SERVER = 'cr-buildbucket.appspot.com'
+
 tracer = telemetry.get_tracer(__name__)
 
 
@@ -25,26 +28,11 @@ async def get_build_status(
     """
     with tracer.start_as_current_span('chromium.mcp.get_build_status'):
         await ctx.info(f'Received request {build_id}')
-        command = [
-            'prpc',
-            'call',
-            'cr-buildbucket.appspot.com',
-            'buildbucket.v2.Builds.GetBuildStatus',
-        ]
-        try:
-            result = subprocess.run(
-                command,
-                capture_output=True,
-                input=json.dumps({'id': build_id}),
-                check=False,
-                text=True,
-            )
-            await ctx.info(result.stdout)
-            await ctx.info(result.stderr)
-            return json.loads(result.stdout)['status']
-        except Exception as e:
-            await ctx.info('Exception calling prpc')
-            return f'Exception calling prpc return {e}'
+        request = {'id': build_id}
+        response = await common.run_prpc_call(
+            ctx, BUILDBUCKET_SERVER, 'buildbucket.v2.Builds.GetBuildStatus',
+            request)
+        return json.loads(response)['status']
 
 
 async def get_build_from_id(
@@ -87,26 +75,10 @@ async def get_build_from_id(
     """
     with tracer.start_as_current_span('chromium.mcp.get_build_from_id'):
         request = {'id': build_id, 'mask': {'fields': ','.join(fields)}}
-        command = [
-            'prpc',
-            'call',
-            'cr-buildbucket.appspot.com',
-            'buildbucket.v2.Builds.GetBuild',
-        ]
-        try:
-            result = subprocess.run(
-                command,
-                capture_output=True,
-                input=json.dumps(request),
-                check=False,
-                text=True,
-            )
-            await ctx.info(result.stdout)
-            await ctx.info(result.stderr)
-        except Exception as e:
-            await ctx.info('Exception calling prpc')
-            return f'Exception calling prpc return {e}'
-        return result.stdout
+        response = await common.run_prpc_call(ctx, BUILDBUCKET_SERVER,
+                                              'buildbucket.v2.Builds.GetBuild',
+                                              request)
+        return response
 
 
 async def get_build_from_build_number(
@@ -162,26 +134,10 @@ async def get_build_from_build_number(
                 'fields': ','.join(fields)
             }
         }
-        command = [
-            'prpc',
-            'call',
-            'cr-buildbucket.appspot.com',
-            'buildbucket.v2.Builds.GetBuild',
-        ]
-        try:
-            result = subprocess.run(
-                command,
-                capture_output=True,
-                input=json.dumps(request),
-                check=False,
-                text=True,
-            )
-            await ctx.info(result.stdout)
-            await ctx.info(result.stderr)
-        except Exception as e:
-            await ctx.info('Exception calling prpc')
-            return f'Exception calling prpc return {e}'
-        return result.stdout
+        response = await common.run_prpc_call(ctx, BUILDBUCKET_SERVER,
+                                              'buildbucket.v2.Builds.GetBuild',
+                                              request)
+        return response
 
 
 async def get_build(
@@ -237,26 +193,10 @@ async def get_build(
     """
     with tracer.start_as_current_span('chromium.mcp.get_build'):
         await ctx.info(f'Received request {request}')
-        command = [
-            'prpc',
-            'call',
-            'cr-buildbucket.appspot.com',
-            'buildbucket.v2.Builds.GetBuild',
-        ]
-        try:
-            result = subprocess.run(
-                command,
-                capture_output=True,
-                input=json.dumps(request),
-                check=False,
-                text=True,
-            )
-            await ctx.info(result.stdout)
-            await ctx.info(result.stderr)
-        except Exception as e:
-            await ctx.info('Exception calling prpc')
-            return f'Exception calling prpc return {e}'
-        return result.stdout
+        response = await common.run_prpc_call(ctx, BUILDBUCKET_SERVER,
+                                              'buildbucket.v2.Builds.GetBuild',
+                                              request)
+        return response
 
 
 async def get_recent_builds(
@@ -378,23 +318,7 @@ async def _get_recent_builds(
         },
         'page_size': f'{num_builds}'
     }
-    command = [
-        'prpc',
-        'call',
-        'cr-buildbucket.appspot.com',
-        'buildbucket.v2.Builds.SearchBuilds',
-    ]
-    try:
-        result = subprocess.run(
-            command,
-            capture_output=True,
-            input=json.dumps(request),
-            check=True,
-            text=True,
-        )
-        await ctx.info(result.stdout)
-        await ctx.info(result.stderr)
-    except Exception as e:
-        raise fastmcp.exceptions.ToolError(
-            f'Exception calling prpc: {e}') from e
-    return result.stdout
+    response = await common.run_prpc_call(ctx, BUILDBUCKET_SERVER,
+                                          'buildbucket.v2.Builds.SearchBuilds',
+                                          request)
+    return response

mcp/buildbucket_test.py

@@ -16,7 +16,6 @@ sys.path.insert(
     os.path.abspath(
         pathlib.Path(__file__).resolve().parent.parent.joinpath(
             pathlib.Path('infra_lib'))))
-from mcp.server import fastmcp
 import buildbucket
 
 
@@ -50,7 +49,7 @@ class BuildbucketTest(unittest.IsolatedAsyncioTestCase):
             expected_command,
             capture_output=True,
             input=json.dumps({'id': build_id}),
-            check=False,
+            check=True,
             text=True,
         )
 
@@ -59,10 +58,8 @@ class BuildbucketTest(unittest.IsolatedAsyncioTestCase):
         build_id = '12345'
         mock_subprocess_run.side_effect = Exception('PRPC call failed')
 
-        result = await buildbucket.get_build_status(self.mock_context, build_id)
-
-        self.assertIn('Exception calling prpc', result)
-        self.assertIn('PRPC call failed', result)
+        with self.assertRaisesRegex(Exception, 'PRPC call failed'):
+            await buildbucket.get_build_status(self.mock_context, build_id)
 
     @mock.patch('subprocess.run')
     async def test_get_build_from_id_success(self, mock_subprocess_run):
@@ -88,7 +85,7 @@ class BuildbucketTest(unittest.IsolatedAsyncioTestCase):
             expected_command,
             capture_output=True,
             input=json.dumps(expected_request),
-            check=False,
+            check=True,
             text=True)
 
     @mock.patch('subprocess.run')
@@ -97,14 +94,12 @@ class BuildbucketTest(unittest.IsolatedAsyncioTestCase):
         fields = ['steps']
         mock_subprocess_run.side_effect = Exception('PRPC call failed')
 
-        result = await buildbucket.get_build_from_id(
-            self.mock_context,
-            build_id,
-            fields,
-        )
-
-        self.assertIn('Exception calling prpc', result)
-        self.assertIn('PRPC call failed', result)
+        with self.assertRaisesRegex(Exception, 'PRPC call failed'):
+            await buildbucket.get_build_from_id(
+                self.mock_context,
+                build_id,
+                fields,
+            )
 
     @mock.patch('subprocess.run')
     async def test_get_build_from_build_number_success(
@@ -144,7 +139,7 @@ class BuildbucketTest(unittest.IsolatedAsyncioTestCase):
             expected_command,
             capture_output=True,
             input=json.dumps(expected_request),
-            check=False,
+            check=True,
             text=True)
 
     @mock.patch('subprocess.run')
@@ -157,12 +152,10 @@ class BuildbucketTest(unittest.IsolatedAsyncioTestCase):
         fields = ['status']
         mock_subprocess_run.side_effect = Exception('PRPC call failed')
 
-        result = await buildbucket.get_build_from_build_number(
-            self.mock_context, build_number, builder_name, builder_bucket,
-            builder_project, fields)
-
-        self.assertIn('Exception calling prpc', result)
-        self.assertIn('PRPC call failed', result)
+        with self.assertRaisesRegex(Exception, 'PRPC call failed'):
+            await buildbucket.get_build_from_build_number(
+                self.mock_context, build_number, builder_name, builder_bucket,
+                builder_project, fields)
 
     @mock.patch('subprocess.run')
     async def test_get_build_success(self, mock_subprocess_run):
@@ -186,7 +179,7 @@ class BuildbucketTest(unittest.IsolatedAsyncioTestCase):
             expected_command,
             capture_output=True,
             input=json.dumps(request),
-            check=False,
+            check=True,
             text=True,
         )
 
@@ -195,10 +188,8 @@ class BuildbucketTest(unittest.IsolatedAsyncioTestCase):
         request = {"id": "12345"}
         mock_subprocess_run.side_effect = Exception('PRPC call failed')
 
-        result = await buildbucket.get_build(self.mock_context, request)
-
-        self.assertIn('Exception calling prpc', result)
-        self.assertIn('PRPC call failed', result)
+        with self.assertRaisesRegex(Exception, 'PRPC call failed'):
+            await buildbucket.get_build(self.mock_context, request)
 
     @mock.patch('subprocess.run')
     async def test_get_recent_builds_success(self, mock_subprocess_run):
@@ -295,7 +286,7 @@ class BuildbucketTest(unittest.IsolatedAsyncioTestCase):
     async def test_get_recent_builds_exception(self, mock_subprocess_run):
         mock_subprocess_run.side_effect = Exception('PRPC call failed')
 
-        with self.assertRaises(fastmcp.exceptions.ToolError) as e:
+        with self.assertRaisesRegex(Exception, 'PRPC call failed'):
             await buildbucket.get_recent_builds(
                 self.mock_context,
                 'test_builder',
@@ -303,8 +294,6 @@ class BuildbucketTest(unittest.IsolatedAsyncioTestCase):
                 'chromium',
                 10,
             )
-        self.assertIn('Exception calling prpc', str(e.exception))
-        self.assertIn('PRPC call failed', str(e.exception))
 
     async def test_get_recent_builds_invalid_num_builds(self):
         with self.assertRaisesRegex(ValueError,
@@ -366,7 +355,7 @@ class BuildbucketTest(unittest.IsolatedAsyncioTestCase):
                                                       mock_subprocess_run):
         mock_subprocess_run.side_effect = Exception('PRPC call failed')
 
-        with self.assertRaises(fastmcp.exceptions.ToolError) as e:
+        with self.assertRaisesRegex(Exception, 'PRPC call failed'):
             await buildbucket.get_recent_failed_builds(
                 self.mock_context,
                 'test_builder',
@@ -374,8 +363,6 @@ class BuildbucketTest(unittest.IsolatedAsyncioTestCase):
                 'chromium',
                 10,
             )
-        self.assertIn('Exception calling prpc', str(e.exception))
-        self.assertIn('PRPC call failed', str(e.exception))
 
     async def test_get_recent_failed_builds_invalid_num_builds(self):
         with self.assertRaisesRegex(ValueError,

mcp/common.py

@@ -0,0 +1,39 @@
+# Copyright 2025 The Chromium Authors
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+"""Common code shared between different MCP server areas."""
+
+import json
+import subprocess
+
+from mcp.server import fastmcp
+
+
+async def run_prpc_call(ctx: fastmcp.Context, server: str, service: str,
+                        message: dict) -> str:
+    """Runs 'prpc call' with the given parameters.
+
+    Args:
+        server: The server the request is for, e.g. cr-buildbucket.appspot.com.
+        service: The specific RPC service to call.
+        message: The RPC message to send to the service.
+
+    Returns:
+        A string containing the JSON response of the call.
+    """
+    command = [
+        'prpc',
+        'call',
+        server,
+        service,
+    ]
+    result = subprocess.run(
+        command,
+        capture_output=True,
+        input=json.dumps(message),
+        check=True,
+        text=True,
+    )
+    await ctx.info(result.stdout)
+    await ctx.info(result.stderr)
+    return result.stdout

mcp/git_cl.py

@@ -19,7 +19,8 @@ async def try_builder_results(
       checkout: Location of the current checkout.
 
     Returns:
-      A json list of builds that either ran or are still running on the current CL
+      A json list of builds that either ran or are still running on the current
+      CL
     """
     with tracer.start_as_current_span('chromium.mcp.try_builder_results'):
         command = [
@@ -70,7 +71,7 @@ async def get_current_changes(
         return result.stdout
 
 
-async def format(
+async def format_checkout(
     ctx: fastmcp.Context,
     checkout: str,
 ) -> None:

mcp/git_cl_test.py

@@ -60,12 +60,12 @@ class GitClTest(unittest.IsolatedAsyncioTestCase):
                                                     cwd=self.checkout)
 
     @mock.patch('subprocess.run')
-    async def test_format_success(self, mock_subprocess_run):
+    async def test_format_checkout_success(self, mock_subprocess_run):
         expected_output = 'Formatted 1 file'
         mock_subprocess_run.return_value = subprocess.CompletedProcess(
             args=[], returncode=0, stdout=expected_output, stderr='')
 
-        output = await git_cl.format(self.mock_context, self.checkout)
+        output = await git_cl.format_checkout(self.mock_context, self.checkout)
 
         self.assertEqual(output, expected_output)
         expected_command = ["git", "cl", "format"]

mcp/resultdb.py

@@ -6,14 +6,17 @@
 import json
 import posixpath
 import requests
-import subprocess
 
 import bs4
 from mcp.server import fastmcp
 import telemetry
 
+import common
+
 tracer = telemetry.get_tracer(__name__)
 
+RESULTDB_SERVER = 'results.api.luci.app'
+
 
 async def get_non_exonerated_unexpected_results_from_build(
     ctx: fastmcp.Context,
@@ -50,26 +53,10 @@ async def get_non_exonerated_unexpected_results_from_build(
             'read_mask': ('name,resultId,variant,status,statusV2,duration,'
                           'failureReason,summary_html'),
         }
-        command = [
-            'prpc',
-            'call',
-            'results.api.luci.app',
-            'luci.resultdb.v1.ResultDB.QueryTestResults',
-        ]
-        try:
-            result = subprocess.run(
-                command,
-                capture_output=True,
-                input=json.dumps(request),
-                check=True,
-                text=True,
-            )
-            await ctx.info(result.stdout)
-            await ctx.info(result.stderr)
-        except Exception as e:
-            raise fastmcp.exceptions.ToolError(
-                f'Exception calling prpc: {e}') from e
-        return result.stdout
+        response = await common.run_prpc_call(
+            ctx, RESULTDB_SERVER, 'luci.resultdb.v1.ResultDB.QueryTestResults',
+            request)
+        return response
 
 
 async def expand_summary_html(
@@ -137,26 +124,10 @@ async def get_test_level_text_artifact(
         request = {
             'name': artifact_name,
         }
-        command = [
-            'prpc',
-            'call',
-            'results.api.luci.app',
-            'luci.resultdb.v1.ResultDB.GetArtifact',
-        ]
-        try:
-            result = subprocess.run(
-                command,
-                capture_output=True,
-                input=json.dumps(request),
-                check=True,
-                text=True,
-            )
-            await ctx.info(result.stdout)
-            await ctx.info(result.stderr)
-        except Exception as e:
-            raise fastmcp.exceptions.ToolError(
-                f'Exception calling prpc: {e}') from e
-        response = json.loads(result.stdout)
+        prpc_response = await common.run_prpc_call(
+            ctx, RESULTDB_SERVER, 'luci.resultdb.v1.ResultDB.GetArtifact',
+            request)
+        response = json.loads(prpc_response)
 
         content_type = response['contentType']
         if 'text/plain' not in content_type:

mcp/resultdb_test.py

@@ -16,7 +16,6 @@ sys.path.insert(
     os.path.abspath(
         pathlib.Path(__file__).resolve().parent.parent.joinpath(
             pathlib.Path('infra_lib'))))
-from mcp.server import fastmcp
 import requests
 
 import resultdb
@@ -79,11 +78,9 @@ class GetNonExoneratedUnexpectedResultsFromBuildTest(
         build_id = '12345'
         mock_subprocess_run.side_effect = Exception('PRPC call failed')
 
-        with self.assertRaises(fastmcp.exceptions.ToolError) as e:
+        with self.assertRaisesRegex(Exception, 'PRPC call failed'):
             await resultdb.get_non_exonerated_unexpected_results_from_build(
                 self.mock_context, build_id)
-        self.assertIn('Exception calling prpc', str(e.exception))
-        self.assertIn('PRPC call failed', str(e.exception))
 
     async def test_get_non_exonerated_unexpected_results_from_build_invalid_id(
             self):
@@ -200,12 +197,10 @@ class GetTestLevelTextArtifactTest(unittest.IsolatedAsyncioTestCase):
         artifact_id = 'artifact1'
         mock_subprocess_run.side_effect = Exception('PRPC call failed')
 
-        with self.assertRaises(fastmcp.exceptions.ToolError) as e:
+        with self.assertRaisesRegex(Exception, 'PRPC call failed'):
             await resultdb.get_test_level_text_artifact(self.mock_context,
                                                         result_name,
                                                         artifact_id)
-        self.assertIn('Exception calling prpc', str(e.exception))
-        self.assertIn('PRPC call failed', str(e.exception))
 
     @mock.patch('subprocess.run')
     async def test_get_test_level_text_artifact_wrong_content_type(

mcp/server.py

@@ -45,7 +45,7 @@ def main(argv: Sequence[str]) -> None:
     mcp.add_tool(resultdb.get_test_level_text_artifact)
     mcp.add_tool(git_cl.try_builder_results)
     mcp.add_tool(git_cl.get_current_changes)
-    mcp.add_tool(git_cl.format)
+    mcp.add_tool(git_cl.format_checkout)
     mcp.add_tool(git_cl.upload_change_list)
     mcp.run()