mcp: Add some buildbucket tools and the mcp server

These tools are still being evaluated and may change. get_build vs get_build_from_id and get_build_from_build_number in particular are two distinct ways of implementing the same tool that need to be compared. get_build gives flexibility but is extremely verbose which means it might be too flaky or consume unnecessary tokens. The return of these might also need filtering for unnecessary output. Particularly steps which can be very verbose without much information. Bug: 430952168 Change-Id: Ia9b092a6b6c546198939deb36ea1b89ff33d7fc6 Reviewed-on: https://chromium-review.googlesource.com/c/chromium/tools/depot_tools/+/6758660 Auto-Submit: Struan Shrimpton <sshrimp@google.com> Reviewed-by: Gavin Mak <gavinmak@google.com> Reviewed-by: Erik Staab <estaab@chromium.org> Commit-Queue: Struan Shrimpton <sshrimp@google.com>
4 months ago · 75449d2000
parent 068f8379e9
commit 75449d2000
6 changed files with 692 additions and 0 deletions
--- a/infra_lib/telemetry/init.py
+++ b/infra_lib/telemetry/init.py
@ -113,3 +113,23 @@ def initialize(service_name,

 def get_tracer(name: str, version: Optional[str] = None):
    return otel_trace_api.get_tracer(name, version)
+
+
+def opted_in(cfg_file=config.DEFAULT_CONFIG_FILE):
+    """Get if the user is opted-in
+
+    Unlike initialize which continues when not explicitly opted out this will
+    return if the user is opted in, either by user or automatically after the
+    banner display period.
+    """
+    cfg = config.Config(cfg_file)
+    if cfg.trace_config.disabled():
+        return False
+
+    bot_enabled = (cfg.trace_config.has_enabled()
+                   and cfg.trace_config.enabled_reason == 'BOT_USER')
+    if not is_google_host() and not bot_enabled:
+        return False
+
+    cfg = config.Config(cfg_file)
+    return cfg.trace_config.enabled
--- a/mcp/.vpython3
+++ b/mcp/.vpython3
@ -0,0 +1,162 @@
+# This includes the wheels necessary to fast mcp 1.9.4 and telemetry
+#
+# Read more about `vpython` and how to modify this file here:
+#   https://chromium.googlesource.com/infra/infra/+/main/doc/users/vpython.md
+
+python_version: "3.11"
+
+wheel: <
+  name: "infra/python/wheels/certifi-py3"
+  version: "version:2025.4.26"
+>
+
+wheel: <
+  name: "infra/python/wheels/idna-py3"
+  version: "version:3.4"
+>
+
+wheel: <
+  name: "infra/python/wheels/sniffio-py3"
+  version: "version:1.3.0"
+>
+
+wheel: <
+  name: "infra/python/wheels/typing-extensions-py3"
+  version: "version:4.13.2"
+>
+
+wheel: <
+  name: "infra/python/wheels/click-py3"
+  version: "version:8.0.3"
+>
+
+wheel: <
+  name: "infra/python/wheels/absl-py-py3"
+  version: "version:2.1.0"
+>
+
+wheel: <
+  name: "infra/python/wheels/mcp-py3"
+  version: "version:1.9.4"
+>
+
+wheel: <
+  name: "infra/python/wheels/pydantic-py3"
+  version: "version:2.11.7"
+>
+
+wheel: <
+  name: "infra/python/wheels/httpx-py3"
+  version: "version:0.28.1"
+>
+
+wheel: <
+  name: "infra/python/wheels/httpcore-py3"
+  version: "version:1.0.9"
+>
+
+wheel: <
+  name: "infra/python/wheels/h11-py3"
+  version: "version:0.16.0"
+>
+
+wheel: <
+  name: "infra/python/wheels/anyio-py3"
+  version: "version:4.9.0"
+>
+
+wheel: <
+  name: "infra/python/wheels/httpx_sse-py3"
+  version: "version:0.4.1"
+>
+
+wheel: <
+  name: "infra/python/wheels/pydantic-settings-py3"
+  version: "version:2.10.1"
+>
+
+wheel: <
+  name: "infra/python/wheels/python-multipart-py3"
+  version: "version:0.0.20"
+>
+
+wheel: <
+  name: "infra/python/wheels/sse-starlette-py3"
+  version: "version:2.4.1"
+>
+
+wheel: <
+  name: "infra/python/wheels/starlette-py3"
+  version: "version:0.47.1"
+>
+
+wheel: <
+  name: "infra/python/wheels/uvicorn-py3"
+  version: "version:0.35.0"
+>
+
+wheel: <
+  name: "infra/python/wheels/annotated-types-py3"
+  version: "version:0.7.0"
+>
+
+wheel: <
+  name: "infra/python/wheels/pydantic_core-py3"
+  version: "version:2.33.2"
+>
+
+wheel: <
+  name: "infra/python/wheels/typing-inspection-py3"
+  version: "version:0.4.1"
+>
+
+wheel: <
+  name: "infra/python/wheels/python-dotenv-py3"
+  version: "version:1.1.1"
+>
+
+# Open Telemetry wheels
+wheel: <
+  name: "infra/python/wheels/opentelemetry-api-py3"
+  version: "version:1.18.0"
+>
+
+wheel: <
+  name: "infra/python/wheels/opentelemetry-sdk-py3"
+  version: "version:1.18.0"
+>
+
+wheel: <
+  name: "infra/python/wheels/deprecated-py3"
+  version: "version:1.2.13"
+>
+
+wheel: <
+  name: "infra/python/wheels/wrapt-py3"
+  version: "version:1.15.0"
+>
+
+wheel: <
+  name: "infra/python/wheels/importlib-metadata-py3"
+  version: "version:6.0.0"
+>
+
+wheel: <
+  name: "infra/python/wheels/zipp-py3"
+  version: "version:3.17.0"
+>
+
+wheel: <
+  name: "infra/python/wheels/opentelemetry-semantic-conventions-py3"
+  version: "version:0.39b0"
+>
+
+wheel: <
+  name: "infra/python/wheels/googleapis-common-protos-py2_py3"
+  version: "version:1.61.0"
+>
+
+wheel: <
+  name: "infra/python/wheels/protobuf-py3"
+  version: "version:4.25.1"
+>
--- a/mcp/OWNERS
+++ b/mcp/OWNERS
@ -0,0 +1,3 @@
+estaab@google.com
+jiesheng@google.com
+sshrimp@google.com
--- a/mcp/buildbucket.py
+++ b/mcp/buildbucket.py
@ -0,0 +1,258 @@
+# 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.
+"""Tools for interacting with buildbucket"""
+import json
+import subprocess
+
+from mcp.server import fastmcp
+import telemetry
+
+tracer = telemetry.get_tracer(__name__)
+
+
+async def get_build_status(
+    ctx: fastmcp.Context,
+    build_id: str,
+) -> str:
+    """Gets the build status from the provided build_id
+
+    Args:
+      build_id: The buildbucket id of the build. This is not the build number.
+    Return:
+      The status of the build as a string
+    """
+    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}'
+
+
+async def get_build_from_id(
+    ctx: fastmcp.Context,
+    build_id: str,
+    fields: list[str],
+) -> str:
+    """Gets a buildbucket build from its ID
+
+    The url of a build can be deconstructed and used to get more details about
+    the build. e.g.
+    https://ci.chromium.org/b/<build_id>
+
+    Args:
+      build_id: The request body for the RPC. All fields should be represented
+        by strings. Integer fields will be parsed later.
+        https://chromium.googlesource.com/infra/luci/luci-go/+/main/buildbucket/proto/builds_service.proto
+        for more details.
+
+        The request's mask can be set to get more information. By default only
+        high level statuses will be returned. Some useful fields to include in
+        this mask are:
+        status, input, output, id, builder, builder_info, tags, steps, infra
+        Multiple fields in the mask can be included as a comma separated string
+        e.g.
+        The build_number is mutually exclusive with the build_id. To get the
+        build from a build_id, only the build_id is needed. e.g.
+        {
+          "id": "<build id>",
+          "mask": {
+            "fields": "steps,tags"
+          }
+        }
+      fields: A list of fields to return. Options are:
+        status, input, output, id, builder, builder_info, tags, steps, infra
+
+    Returns:
+      The build in json format including the requested fields. See:
+      https://chromium.googlesource.com/infra/luci/recipes-py/+/main/recipe_proto/go.chromium.org/luci/buildbucket/proto/build.proto
+    """
+    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
+
+
+async def get_build_from_build_number(
+    ctx: fastmcp.Context,
+    build_number: int,
+    builder_name: int,
+    builder_bucket: int,
+    builder_project: int,
+    fields: list[str],
+) -> str:
+    """Gets a buildbucket build from its build number and builder
+
+    The url of a build can be deconstructed and used to get more details about
+    the build. e.g.
+    https://ci.chromium.org/b/<build_id>
+
+    Args:
+      build_id: The request body for the RPC. All fields should be represented
+        by strings. Integer fields will be parsed later.
+        https://chromium.googlesource.com/infra/luci/luci-go/+/main/buildbucket/proto/builds_service.proto
+        for more details.
+
+        The request's mask can be set to get more information. By default only
+        high level statuses will be returned. Some useful fields to include in
+        this mask are:
+        status, input, output, id, builder, builder_info, tags, steps, infra
+        Multiple fields in the mask can be included as a comma separated string
+        The build_number is mutually exclusive with the build_id. To get the
+        build from a build_id only the build_id is needed. e.g.
+        {
+          "id": "<build_id>",
+          "mask": {
+            "fields": "steps,tags"
+          }
+        }
+      fields: A list of fields to return. Options are:
+        status, input, output, id, builder, builder_info, tags, steps, infra
+
+    Returns:
+      The build in json format including the requested fields. See
+      https://chromium.googlesource.com/infra/luci/recipes-py/+/main/recipe_proto/go.chromium.org/luci/buildbucket/proto/build.proto
+    """
+    with tracer.start_as_current_span(
+            'chromium.mcp.get_build_from_build_number'):
+        request = {
+            'buildNumber': build_number,
+            'builder': {
+                'builder': builder_name,
+                'bucket': builder_bucket,
+                'project': builder_project
+            },
+            '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
+
+
+async def get_build(
+    ctx: fastmcp.Context,
+    request: dict,
+) -> str:
+    """Calls the buildbucket.v2.Builds.GetBuild RPC to fetch a build
+
+    The url of a build can be deconstructed and used to get more details about
+    the
+    build. e.g.
+    https://ci.chromium.org/ui/p/chromium/builders/<bucket>/<builder_name>/<build_number>/infra
+    Builds can also be in the form:
+    https://ci.chromium.org/b/<build id>
+
+    Args:
+      request: The request body for the RPC. All fields should be represented
+        by strings. Integer fields will be parsed later.
+        https://chromium.googlesource.com/infra/luci/luci-go/+/main/buildbucket/proto/builds_service.proto
+        for more details.
+
+        The request's mask can be set to get more information. By default only
+        high level statuses will be returned. Some useful fields to include in
+        this mask are:
+        status, input, output, id, builder, builder_info, tags, steps, infra
+        Multiple fields in the mask can be included as a comma separated string
+        e.g.
+        {
+          "build_number": "<build number>",
+          "builder": {
+            "bucket": "<bucket>",
+            "builder": "<builder name>",
+            "project": "chromium"
+          },
+          "mask": {
+            "fields": "steps,tags"
+          }
+        }
+        The build_number is mutually exclusive with the build_id. To get the
+        build from a build_id, only the build_id is needed. e.g.
+        {
+          "id": "<build id>",
+          "mask": {
+            "fields": "steps,tags"
+          }
+        }
+
+    Returns:
+      The stdout of the prpc command which should be a JSON string for a
+      buildbucket.v2.Build proto. See
+      https://chromium.googlesource.com/infra/luci/recipes-py/+/main/recipe_proto/go.chromium.org/luci/buildbucket/proto/build.proto
+      for more details.
+    """
+    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
--- a/mcp/buildbucket_test.py
+++ b/mcp/buildbucket_test.py
@ -0,0 +1,204 @@
+# 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.
+"""Tests for buildbucket tools."""
+
+import json
+import os
+import pathlib
+import subprocess
+import sys
+import unittest
+from unittest.mock import AsyncMock, patch
+
+sys.path.insert(
+    0,
+    os.path.abspath(
+        pathlib.Path(__file__).resolve().parent.parent.joinpath(
+            pathlib.Path('infra_lib'))))
+import buildbucket
+
+
+class BuildbucketTest(unittest.IsolatedAsyncioTestCase):
+
+    def setUp(self):
+        self.mock_context = AsyncMock()
+        self.mock_context.info = AsyncMock()
+
+    @patch('subprocess.run')
+    async def test_get_build_status_success(self, mock_subprocess_run):
+        build_id = '12345'
+        expected_status = 'SUCCESS'
+        mock_subprocess_run.return_value = subprocess.CompletedProcess(
+            args=[],
+            returncode=0,
+            stdout=json.dumps({'status': expected_status}),
+            stderr='')
+
+        status = await buildbucket.get_build_status(
+            self.mock_context,
+            build_id,
+        )
+
+        self.assertEqual(status, expected_status)
+        expected_command = [
+            'prpc', 'call', 'cr-buildbucket.appspot.com',
+            'buildbucket.v2.Builds.GetBuildStatus'
+        ]
+        mock_subprocess_run.assert_called_once_with(
+            expected_command,
+            capture_output=True,
+            input=json.dumps({'id': build_id}),
+            check=False,
+            text=True,
+        )
+
+    @patch('subprocess.run')
+    async def test_get_build_status_exception(self, mock_subprocess_run):
+        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)
+
+    @patch('subprocess.run')
+    async def test_get_build_from_id_success(self, mock_subprocess_run):
+        build_id = '12345'
+        fields = ['steps', 'tags']
+        expected_output = '{"id": "12345", "steps": [], "tags": []}'
+        mock_subprocess_run.return_value = subprocess.CompletedProcess(
+            args=[], returncode=0, stdout=expected_output, stderr='')
+
+        output = await buildbucket.get_build_from_id(
+            self.mock_context,
+            build_id,
+            fields,
+        )
+
+        self.assertEqual(output, expected_output)
+        expected_command = [
+            'prpc', 'call', 'cr-buildbucket.appspot.com',
+            'buildbucket.v2.Builds.GetBuild'
+        ]
+        expected_request = {'id': build_id, 'mask': {'fields': 'steps,tags'}}
+        mock_subprocess_run.assert_called_once_with(
+            expected_command,
+            capture_output=True,
+            input=json.dumps(expected_request),
+            check=False,
+            text=True)
+
+    @patch('subprocess.run')
+    async def test_get_build_from_id_exception(self, mock_subprocess_run):
+        build_id = '12345'
+        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)
+
+    @patch('subprocess.run')
+    async def test_get_build_from_build_number_success(
+        self,
+        mock_subprocess_run,
+    ):
+        build_number = 987
+        builder_name = 'test_builder'
+        builder_bucket = 'try'
+        builder_project = 'chromium'
+        fields = ['status']
+        expected_output = '{"status": "SUCCESS"}'
+        mock_subprocess_run.return_value = subprocess.CompletedProcess(
+            args=[], returncode=0, stdout=expected_output, stderr='')
+
+        output = await buildbucket.get_build_from_build_number(
+            self.mock_context, build_number, builder_name, builder_bucket,
+            builder_project, fields)
+
+        self.assertEqual(output, expected_output)
+        expected_command = [
+            'prpc', 'call', 'cr-buildbucket.appspot.com',
+            'buildbucket.v2.Builds.GetBuild'
+        ]
+        expected_request = {
+            'buildNumber': build_number,
+            'builder': {
+                'builder': builder_name,
+                'bucket': builder_bucket,
+                'project': builder_project
+            },
+            'mask': {
+                'fields': 'status'
+            }
+        }
+        mock_subprocess_run.assert_called_once_with(
+            expected_command,
+            capture_output=True,
+            input=json.dumps(expected_request),
+            check=False,
+            text=True)
+
+    @patch('subprocess.run')
+    async def test_get_build_from_build_number_exception(
+            self, mock_subprocess_run):
+        build_number = 987
+        builder_name = 'test_builder'
+        builder_bucket = 'try'
+        builder_project = 'chromium'
+        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)
+
+    @patch('subprocess.run')
+    async def test_get_build_success(self, mock_subprocess_run):
+        request = {"id": "12345"}
+        expected_output = '{"id": "12345"}'
+        mock_subprocess_run.return_value = subprocess.CompletedProcess(
+            args=[],
+            returncode=0,
+            stdout=expected_output,
+            stderr='',
+        )
+
+        output = await buildbucket.get_build(self.mock_context, request)
+
+        self.assertEqual(output, expected_output)
+        expected_command = [
+            'prpc', 'call', 'cr-buildbucket.appspot.com',
+            'buildbucket.v2.Builds.GetBuild'
+        ]
+        mock_subprocess_run.assert_called_once_with(
+            expected_command,
+            capture_output=True,
+            input=json.dumps(request),
+            check=False,
+            text=True,
+        )
+
+    @patch('subprocess.run')
+    async def test_get_build_exception(self, mock_subprocess_run):
+        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)
+
+
+if __name__ == '__main__':
+    unittest.main()
--- a/mcp/server.py
+++ b/mcp/server.py
@ -0,0 +1,45 @@
+#!/bin/env vpython3
+# 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.
+"""The MCP server that provides tools"""
+from collections.abc import Sequence
+import pathlib
+import os
+import sys
+
+sys.path.insert(
+    0,
+    os.path.abspath(
+        pathlib.Path(__file__).resolve().parent.parent.joinpath(
+            pathlib.Path('infra_lib'))))
+import telemetry
+
+import buildbucket
+
+from absl import app
+
+from mcp.server import fastmcp
+
+mcp = fastmcp.FastMCP('chrome-infra-mcp')
+
+
+def main(argv: Sequence[str]) -> None:
+    if len(argv) > 1:
+        raise app.UsageError('Too many command-line arguments.')
+
+    # Only initialize telemetry if the user is opted in. The MCP does not
+    # currently have the ability to show the banner so we need to rely on other
+    # tools to get consent
+    if telemetry.opted_in():
+        telemetry.initialize('chromium.mcp')
+
+    mcp.add_tool(buildbucket.get_build)
+    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.run()
+
+
+if __name__ == '__main__':
+    app.run(main)