From 75449d2000f1bbfdae9e5cfbc13bedac876fa6b0 Mon Sep 17 00:00:00 2001 From: Struan Shrimpton Date: Mon, 21 Jul 2025 10:33:13 -0700 Subject: [PATCH] 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 Reviewed-by: Gavin Mak Reviewed-by: Erik Staab Commit-Queue: Struan Shrimpton --- infra_lib/telemetry/__init__.py | 20 +++ mcp/.vpython3 | 162 ++++++++++++++++++++ mcp/OWNERS | 3 + mcp/buildbucket.py | 258 ++++++++++++++++++++++++++++++++ mcp/buildbucket_test.py | 204 +++++++++++++++++++++++++ mcp/server.py | 45 ++++++ 6 files changed, 692 insertions(+) create mode 100644 mcp/.vpython3 create mode 100644 mcp/OWNERS create mode 100644 mcp/buildbucket.py create mode 100644 mcp/buildbucket_test.py create mode 100755 mcp/server.py diff --git a/infra_lib/telemetry/__init__.py b/infra_lib/telemetry/__init__.py index 5d7bbed635..e707d691d9 100644 --- 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 \ No newline at end of file diff --git a/mcp/.vpython3 b/mcp/.vpython3 new file mode 100644 index 0000000000..9994575626 --- /dev/null +++ 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" +> diff --git a/mcp/OWNERS b/mcp/OWNERS new file mode 100644 index 0000000000..29a5b4675e --- /dev/null +++ b/mcp/OWNERS @@ -0,0 +1,3 @@ +estaab@google.com +jiesheng@google.com +sshrimp@google.com \ No newline at end of file diff --git a/mcp/buildbucket.py b/mcp/buildbucket.py new file mode 100644 index 0000000000..bdc83901cd --- /dev/null +++ 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/ + + 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": "", + "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/ + + 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": "", + "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////infra + Builds can also be in the form: + https://ci.chromium.org/b/ + + 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": "", + "builder": { + "bucket": "", + "builder": "", + "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": "", + "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 diff --git a/mcp/buildbucket_test.py b/mcp/buildbucket_test.py new file mode 100644 index 0000000000..5e3178dc3c --- /dev/null +++ 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() diff --git a/mcp/server.py b/mcp/server.py new file mode 100755 index 0000000000..0379aa0482 --- /dev/null +++ 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)