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..852a10969a28 --- /dev/null +++ b/awscli/customizations/mcp/_mcp_server.py @@ -0,0 +1,75 @@ +import json +import logging +import re +import shlex +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 awscli.clidriver import create_clidriver +from awscli.compat import StringIO +from awscli.help import PosixHelpRenderer + +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 + + +_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( + "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: + 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: + return json.loads(output) + except json.JSONDecodeError: + return {"output": output} + + 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)