From 32adf2315ae6174b4a37e74f6a4f8599083b58b1 Mon Sep 17 00:00:00 2001 From: Tony Xing <16152581+tonyxwz@users.noreply.github.com> Date: Sun, 10 Aug 2025 17:13:07 +0200 Subject: [PATCH 1/3] feat: implement AWS CLI MCP as a customization This is an MCP server that can execute any AWS CLI command supported by the current user's installed aws cli. --- awscli/customizations/mcp/__init__.py | 1 + awscli/customizations/mcp/_initialize.py | 10 ++++ awscli/customizations/mcp/_mcp_command.py | 17 ++++++ awscli/customizations/mcp/_mcp_server.py | 58 +++++++++++++++++++ awscli/customizations/mcp/_mcp_subcommands.py | 44 ++++++++++++++ awscli/handlers.py | 2 + 6 files changed, 132 insertions(+) create mode 100644 awscli/customizations/mcp/__init__.py create mode 100644 awscli/customizations/mcp/_initialize.py create mode 100644 awscli/customizations/mcp/_mcp_command.py create mode 100644 awscli/customizations/mcp/_mcp_server.py create mode 100644 awscli/customizations/mcp/_mcp_subcommands.py diff --git a/awscli/customizations/mcp/__init__.py b/awscli/customizations/mcp/__init__.py new file mode 100644 index 000000000000..343cf60d15c3 --- /dev/null +++ b/awscli/customizations/mcp/__init__.py @@ -0,0 +1 @@ +from ._initialize import awscli_initialize as awscli_initialize diff --git a/awscli/customizations/mcp/_initialize.py b/awscli/customizations/mcp/_initialize.py new file mode 100644 index 000000000000..a8b1741f871e --- /dev/null +++ b/awscli/customizations/mcp/_initialize.py @@ -0,0 +1,10 @@ +from ._mcp_command import Mcp + + +def awscli_initialize(cli): + """The entry point to the AWS CLI MCP server.""" + cli.register('building-command-table.main', add_mcp_command) + + +def add_mcp_command(command_table, session, **kwargs): + command_table['mcp-server'] = Mcp(session) diff --git a/awscli/customizations/mcp/_mcp_command.py b/awscli/customizations/mcp/_mcp_command.py new file mode 100644 index 000000000000..9632cab23589 --- /dev/null +++ b/awscli/customizations/mcp/_mcp_command.py @@ -0,0 +1,17 @@ +from typing import override + +from awscli.customizations.commands import BasicCommand + +from ._mcp_subcommands import StartCommand + + +class Mcp(BasicCommand): + NAME = "mcp-server" + DESCRIPTION = "Use AWS CPI as a MCP server" + SYNOPSIS = "aws mcp [ ...]" + SUBCOMMANDS = [{'name': 'start', 'command_class': StartCommand}] + + @override + def _run_main(self, parsed_args, parsed_globals): + if parsed_args.subcommand is None: + self._raise_usage_error() diff --git a/awscli/customizations/mcp/_mcp_server.py b/awscli/customizations/mcp/_mcp_server.py new file mode 100644 index 000000000000..46c143ad931d --- /dev/null +++ b/awscli/customizations/mcp/_mcp_server.py @@ -0,0 +1,58 @@ +import json +import logging +import shlex +import sys +from contextlib import redirect_stdout + +from mcp.server.fastmcp import FastMCP +from mcp.server.fastmcp.exceptions import ToolError +from mcp.server.fastmcp.server import Context +from pydantic import BaseModel + +from awscli.clidriver import create_clidriver +from awscli.compat import StringIO + +LOG = logging.getLogger(__name__) + + +def get_aws_cli_mcp_server(*, port: str, host: str) -> FastMCP: + mcp_server = FastMCP( + name="AWS CLI MCP server", + instructions="MCP server to execute AWS CLI commands", + port=int(port), + host=host, + ) + mcp_server.add_tool(run_aws_cli) + return mcp_server + + +def run_aws_cli(ctx: Context, aws_cli: str) -> dict: + """Run the aws cli command and return the JSON response.""" + + tokens = shlex.split(aws_cli) + if set(tokens) & set(("|", ">", ">>", "||", "&&", "&")): + raise ToolError( + "run_aws_cli cannot handle pipelines. You can only run one AWS CLI command at one time." + ) + + sub_command = tokens[1:] + LOG.debug(f"Running AWS CLI command {sub_command!r}") + try: + output = StringIO() + with redirect_stdout(output): + clidriver = create_clidriver(sub_command) + return_code = clidriver.main(sub_command) + try: + parsed_json = json.load(output) + return parsed_json + except json.JSONDecodeError: + return {"output": output.getvalue()} + + except SystemExit as sys_exit: + if sys_exit.code != 0: + raise ToolError("Non 0 return code from aws cli") + else: + raise ToolError(f"Unknown error occured when executing {aws_cli}") + + except BaseException as e: + raise ToolError(str(e)) diff --git a/awscli/customizations/mcp/_mcp_subcommands.py b/awscli/customizations/mcp/_mcp_subcommands.py new file mode 100644 index 000000000000..7bf72ec92754 --- /dev/null +++ b/awscli/customizations/mcp/_mcp_subcommands.py @@ -0,0 +1,44 @@ +import logging + +from awscli.customizations.commands import BasicCommand + +from ._mcp_server import get_aws_cli_mcp_server + +LOG = logging.getLogger(__name__) + + +class StartCommand(BasicCommand): + NAME = "start" + DESCRIPTION = "Start the AWS CLI MCP server." + USAGE = "aws mcp-server start" + ARG_TABLE = [ + { + "name": "transport", + "nargs": "?", + "choices": ["stdio", "streamable-http"], + "default": "stdio", + 'positional_arg': False, + 'help_text': "Specify the transport to run the JSON-RPC server used by MCP. (default: stdio)", + }, + { + "name": "port", + "nargs": "?", + "default": "8080", + 'positional_arg': False, + 'help_text': "Specify the port that the MCP server listens to if protocol is streamable-http", + }, + { + "name": "host", + "nargs": "?", + "default": "127.0.0.1", + 'positional_arg': False, + 'help_text': "Specify the host that the MCP server listens to if protocol is streamable-http", + }, + ] + + def _run_main(self, parsed_args, parsed_globals): + mcp_server = get_aws_cli_mcp_server( + port=parsed_args.port, + host=parsed_args.host, + ) + mcp_server.run(transport=parsed_args.transport) diff --git a/awscli/handlers.py b/awscli/handlers.py index ef3abe07fc3a..31497e187c4c 100644 --- a/awscli/handlers.py +++ b/awscli/handlers.py @@ -94,6 +94,7 @@ from awscli.customizations.kms import register_fix_kms_create_grant_docs from awscli.customizations.lightsail import initialize as lightsail_initialize from awscli.customizations.logs import register_logs_commands +from awscli.customizations.mcp import awscli_initialize as mcp_initialize from awscli.customizations.opsworks import initialize as opsworks_init from awscli.customizations.opsworkscm import register_alias_opsworks_cm from awscli.customizations.paginate import register_pagination @@ -237,3 +238,4 @@ def awscli_initialize(event_handlers): register_kinesis_list_streams_pagination_backcompat(event_handlers) register_quicksight_asset_bundle_customizations(event_handlers) register_ec2_instance_connect_commands(event_handlers) + mcp_initialize(event_handlers) From 3812607910ca61af81cf96431fac7636e61541f7 Mon Sep 17 00:00:00 2001 From: Tony Xing <16152581+tonyxwz@users.noreply.github.com> Date: Sun, 10 Aug 2025 20:03:18 +0200 Subject: [PATCH 2/3] fix: show help in aws cli --- awscli/customizations/mcp/_mcp_server.py | 23 ++++++++++++++++++++--- 1 file changed, 20 insertions(+), 3 deletions(-) diff --git a/awscli/customizations/mcp/_mcp_server.py b/awscli/customizations/mcp/_mcp_server.py index 46c143ad931d..81c78f866f44 100644 --- a/awscli/customizations/mcp/_mcp_server.py +++ b/awscli/customizations/mcp/_mcp_server.py @@ -1,16 +1,16 @@ import json import logging +import re import shlex -import sys from contextlib import redirect_stdout from mcp.server.fastmcp import FastMCP from mcp.server.fastmcp.exceptions import ToolError from mcp.server.fastmcp.server import Context -from pydantic import BaseModel from awscli.clidriver import create_clidriver from awscli.compat import StringIO +from awscli.help import PosixHelpRenderer LOG = logging.getLogger(__name__) @@ -26,9 +26,25 @@ def get_aws_cli_mcp_server(*, port: str, host: str) -> FastMCP: return mcp_server +_CONTROL_SEQUENCE = re.compile("\x1b\\[\\d+m") + + +def _remove_control_sequence(content: str): + return _CONTROL_SEQUENCE.sub("", content) + + +def _patch_aws_cli_pager(buf: StringIO): + def patched_send_output_to_pager(self, output): + content = output.decode('utf-8') + buf.write(_remove_control_sequence(content) + "\n") + buf.flush() + return + + PosixHelpRenderer._send_output_to_pager = patched_send_output_to_pager + + def run_aws_cli(ctx: Context, aws_cli: str) -> dict: """Run the aws cli command and return the JSON response.""" - tokens = shlex.split(aws_cli) if set(tokens) & set(("|", ">", ">>", "||", "&&", "&")): raise ToolError( @@ -39,6 +55,7 @@ def run_aws_cli(ctx: Context, aws_cli: str) -> dict: LOG.debug(f"Running AWS CLI command {sub_command!r}") try: output = StringIO() + _patch_aws_cli_pager(output) with redirect_stdout(output): clidriver = create_clidriver(sub_command) return_code = clidriver.main(sub_command) From c42e506038c33aa7e3896646d0f4218ac097d5d4 Mon Sep 17 00:00:00 2001 From: Tony Xing <16152581+tonyxwz@users.noreply.github.com> Date: Sun, 10 Aug 2025 20:29:44 +0200 Subject: [PATCH 3/3] fix parsing output --- awscli/customizations/mcp/_mcp_server.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/awscli/customizations/mcp/_mcp_server.py b/awscli/customizations/mcp/_mcp_server.py index 81c78f866f44..852a10969a28 100644 --- a/awscli/customizations/mcp/_mcp_server.py +++ b/awscli/customizations/mcp/_mcp_server.py @@ -54,16 +54,16 @@ def run_aws_cli(ctx: Context, aws_cli: str) -> dict: sub_command = tokens[1:] LOG.debug(f"Running AWS CLI command {sub_command!r}") try: - output = StringIO() - _patch_aws_cli_pager(output) - with redirect_stdout(output): + buf = StringIO() + _patch_aws_cli_pager(buf) + with redirect_stdout(buf): clidriver = create_clidriver(sub_command) return_code = clidriver.main(sub_command) + output = buf.getvalue() try: - parsed_json = json.load(output) - return parsed_json + return json.loads(output) except json.JSONDecodeError: - return {"output": output.getvalue()} + return {"output": output} except SystemExit as sys_exit: if sys_exit.code != 0: