diff --git a/src/iam-mcp-server/awslabs/iam_mcp_server/server.py b/src/iam-mcp-server/awslabs/iam_mcp_server/server.py index 8fc7f7a1c4..d98c3d3c55 100644 --- a/src/iam-mcp-server/awslabs/iam_mcp_server/server.py +++ b/src/iam-mcp-server/awslabs/iam_mcp_server/server.py @@ -17,7 +17,7 @@ import argparse import json from awslabs.iam_mcp_server.aws_client import get_iam_client -from awslabs.iam_mcp_server.context import Context +from awslabs.iam_mcp_server.context import Context as ServerContext from awslabs.iam_mcp_server.errors import IamClientError, IamValidationError, handle_iam_error from awslabs.iam_mcp_server.models import ( AccessKey, @@ -37,8 +37,7 @@ UsersListResponse, ) from loguru import logger -from mcp.server.fastmcp import FastMCP -from mcp.types import CallToolResult +from mcp.server.fastmcp import Context, FastMCP from pydantic import Field from typing import Any, Dict, List, Optional, Union @@ -86,7 +85,7 @@ @mcp.tool() async def list_users( - ctx: CallToolResult, + ctx: Context, path_prefix: Optional[str] = Field( description='Path prefix to filter users (e.g., "/division_abc/")', default=None ), @@ -154,7 +153,7 @@ async def list_users( @mcp.tool() async def get_user( - ctx: CallToolResult, user_name: str = Field(description='The name of the IAM user to retrieve') + ctx: Context, user_name: str = Field(description='The name of the IAM user to retrieve') ) -> UserDetailsResponse: """Get detailed information about a specific IAM user. @@ -242,7 +241,7 @@ async def get_user( @mcp.tool() async def create_user( - ctx: CallToolResult, + ctx: Context, user_name: str = Field(description='The name of the new IAM user'), path: str = Field(description='The path for the user', default='/'), permissions_boundary: Optional[str] = Field( @@ -273,7 +272,7 @@ async def create_user( logger.info(f'Creating IAM user: {user_name}') # Check if server is in read-only mode - if Context.is_readonly(): + if ServerContext.is_readonly(): raise IamClientError('Cannot create user: server is running in read-only mode') if not user_name: @@ -332,7 +331,7 @@ async def delete_user( """ try: # Check if server is in read-only mode - if Context.is_readonly(): + if ServerContext.is_readonly(): raise IamClientError('Cannot delete user: server is running in read-only mode') iam = get_iam_client() @@ -448,7 +447,7 @@ async def create_role( """ try: # Check if server is in read-only mode - if Context.is_readonly(): + if ServerContext.is_readonly(): raise IamClientError('Cannot create role: server is running in read-only mode') iam = get_iam_client() @@ -633,7 +632,7 @@ async def attach_user_policy( """ try: # Check if server is in read-only mode - if Context.is_readonly(): + if ServerContext.is_readonly(): raise IamClientError('Cannot attach policy: server is running in read-only mode') iam = get_iam_client() @@ -666,7 +665,7 @@ async def detach_user_policy( """ try: # Check if server is in read-only mode - if Context.is_readonly(): + if ServerContext.is_readonly(): raise IamClientError('Cannot detach policy: server is running in read-only mode') iam = get_iam_client() @@ -697,7 +696,7 @@ async def create_access_key( """ try: # Check if server is in read-only mode - if Context.is_readonly(): + if ServerContext.is_readonly(): raise IamClientError('Cannot create access key: server is running in read-only mode') iam = get_iam_client() @@ -737,7 +736,7 @@ async def delete_access_key( """ try: # Check if server is in read-only mode - if Context.is_readonly(): + if ServerContext.is_readonly(): raise IamClientError('Cannot delete access key: server is running in read-only mode') iam = get_iam_client() @@ -840,7 +839,7 @@ async def list_groups( Returns: GroupsListResponse containing list of groups and metadata """ - if Context.is_readonly(): + if ServerContext.is_readonly(): # List operations are allowed in read-only mode pass @@ -896,7 +895,7 @@ async def get_group( Returns: GroupDetailsResponse containing comprehensive group information """ - if Context.is_readonly(): + if ServerContext.is_readonly(): # Get operations are allowed in read-only mode pass @@ -962,7 +961,7 @@ async def create_group( Returns: CreateGroupResponse containing the created group details """ - if Context.is_readonly(): + if ServerContext.is_readonly(): raise IamValidationError('Cannot create group in read-only mode') try: @@ -1003,7 +1002,7 @@ async def delete_group( Returns: Dictionary containing deletion status """ - if Context.is_readonly(): + if ServerContext.is_readonly(): raise IamValidationError('Cannot delete group in read-only mode') try: @@ -1048,7 +1047,7 @@ async def add_user_to_group( Returns: GroupMembershipResponse containing operation status """ - if Context.is_readonly(): + if ServerContext.is_readonly(): raise IamValidationError('Cannot add user to group in read-only mode') try: @@ -1079,7 +1078,7 @@ async def remove_user_from_group( Returns: GroupMembershipResponse containing operation status """ - if Context.is_readonly(): + if ServerContext.is_readonly(): raise IamValidationError('Cannot remove user from group in read-only mode') try: @@ -1110,7 +1109,7 @@ async def attach_group_policy( Returns: GroupPolicyAttachmentResponse containing operation status """ - if Context.is_readonly(): + if ServerContext.is_readonly(): raise IamValidationError('Cannot attach policy to group in read-only mode') try: @@ -1141,7 +1140,7 @@ async def detach_group_policy( Returns: GroupPolicyAttachmentResponse containing operation status """ - if Context.is_readonly(): + if ServerContext.is_readonly(): raise IamValidationError('Cannot detach policy from group in read-only mode') try: @@ -1193,7 +1192,7 @@ async def put_user_policy( logger.info(f'Creating/updating inline policy {policy_name} for user: {user_name}') # Check if server is in read-only mode - if Context.is_readonly(): + if ServerContext.is_readonly(): raise IamClientError( 'Cannot create/update inline policy: server is running in read-only mode' ) @@ -1299,7 +1298,7 @@ async def delete_user_policy( logger.info(f'Deleting inline policy {policy_name} from user: {user_name}') # Check if server is in read-only mode - if Context.is_readonly(): + if ServerContext.is_readonly(): raise IamClientError( 'Cannot delete inline policy: server is running in read-only mode' ) @@ -1355,7 +1354,7 @@ async def put_role_policy( logger.info(f'Creating/updating inline policy {policy_name} for role: {role_name}') # Check if server is in read-only mode - if Context.is_readonly(): + if ServerContext.is_readonly(): raise IamClientError( 'Cannot create/update inline policy: server is running in read-only mode' ) @@ -1461,7 +1460,7 @@ async def delete_role_policy( logger.info(f'Deleting inline policy {policy_name} from role: {role_name}') # Check if server is in read-only mode - if Context.is_readonly(): + if ServerContext.is_readonly(): raise IamClientError( 'Cannot delete inline policy: server is running in read-only mode' ) @@ -1583,7 +1582,7 @@ def main(): # Set read-only mode if specified if args.readonly: - Context.set_readonly(True) + ServerContext.set_readonly(True) logger.info('Server started in READ-ONLY mode - all mutating operations are disabled') else: logger.info('Server started in FULL ACCESS mode') diff --git a/test-iam-ctx-fix.py b/test-iam-ctx-fix.py new file mode 100755 index 0000000000..bcd8969cbe --- /dev/null +++ b/test-iam-ctx-fix.py @@ -0,0 +1,92 @@ +#!/usr/bin/env python3 +"""Test script to verify the ctx parameter fix.""" + +import asyncio +import os +import sys + +# Add the src directory to the path so we can import the modified server +sys.path.insert(0, '/home/plex/development/repos/aws/mcp/src/iam-mcp-server') + +from mcp import ClientSession, StdioServerParameters +from mcp.client.stdio import stdio_client + + +async def main(): + # Use the local modified version via python -m + server_params = StdioServerParameters( + command="python", + args=["-m", "awslabs.iam_mcp_server.server", "--readonly"], + env={ + "AWS_PROFILE": os.environ.get("AWS_PROFILE", "default"), + "AWS_REGION": os.environ.get("AWS_REGION", "us-east-1"), + "PYTHONPATH": "/home/plex/development/repos/aws/mcp/src/iam-mcp-server" + } + ) + + print("=" * 70) + print("Testing IAM MCP Server ctx Parameter Fix") + print("=" * 70) + print() + + async with stdio_client(server_params) as (read, write): + async with ClientSession(read, write) as session: + await session.initialize() + print("✓ Connected to server\n") + + # Test 1: List tools and check schema + print("Test 1: Check list_users tool schema") + print("-" * 50) + tools_result = await session.list_tools() + list_users_tool = [t for t in tools_result.tools if t.name == "list_users"][0] + + schema = list_users_tool.inputSchema + if "properties" in schema: + props = schema["properties"] + if "ctx" in props: + print("❌ FAIL: ctx is still in the schema!") + print(f" Properties: {list(props.keys())}") + return False + else: + print("✅ PASS: ctx is not in the schema") + print(f" Properties: {list(props.keys())}") + print() + + # Test 2: Call the tool + print("Test 2: Call list_users tool") + print("-" * 50) + result = await session.call_tool("list_users", arguments={"max_items": 3}) + + if result.isError: + print(f"❌ FAIL: Tool returned error") + print(f" Error: {result.content[0].text if result.content else 'Unknown'}") + return False + else: + print("✅ PASS: Tool call succeeded!") + if result.content: + import json + from mcp import types + content_block = result.content[0] + if isinstance(content_block, types.TextContent): + data = json.loads(content_block.text) + users = data.get('Users', []) + print(f" Retrieved {len(users)} users") + if users: + print(f" First user: {users[0].get('UserName')}") + + print() + print("=" * 70) + print("✅ ALL TESTS PASSED - Fix verified!") + print("=" * 70) + return True + + +if __name__ == "__main__": + try: + success = asyncio.run(main()) + sys.exit(0 if success else 1) + except Exception as e: + print(f"\n❌ Test failed with exception: {e}") + import traceback + traceback.print_exc() + sys.exit(1)