agent-claw/agentclaw/app/agent_claw_main/tools/mcp_tools.py

"""MCP connection management tool — add/remove/enable/disable user MCP servers."""
import boto3
from strands import tool

USERS_TABLE_NAME = 'agent-claw-users'

# Set per-invocation by main.py
_current_actor_id: str = ''


@tool
def manage_mcp_connection(action: str, name: str = '', url: str = '',
                          auth_type: str = 'none', cognito_token_url: str = '',
                          client_id: str = '', client_secret: str = '',
                          scope: str = '', token: str = '') -> str:
    """Add, remove, enable, disable, or list MCP server connections.

    Actions: add, remove, enable, disable, list

    For add with auth_type=oauth_client_credentials, provide:
      - cognito_token_url: Cognito token endpoint
      - client_id: OAuth client ID
      - client_secret: Secret value (stored securely in SSM, not in database)
      - scope: Space-separated OAuth scopes

    For add with auth_type=bearer, provide:
      - token: Bearer token value (stored securely in SSM, not in database)

    For add with auth_type=none, only name and url are required.

    Args:
        action: One of "add", "remove", "enable", "disable", "list".
        name: Connection name (required for add/remove/enable/disable).
        url: MCP server URL (required for add).
        auth_type: One of "none", "bearer", "oauth_client_credentials".
        cognito_token_url: Token endpoint for oauth_client_credentials.
        client_id: OAuth client ID for oauth_client_credentials.
        client_secret: OAuth client secret (will be stored in SSM).
        scope: OAuth scopes for oauth_client_credentials.
        token: Bearer token (will be stored in SSM).
    """
    actor_id = _current_actor_id
    if not actor_id:
        return 'Cannot determine actor_id.'

    ddb = boto3.resource('dynamodb', region_name='us-east-1')
    table = ddb.Table(USERS_TABLE_NAME)

    if action == 'list':
        resp = table.get_item(Key={'actor_id': actor_id})
        connections = resp.get('Item', {}).get('services', {}).get('mcp_connections', [])
        if not connections:
            return 'No MCP connections configured.'
        lines = []
        for c in connections:
            status = '✓' if c.get('enabled', True) else '✗'
            lines.append(f"  {status} {c['name']}: {c['url']} (auth: {c.get('auth_type', 'none')})")
        return 'MCP connections:\n' + '\n'.join(lines)

    if not name:
        return 'name is required for this action.'

    if action == 'add':
        if not url:
            return 'url is required for add.'
        if auth_type not in ('none', 'bearer', 'oauth_client_credentials'):
            return f'Invalid auth_type: {auth_type}. Use none, bearer, or oauth_client_credentials.'

        ssm = boto3.client('ssm', region_name='us-east-1')
        safe_actor = actor_id.replace(':', '-')
        ssm_prefix = f'/agent-claw/mcp/{safe_actor}/{name}'

        conn = {'name': name, 'url': url, 'auth_type': auth_type, 'enabled': True}

        if auth_type == 'oauth_client_credentials':
            if not cognito_token_url or not client_id or not client_secret:
                return 'oauth_client_credentials requires cognito_token_url, client_id, and client_secret.'
            ssm.put_parameter(Name=f'{ssm_prefix}/client-secret', Value=client_secret,
                              Type='SecureString', Overwrite=True)
            conn['cognito_token_url'] = cognito_token_url
            conn['client_id'] = client_id
            conn['client_secret_ssm'] = f'{ssm_prefix}/client-secret'
            conn['scope'] = scope
        elif auth_type == 'bearer':
            if not token:
                return 'bearer auth requires token.'
            ssm.put_parameter(Name=f'{ssm_prefix}/token', Value=token,
                              Type='SecureString', Overwrite=True)
            conn['token_ssm'] = f'{ssm_prefix}/token'

        # Upsert into mcp_connections list
        resp = table.get_item(Key={'actor_id': actor_id})
        services = resp.get('Item', {}).get('services', {})
        connections = services.get('mcp_connections', [])
        connections = [c for c in connections if c['name'] != name]
        connections.append(conn)

        table.update_item(
            Key={'actor_id': actor_id},
            UpdateExpression='SET services = if_not_exists(services, :empty), services.mcp_connections = :conns',
            ExpressionAttributeValues={':conns': connections, ':empty': {}},
        )
        return f'MCP connection "{name}" added ({auth_type} auth). It will be available on your next message.'

    elif action == 'remove':
        resp = table.get_item(Key={'actor_id': actor_id})
        connections = resp.get('Item', {}).get('services', {}).get('mcp_connections', [])
        found = [c for c in connections if c['name'] == name]
        if not found:
            return f'No connection named "{name}" found.'

        # Clean up SSM secrets
        ssm = boto3.client('ssm', region_name='us-east-1')
        safe_actor = actor_id.replace(':', '-')
        for key in ('client-secret', 'token'):
            try:
                ssm.delete_parameter(Name=f'/agent-claw/mcp/{safe_actor}/{name}/{key}')
            except ssm.exceptions.ParameterNotFound:
                pass

        connections = [c for c in connections if c['name'] != name]
        table.update_item(
            Key={'actor_id': actor_id},
            UpdateExpression='SET services.mcp_connections = :conns',
            ExpressionAttributeValues={':conns': connections},
        )
        return f'MCP connection "{name}" removed.'

    elif action in ('enable', 'disable'):
        resp = table.get_item(Key={'actor_id': actor_id})
        connections = resp.get('Item', {}).get('services', {}).get('mcp_connections', [])
        updated = False
        for c in connections:
            if c['name'] == name:
                c['enabled'] = (action == 'enable')
                updated = True
                break
        if not updated:
            return f'No connection named "{name}" found.'
        table.update_item(
            Key={'actor_id': actor_id},
            UpdateExpression='SET services.mcp_connections = :conns',
            ExpressionAttributeValues={':conns': connections},
        )
        return f'MCP connection "{name}" {action}d.'

    else:
        return f'Unknown action: {action}. Use add, remove, enable, disable, or list.'