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>
This commit is contained in:
Brian Sheedy
2025-08-13 17:03:24 -07:00
committed by LUCI CQ
parent 09f005050e
commit 569c95044f
8 changed files with 104 additions and 187 deletions

View File

@@ -3,12 +3,15 @@
# found in the LICENSE file. # found in the LICENSE file.
"""Tools for interacting with buildbucket""" """Tools for interacting with buildbucket"""
import json import json
import subprocess
import urllib.parse import urllib.parse
from mcp.server import fastmcp from mcp.server import fastmcp
import telemetry import telemetry
import common
BUILDBUCKET_SERVER = 'cr-buildbucket.appspot.com'
tracer = telemetry.get_tracer(__name__) 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'): with tracer.start_as_current_span('chromium.mcp.get_build_status'):
await ctx.info(f'Received request {build_id}') await ctx.info(f'Received request {build_id}')
command = [ request = {'id': build_id}
'prpc', response = await common.run_prpc_call(
'call', ctx, BUILDBUCKET_SERVER, 'buildbucket.v2.Builds.GetBuildStatus',
'cr-buildbucket.appspot.com', request)
'buildbucket.v2.Builds.GetBuildStatus', return json.loads(response)['status']
]
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( 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'): with tracer.start_as_current_span('chromium.mcp.get_build_from_id'):
request = {'id': build_id, 'mask': {'fields': ','.join(fields)}} request = {'id': build_id, 'mask': {'fields': ','.join(fields)}}
command = [ response = await common.run_prpc_call(ctx, BUILDBUCKET_SERVER,
'prpc', 'buildbucket.v2.Builds.GetBuild',
'call', request)
'cr-buildbucket.appspot.com', return response
'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( async def get_build_from_build_number(
@@ -162,26 +134,10 @@ async def get_build_from_build_number(
'fields': ','.join(fields) 'fields': ','.join(fields)
} }
} }
command = [ response = await common.run_prpc_call(ctx, BUILDBUCKET_SERVER,
'prpc', 'buildbucket.v2.Builds.GetBuild',
'call', request)
'cr-buildbucket.appspot.com', return response
'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( async def get_build(
@@ -237,26 +193,10 @@ async def get_build(
""" """
with tracer.start_as_current_span('chromium.mcp.get_build'): with tracer.start_as_current_span('chromium.mcp.get_build'):
await ctx.info(f'Received request {request}') await ctx.info(f'Received request {request}')
command = [ response = await common.run_prpc_call(ctx, BUILDBUCKET_SERVER,
'prpc', 'buildbucket.v2.Builds.GetBuild',
'call', request)
'cr-buildbucket.appspot.com', return response
'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_recent_builds( async def get_recent_builds(
@@ -378,23 +318,7 @@ async def _get_recent_builds(
}, },
'page_size': f'{num_builds}' 'page_size': f'{num_builds}'
} }
command = [ response = await common.run_prpc_call(ctx, BUILDBUCKET_SERVER,
'prpc', 'buildbucket.v2.Builds.SearchBuilds',
'call', request)
'cr-buildbucket.appspot.com', return response
'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

View File

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

39
mcp/common.py Normal file
View File

@@ -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

View File

@@ -19,7 +19,8 @@ async def try_builder_results(
checkout: Location of the current checkout. checkout: Location of the current checkout.
Returns: 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'): with tracer.start_as_current_span('chromium.mcp.try_builder_results'):
command = [ command = [
@@ -70,7 +71,7 @@ async def get_current_changes(
return result.stdout return result.stdout
async def format( async def format_checkout(
ctx: fastmcp.Context, ctx: fastmcp.Context,
checkout: str, checkout: str,
) -> None: ) -> None:

View File

@@ -60,12 +60,12 @@ class GitClTest(unittest.IsolatedAsyncioTestCase):
cwd=self.checkout) cwd=self.checkout)
@mock.patch('subprocess.run') @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' expected_output = 'Formatted 1 file'
mock_subprocess_run.return_value = subprocess.CompletedProcess( mock_subprocess_run.return_value = subprocess.CompletedProcess(
args=[], returncode=0, stdout=expected_output, stderr='') 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) self.assertEqual(output, expected_output)
expected_command = ["git", "cl", "format"] expected_command = ["git", "cl", "format"]

View File

@@ -6,14 +6,17 @@
import json import json
import posixpath import posixpath
import requests import requests
import subprocess
import bs4 import bs4
from mcp.server import fastmcp from mcp.server import fastmcp
import telemetry import telemetry
import common
tracer = telemetry.get_tracer(__name__) tracer = telemetry.get_tracer(__name__)
RESULTDB_SERVER = 'results.api.luci.app'
async def get_non_exonerated_unexpected_results_from_build( async def get_non_exonerated_unexpected_results_from_build(
ctx: fastmcp.Context, ctx: fastmcp.Context,
@@ -50,26 +53,10 @@ async def get_non_exonerated_unexpected_results_from_build(
'read_mask': ('name,resultId,variant,status,statusV2,duration,' 'read_mask': ('name,resultId,variant,status,statusV2,duration,'
'failureReason,summary_html'), 'failureReason,summary_html'),
} }
command = [ response = await common.run_prpc_call(
'prpc', ctx, RESULTDB_SERVER, 'luci.resultdb.v1.ResultDB.QueryTestResults',
'call', request)
'results.api.luci.app', return response
'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
async def expand_summary_html( async def expand_summary_html(
@@ -137,26 +124,10 @@ async def get_test_level_text_artifact(
request = { request = {
'name': artifact_name, 'name': artifact_name,
} }
command = [ prpc_response = await common.run_prpc_call(
'prpc', ctx, RESULTDB_SERVER, 'luci.resultdb.v1.ResultDB.GetArtifact',
'call', request)
'results.api.luci.app', response = json.loads(prpc_response)
'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)
content_type = response['contentType'] content_type = response['contentType']
if 'text/plain' not in content_type: if 'text/plain' not in content_type:

View File

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

View File

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