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)