Implement MCP server

awscli/customizations/mcp/__init__.py

@@ -0,0 +1 @@
+from ._initialize import awscli_initialize as awscli_initialize

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)

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 <Command> [<Arg> ...]"
+    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()

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

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)

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)