Add recent build tools to MCP server

Adds two new tools to the depot_tools MCP server related to retrieving
recent builds from Buildbucket. get_recent_builds retrieves the N most
recent completed builds from the specified builder, while
get_recent_failed_builds retrieves the N most recent failed (i.e. red)
builds.

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

mcp/buildbucket.py

@@ -4,6 +4,7 @@
 """Tools for interacting with buildbucket"""
 import json
 import subprocess
+import urllib.parse
 
 from mcp.server import fastmcp
 import telemetry
@@ -256,3 +257,144 @@ async def get_build(
             await ctx.info('Exception calling prpc')
             return f'Exception calling prpc return {e}'
         return result.stdout
+
+
+async def get_recent_builds(
+    ctx: fastmcp.Context,
+    builder_name: str,
+    builder_bucket: str,
+    builder_project: str,
+    num_builds: int,
+) -> str:
+    """Gets |num_builds| recent completed builds for a builder.
+
+    This will consider any builds that have run to completion, regardless of
+    status.
+
+    The url of a builder can be deconstructed to get the relevant information,
+    e.g.
+    https://ci.chromium.org/ui/p/<project>/builders/<bucket>/<name>
+
+    Args:
+        builder_name: The name of the builder to get builds from. Any URL
+            encoding will automatically be decoded.
+        builder_bucket: The bucket the builder belongs to.
+        builder_project: The project the builder belongs to.
+        num_builds: How many builds to retrieve. Per the proto definition for
+            the underlying request, values >1000 will be treated as 1000.
+
+    Returns:
+        The stdout of the prpc command which should be a JSON string for a
+        buildbucket.v2.SearchBuildsResponse proto. See
+        https://source.chromium.org/chromium/infra/infra_superproject/+/main:infra/go/src/go.chromium.org/luci/buildbucket/proto/builds_service.proto
+        for more details.
+    """
+    with tracer.start_as_current_span('chromium.mcp.get_recent_builds'):
+        return await _get_recent_builds(
+            ctx,
+            builder_name,
+            builder_bucket,
+            builder_project,
+            num_builds,
+            failed_builds_only=False,
+        )
+
+
+async def get_recent_failed_builds(
+    ctx: fastmcp.Context,
+    builder_name: str,
+    builder_bucket: str,
+    builder_project: str,
+    num_builds: int,
+) -> str:
+    """Gets |num_builds| recent failed builds for a builder.
+
+    This will only consider builds that have run to completion and exited with
+    the FAILURE status, i.e. builds that show up as red in Milo.
+
+    The url of a builder can be deconstructed to get the relevant information,
+    e.g.
+    https://ci.chromium.org/ui/p/<project>/builders/<bucket>/<name>
+
+    Args:
+        builder_name: The name of the builder to get builds from. Any URL
+            encoding will automatically be decoded.
+        builder_bucket: The bucket the builder belongs to.
+        builder_project: The project the builder belongs to.
+        num_builds: How many builds to retrieve. Per the proto definition for
+            the underlying request, values >1000 will be treated as 1000.
+
+    Returns:
+        The stdout of the prpc command which should be a JSON string for a
+        buildbucket.v2.SearchBuildsResponse proto. See
+        https://source.chromium.org/chromium/infra/infra_superproject/+/main:infra/go/src/go.chromium.org/luci/buildbucket/proto/builds_service.proto
+        for more details.
+    """
+    with tracer.start_as_current_span('chromium.mcp.get_recent_failed_builds'):
+        return await _get_recent_builds(
+            ctx,
+            builder_name,
+            builder_bucket,
+            builder_project,
+            num_builds,
+            failed_builds_only=True,
+        )
+
+
+async def _get_recent_builds(
+    ctx: fastmcp.Context,
+    builder_name: str,
+    builder_bucket: str,
+    builder_project: str,
+    num_builds: int,
+    failed_builds_only: bool,
+) -> str:
+    """Helper function to get recent builds for a builder.
+
+    See docstrings for get_recent_builds/get_recent_failed_builds for more
+    information.
+
+    Args:
+        builder_name: Same as caller.
+        builder_bucket: Same as caller.
+        builder_project: Same as caller.
+        num_builds: Same as caller.
+        failed_builds_only: Whether to only search for failed builds instead of
+            all completed builds.
+
+    Returns:
+        Same as caller.
+    """
+    if num_builds < 1:
+        raise ValueError(f'Provided num_builds {num_builds} is not positive')
+    request = {
+        'predicate': {
+            'builder': {
+                'project': builder_project,
+                'bucket': builder_bucket,
+                'builder': urllib.parse.unquote(builder_name),
+            },
+            'status': 'FAILURE' if failed_builds_only else 'ENDED_MASK',
+        },
+        '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

mcp/buildbucket_test.py

@@ -16,6 +16,7 @@ sys.path.insert(
     os.path.abspath(
         pathlib.Path(__file__).resolve().parent.parent.joinpath(
             pathlib.Path('infra_lib'))))
+from mcp.server import fastmcp
 import buildbucket
 
 
@@ -199,6 +200,194 @@ class BuildbucketTest(unittest.IsolatedAsyncioTestCase):
         self.assertIn('Exception calling prpc', result)
         self.assertIn('PRPC call failed', result)
 
+    @patch('subprocess.run')
+    async def test_get_recent_builds_success(self, mock_subprocess_run):
+        expected_output = '{"builds": [{"id": "1"}]}'
+        mock_subprocess_run.return_value = subprocess.CompletedProcess(
+            args=[],
+            returncode=0,
+            stdout=expected_output,
+            stderr='',
+        )
+
+        output = await buildbucket.get_recent_builds(
+            self.mock_context,
+            'test_builder',
+            'try',
+            'chromium',
+            10,
+        )
+
+        self.assertEqual(output, expected_output)
+        expected_command = [
+            'prpc',
+            'call',
+            'cr-buildbucket.appspot.com',
+            'buildbucket.v2.Builds.SearchBuilds',
+        ]
+        expected_request = {
+            'predicate': {
+                'builder': {
+                    'project': 'chromium',
+                    'bucket': 'try',
+                    'builder': 'test_builder',
+                },
+                'status': 'ENDED_MASK',
+            },
+            'page_size': '10'
+        }
+        mock_subprocess_run.assert_called_once_with(
+            expected_command,
+            capture_output=True,
+            input=json.dumps(expected_request),
+            check=True,
+            text=True,
+        )
+
+    @patch('subprocess.run')
+    async def test_get_recent_builds_with_url_encoding_success(
+            self, mock_subprocess_run):
+        builder_name_encoded = 'test%20builder'
+        builder_name_decoded = 'test builder'
+        expected_output = '{"builds": [{"id": "1"}]}'
+        mock_subprocess_run.return_value = subprocess.CompletedProcess(
+            args=[],
+            returncode=0,
+            stdout=expected_output,
+            stderr='',
+        )
+
+        output = await buildbucket.get_recent_builds(
+            self.mock_context,
+            builder_name_encoded,
+            'try',
+            'chromium',
+            10,
+        )
+
+        self.assertEqual(output, expected_output)
+        expected_command = [
+            'prpc',
+            'call',
+            'cr-buildbucket.appspot.com',
+            'buildbucket.v2.Builds.SearchBuilds',
+        ]
+        expected_request = {
+            'predicate': {
+                'builder': {
+                    'project': 'chromium',
+                    'bucket': 'try',
+                    'builder': builder_name_decoded,
+                },
+                'status': 'ENDED_MASK',
+            },
+            'page_size': '10'
+        }
+        mock_subprocess_run.assert_called_once_with(
+            expected_command,
+            capture_output=True,
+            input=json.dumps(expected_request),
+            check=True,
+            text=True,
+        )
+
+    @patch('subprocess.run')
+    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:
+            await buildbucket.get_recent_builds(
+                self.mock_context,
+                'test_builder',
+                'try',
+                '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,
+                                    'Provided num_builds 0 is not positive'):
+            await buildbucket.get_recent_builds(
+                self.mock_context,
+                'test_builder',
+                'try',
+                'chromium',
+                0,
+            )
+
+    @patch('subprocess.run')
+    async def test_get_recent_failed_builds_success(self, mock_subprocess_run):
+        expected_output = '{"builds": [{"id": "1", "status": "FAILURE"}]}'
+        mock_subprocess_run.return_value = subprocess.CompletedProcess(
+            args=[],
+            returncode=0,
+            stdout=expected_output,
+            stderr='',
+        )
+
+        output = await buildbucket.get_recent_failed_builds(
+            self.mock_context,
+            'test_builder',
+            'try',
+            'chromium',
+            10,
+        )
+
+        self.assertEqual(output, expected_output)
+        expected_command = [
+            'prpc',
+            'call',
+            'cr-buildbucket.appspot.com',
+            'buildbucket.v2.Builds.SearchBuilds',
+        ]
+        expected_request = {
+            'predicate': {
+                'builder': {
+                    'project': 'chromium',
+                    'bucket': 'try',
+                    'builder': 'test_builder',
+                },
+                'status': 'FAILURE',
+            },
+            'page_size': '10'
+        }
+        mock_subprocess_run.assert_called_once_with(
+            expected_command,
+            capture_output=True,
+            input=json.dumps(expected_request),
+            check=True,
+            text=True,
+        )
+
+    @patch('subprocess.run')
+    async def test_get_recent_failed_builds_exception(self,
+                                                      mock_subprocess_run):
+        mock_subprocess_run.side_effect = Exception('PRPC call failed')
+
+        with self.assertRaises(fastmcp.exceptions.ToolError) as e:
+            await buildbucket.get_recent_failed_builds(
+                self.mock_context,
+                'test_builder',
+                'try',
+                '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,
+                                    'Provided num_builds -1 is not positive'):
+            await buildbucket.get_recent_failed_builds(
+                self.mock_context,
+                'test_builder',
+                'try',
+                'chromium',
+                -1,
+            )
+
 
 if __name__ == '__main__':
     unittest.main()

mcp/server.py

@@ -38,6 +38,8 @@ def main(argv: Sequence[str]) -> None:
     mcp.add_tool(buildbucket.get_build_from_build_number)
     mcp.add_tool(buildbucket.get_build_from_id)
     mcp.add_tool(buildbucket.get_build_status)
+    mcp.add_tool(buildbucket.get_recent_builds)
+    mcp.add_tool(buildbucket.get_recent_failed_builds)
     mcp.run()