awslabs.prometheus-mcp-server 0.1.1__py3-none-any.whl → 0.2.1__py3-none-any.whl

awslabs/prometheus_mcp_server/server.py

@@ -29,25 +29,22 @@ from awslabs.prometheus_mcp_server.consts import (
     DEFAULT_SERVICE_NAME,
     ENV_AWS_PROFILE,
     ENV_AWS_REGION,
-    ENV_AWS_SERVICE_NAME,
     ENV_LOG_LEVEL,
-    ENV_PROMETHEUS_URL,
     SERVER_INSTRUCTIONS,
 )
 from awslabs.prometheus_mcp_server.models import (
     MetricsList,
-    PrometheusConfig,
     ServerInfo,
 )
 from botocore.auth import SigV4Auth
 from botocore.awsrequest import AWSRequest
+from botocore.config import Config
 from botocore.exceptions import ClientError, NoCredentialsError
 from dotenv import load_dotenv
 from loguru import logger
 from mcp.server.fastmcp import Context, FastMCP
 from pydantic import Field
 from typing import Any, Dict, Optional
-from urllib.parse import urlparse
 
 
 # Configure loguru
@@ -55,347 +52,341 @@ logger.remove()
 logger.add(sys.stderr, level=os.getenv(ENV_LOG_LEVEL, 'INFO'))
 
 
-def parse_arguments():
-    """Parse command line arguments."""
-    parser = argparse.ArgumentParser(description='Prometheus MCP Server')
-    parser.add_argument('--profile', type=str, help='AWS profile name to use')
-    parser.add_argument('--region', type=str, help='AWS region to use')
-    parser.add_argument('--url', type=str, help='Prometheus URL')
-    parser.add_argument('--config', type=str, help='Path to configuration file')
-    parser.add_argument('--debug', action='store_true', help='Enable debug logging')
-    return parser.parse_args()
-
-
-def load_config(args):
-    """Load configuration from file, environment variables, and command line arguments."""
-    # Load .env file if it exists
-    load_dotenv()
-
-    # Initialize config with default values
-    config_data = {
-        'aws_profile': None,
-        'aws_region': DEFAULT_AWS_REGION,
-        'prometheus_url': '',
-        'service_name': DEFAULT_SERVICE_NAME,
-        'max_retries': DEFAULT_MAX_RETRIES,
-        'retry_delay': DEFAULT_RETRY_DELAY,
-    }
-
-    # Load from config file if specified
-    if args.config and os.path.exists(args.config):
-        try:
-            with open(args.config, 'r') as f:
-                file_config = json.load(f)
-                config_data.update(file_config)
-            logger.info(f'Loaded configuration from {args.config}')
-        except Exception as e:
-            logger.error(f'Error loading config file: {e}')
-
-    # Override with environment variables
-    if os.getenv(ENV_AWS_PROFILE):
-        config_data['aws_profile'] = os.getenv(ENV_AWS_PROFILE)
-    if os.getenv(ENV_AWS_REGION):
-        config_data['aws_region'] = os.getenv(ENV_AWS_REGION)
-    if os.getenv(ENV_PROMETHEUS_URL):
-        config_data['prometheus_url'] = os.getenv(ENV_PROMETHEUS_URL)
-    if os.getenv(ENV_AWS_SERVICE_NAME):
-        config_data['service_name'] = os.getenv(ENV_AWS_SERVICE_NAME)
-
-    # Override with command line arguments
-    if args.profile:
-        config_data['aws_profile'] = args.profile
-    if args.region:
-        config_data['aws_region'] = args.region
-    if args.url:
-        config_data['prometheus_url'] = args.url
-    if args.debug:
-        logger.level('DEBUG')
-        logger.debug('Debug logging enabled')
-
-    return config_data
-
-
-def setup_environment(config):
-    """Setup and validate environment variables."""
-    logger.info('Setting up environment...')
-
-    # Validate Prometheus URL
-    if not config['prometheus_url']:
-        logger.error(
-            'ERROR: Prometheus URL not configured. Please set using --url parameter or PROMETHEUS_URL environment variable.'
-        )
-        return False
+class ConfigManager:
+    """Configuration management for the application."""
 
-    try:
-        parsed_url = urlparse(config['prometheus_url'])
-        if not all([parsed_url.scheme, parsed_url.netloc]):
-            logger.error(f'ERROR: Invalid Prometheus URL format: {config["prometheus_url"]}')
-            logger.error('URL must include scheme (https://) and hostname')
-            return False
+    @staticmethod
+    def parse_arguments():
+        """Parse command line arguments."""
+        parser = argparse.ArgumentParser(description='Prometheus MCP Server')
+        parser.add_argument('--profile', type=str, help='AWS profile name to use')
+        parser.add_argument('--region', type=str, help='AWS region to use')
+        parser.add_argument('--url', type=str, help='Prometheus URL to use')
+        parser.add_argument('--debug', action='store_true', help='Enable debug logging')
+        return parser.parse_args()
 
-        # Verify URL points to AWS Prometheus
-        if not (
-            parsed_url.netloc.endswith('.amazonaws.com') and 'aps-workspaces' in parsed_url.netloc
-        ):
-            logger.warning(
-                f"WARNING: URL doesn't appear to be an AWS Managed Prometheus endpoint: {config['prometheus_url']}"
-            )
-            logger.warning(
-                'Expected format: https://aps-workspaces.[region].amazonaws.com/workspaces/ws-[id]'
-            )
-    except Exception as e:
-        logger.error(f'ERROR: Error parsing Prometheus URL: {e}')
-        return False
+    @staticmethod
+    def setup_basic_config(args):
+        """Setup basic configuration from command line arguments and environment variables."""
+        # Load .env file if it exists
+        load_dotenv()
 
-    logger.info('Prometheus configuration:')
-    logger.info(f'  Server URL: {config["prometheus_url"]}')
-    logger.info(f'  AWS Region: {config["aws_region"]}')
+        # Set debug logging if requested
+        if args.debug:
+            logger.level('DEBUG')
+            logger.debug('Debug logging enabled')
 
-    # Test AWS credentials
-    try:
-        if not config['aws_region']:
-            logger.error(
-                'ERROR: AWS region not configured. Please set using --region parameter or AWS_REGION environment variable.'
-            )
-            return False
-
-        logger.info(f'  AWS Region: {config["aws_region"]}')
-
-        # Create session with profile if specified
-        if config['aws_profile']:
-            logger.info(f'  Using AWS Profile: {config["aws_profile"]}')
-            session = boto3.Session(
-                profile_name=config['aws_profile'], region_name=config['aws_region']
-            )
-        else:
-            logger.info('  Using default AWS credentials')
-            session = boto3.Session(region_name=config['aws_region'])
-
-        credentials = session.get_credentials()
-        if credentials:
-            logger.info('  AWS Credentials: Available')
-            if credentials.token:
-                logger.info('  Credential Type: Temporary (includes session token)')
-            else:
-                logger.info('  Credential Type: Long-term')
+        # Get region, profile, and URL from args or environment
+        region = args.region or os.getenv(ENV_AWS_REGION) or DEFAULT_AWS_REGION
+        profile = args.profile or os.getenv(ENV_AWS_PROFILE)
+        url = args.url or os.getenv('PROMETHEUS_URL')
 
-            # Test if credentials have necessary permissions
-            try:
-                sts = session.client('sts')
-                identity = sts.get_caller_identity()
-                logger.info(f'  AWS Identity: {identity["Arn"]}')
-            except ClientError as e:
-                logger.warning(f'WARNING: Could not verify AWS identity: {e}')
-                logger.warning(
-                    'This may indicate insufficient permissions for STS:GetCallerIdentity'
-                )
-        else:
-            logger.error('ERROR: AWS Credentials not found')
-            logger.error('Please configure AWS credentials using:')
-            logger.error('  - AWS CLI: aws configure')
-            logger.error('  - Environment variables: AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY')
-            logger.error('  - Or specify a profile with --profile')
-            return False
-    except NoCredentialsError:
-        logger.error('ERROR: AWS credentials not found')
-        logger.error('Please configure AWS credentials using:')
-        logger.error('  - AWS CLI: aws configure')
-        logger.error('  - Environment variables: AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY')
-        logger.error('  - Or specify a profile with --profile')
-        return False
-    except Exception as e:
-        logger.error(f'ERROR: Error setting up AWS session: {e}')
-        return False
+        return {'region': region, 'profile': profile, 'url': url}
 
-    return True
 
+class AWSCredentials:
+    """AWS credentials management."""
 
-def validate_params(params: Dict) -> bool:
-    """Validate request parameters for potential security issues.
+    @staticmethod
+    def validate(region: str, profile: Optional[str] = None) -> bool:
+        """Validate AWS credentials.
 
-    Args:
-        params: The parameters to validate
+        Args:
+            region: AWS region to use
+            profile: AWS profile to use (optional)
 
-    Returns:
-        bool: True if the parameters are safe, False otherwise
-    """
-    if not params:
-        return True
+        Returns:
+            bool: True if credentials are valid, False otherwise
+        """
+        logger.info('Validating AWS credentials...')
 
-    # List of dangerous patterns to check for
-    dangerous_patterns = [
-        # Command injection attempts
-        ';',
-        '&&',
-        '||',
-        '`',
-        '$(',
-        '${',
-        # File access attempts
-        'file://',
-        '/etc/',
-        '/var/log',
-        # Network access attempts
-        'http://',
-        'https://',
-    ]
-
-    # Check each parameter value
-    for key, value in params.items():
-        if not isinstance(value, str):
-            continue
-
-        for pattern in dangerous_patterns:
-            if pattern in value:
-                logger.warning(f'Potentially dangerous parameter detected: {key}={value}')
+        try:
+            # Create session with profile if specified
+            if profile:
+                logger.info(f'Using AWS Profile: {profile}')
+                session = boto3.Session(profile_name=profile, region_name=region)
+            else:
+                logger.info('Using default AWS credentials')
+                session = boto3.Session(region_name=region)
+
+            # Test AWS credentials
+            credentials = session.get_credentials()
+            if not credentials:
+                logger.error('ERROR: AWS credentials not found')
+                logger.error('Please configure AWS credentials using:')
+                logger.error('  - AWS CLI: aws configure')
+                logger.error(
+                    '  - Environment variables: AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY'
+                )
+                logger.error('  - Or specify a profile with --profile')
                 return False
 
-    return True
-
-
-async def make_prometheus_request(
-    endpoint: str, params: Optional[Dict] = None, max_retries: int = 3
-) -> Any:
-    """Make a request to the Prometheus HTTP API with AWS SigV4 authentication.
-
-    Args:
-        endpoint: The Prometheus API endpoint to call
-        params: Query parameters to include in the request
-        max_retries: Maximum number of retry attempts
-
-    Returns:
-        The data portion of the Prometheus API response
-
-    Raises:
-        ValueError: If Prometheus URL or AWS credentials are not configured
-        RuntimeError: If the Prometheus API returns an error status
-        requests.RequestException: If there's a network or HTTP error
-        json.JSONDecodeError: If the response is not valid JSON
-    """
-    if not config or not config.prometheus_url:
-        raise ValueError('Prometheus URL not configured')
-
-    # Validate endpoint
-    if not isinstance(endpoint, str):
-        raise ValueError('Endpoint must be a string')
+            # Test if credentials have necessary permissions
+            sts = session.client('sts', config=Config(user_agent_extra='prometheus-mcp-server'))
+            identity = sts.get_caller_identity()
+            logger.info(f'AWS Identity: {identity["Arn"]}')
+            logger.info(f'AWS Region: {region}')
+            logger.info('AWS credentials validated successfully')
+            return True
+
+        except (NoCredentialsError, ClientError) as e:
+            logger.error(f'ERROR: AWS credentials validation failed: {e}')
+            return False
 
-    if ';' in endpoint or '&&' in endpoint or '||' in endpoint:
-        raise ValueError('Invalid endpoint: potentially dangerous characters detected')
 
-    # Validate parameters
-    if params and not validate_params(params):
-        raise ValueError('Invalid parameters: potentially dangerous values detected')
+# Define dangerous patterns as a constant
+DANGEROUS_PATTERNS = [
+    # Command injection attempts
+    ';',
+    '&&',
+    '||',
+    '`',
+    '$(',
+    '${',
+    # File access attempts
+    'file://',
+    '/etc/',
+    '/var/log',
+    # Network access attempts
+    'http://',
+    'https://',
+]
+
+
+class SecurityValidator:
+    """Security validation utilities."""
+
+    @staticmethod
+    def validate_string(value: str, context: str = 'value') -> bool:
+        """Validate a string for potential security issues.
+
+        Args:
+            value: The string to validate
+            context: Context description for logging (e.g., 'parameter', 'query')
+
+        Returns:
+            bool: True if the string is safe, False otherwise
+        """
+        # Check for dangerous patterns
+        for pattern in DANGEROUS_PATTERNS:
+            if pattern in value:
+                logger.warning(f'Potentially dangerous {context} detected: {pattern}')
+                return False
 
-    # Ensure the URL ends with /api/v1
-    base_url = config.prometheus_url
-    if not base_url.endswith(API_VERSION_PATH):
-        base_url = f'{base_url.rstrip("/")}{API_VERSION_PATH}'
+        return True
 
-    url = f'{base_url}/{endpoint.lstrip("/")}'
+    @staticmethod
+    def validate_params(params: Dict) -> bool:
+        """Validate request parameters for potential security issues.
 
-    # Create AWS request
-    aws_request = AWSRequest(method='GET', url=url, params=params or {})
+        Args:
+            params: The parameters to validate
 
-    # Sign request with SigV4
-    session = boto3.Session(profile_name=config.aws_profile, region_name=config.aws_region)
-    credentials = session.get_credentials()
-    if not credentials:
-        raise ValueError('AWS credentials not found')
+        Returns:
+            bool: True if the parameters are safe, False otherwise
+        """
+        if not params:
+            return True
 
-    SigV4Auth(credentials, config.service_name, config.aws_region).add_auth(aws_request)
+        # Check each parameter value
+        for key, value in params.items():
+            if not isinstance(value, str):
+                continue
 
-    # Convert to requests format
-    prepared_request = requests.Request(
-        method=aws_request.method,
-        url=aws_request.url,
-        headers=dict(aws_request.headers),
-        params=params or {},
-    ).prepare()
+            if not SecurityValidator.validate_string(value, f'parameter {key}'):
+                return False
 
-    # Send request with retry logic
-    retry_count = 0
-    last_exception = None
-    retry_delay_seconds = 1  # Default retry delay if config.retry_delay is None
+        return True
 
-    while retry_count < max_retries:
-        try:
-            with requests.Session() as session:
-                logger.debug(f'Making request to {url} (attempt {retry_count + 1}/{max_retries})')
-                response = session.send(prepared_request)
-                response.raise_for_status()
-                data = response.json()
-
-                if data['status'] != 'success':
-                    error_msg = data.get('error', 'Unknown error')
-                    logger.error(f'Prometheus API request failed: {error_msg}')
-                    raise RuntimeError(f'Prometheus API request failed: {error_msg}')
-
-                return data['data']
-        except (requests.RequestException, json.JSONDecodeError) as e:
-            last_exception = e
-            retry_count += 1
-            if retry_count < max_retries:
-                if config and hasattr(config, 'retry_delay') and config.retry_delay is not None:
-                    retry_delay_seconds = config.retry_delay * (
+    @staticmethod
+    def validate_query(query: str) -> bool:
+        """Validate a PromQL query for potential security issues.
+
+        Args:
+            query: The PromQL query to validate
+
+        Returns:
+            bool: True if the query is safe, False otherwise
+        """
+        return SecurityValidator.validate_string(query, 'query pattern')
+
+
+class PrometheusClient:
+    """Client for interacting with Prometheus API."""
+
+    @staticmethod
+    async def make_request(
+        prometheus_url: str,
+        endpoint: str,
+        params: Optional[Dict] = None,
+        region: str = DEFAULT_AWS_REGION,
+        profile: Optional[str] = None,
+        max_retries: int = DEFAULT_MAX_RETRIES,
+        retry_delay: int = DEFAULT_RETRY_DELAY,
+        service_name: str = DEFAULT_SERVICE_NAME,
+    ) -> Any:
+        """Make a request to the Prometheus HTTP API with AWS SigV4 authentication.
+
+        Args:
+            prometheus_url: The base URL for the Prometheus API
+            endpoint: The Prometheus API endpoint to call
+            params: Query parameters to include in the request
+            region: AWS region to use
+            profile: AWS profile to use
+            max_retries: Maximum number of retry attempts
+            retry_delay: Delay between retry attempts in seconds
+            service_name: AWS service name for SigV4 authentication
+
+        Returns:
+            The data portion of the Prometheus API response
+
+        Raises:
+            ValueError: If Prometheus URL or AWS credentials are not configured
+            RuntimeError: If the Prometheus API returns an error status
+            requests.RequestException: If there's a network or HTTP error
+            json.JSONDecodeError: If the response is not valid JSON
+        """
+        if not prometheus_url:
+            raise ValueError('Prometheus URL not configured')
+
+        # Validate endpoint
+        if not isinstance(endpoint, str):
+            raise ValueError('Endpoint must be a string')
+
+        if ';' in endpoint or '&&' in endpoint or '||' in endpoint:
+            raise ValueError('Invalid endpoint: potentially dangerous characters detected')
+
+        # Validate parameters
+        if params and not SecurityValidator.validate_params(params):
+            raise ValueError('Invalid parameters: potentially dangerous values detected')
+
+        # Ensure the URL ends with /api/v1
+        base_url = prometheus_url
+        if not base_url.endswith(API_VERSION_PATH):
+            base_url = f'{base_url.rstrip("/")}{API_VERSION_PATH}'
+
+        url = f'{base_url}/{endpoint.lstrip("/")}'
+
+        # Send request with retry logic
+        retry_count = 0
+        last_exception = None
+        retry_delay_seconds = retry_delay
+
+        while retry_count < max_retries:
+            try:
+                # Create a fresh session and client for each attempt
+                session = boto3.Session(profile_name=profile, region_name=region)
+                credentials = session.get_credentials()
+                if not credentials:
+                    raise ValueError('AWS credentials not found')
+
+                # Create and sign the request
+                aws_request = AWSRequest(method='GET', url=url, params=params or {})
+                SigV4Auth(credentials, service_name, region).add_auth(aws_request)
+
+                # Convert to requests format
+                prepared_request = requests.Request(
+                    method=aws_request.method,
+                    url=aws_request.url,
+                    headers=dict(aws_request.headers),
+                    params=params or {},
+                ).prepare()
+
+                # Send the request
+                with requests.Session() as req_session:
+                    logger.debug(
+                        f'Making request to {url} (attempt {retry_count + 1}/{max_retries})'
+                    )
+                    response = req_session.send(prepared_request)
+                    response.raise_for_status()
+                    data = response.json()
+
+                    if data['status'] != 'success':
+                        error_msg = data.get('error', 'Unknown error')
+                        logger.error(f'Prometheus API request failed: {error_msg}')
+                        raise RuntimeError(f'Prometheus API request failed: {error_msg}')
+
+                    return data['data']
+            except (requests.RequestException, json.JSONDecodeError) as e:
+                last_exception = e
+                retry_count += 1
+                if retry_count < max_retries:
+                    retry_delay_seconds = retry_delay * (
                         2 ** (retry_count - 1)
                     )  # Exponential backoff
+                    logger.warning(f'Request failed: {e}. Retrying in {retry_delay_seconds}s...')
+                    time.sleep(retry_delay_seconds)
                 else:
-                    retry_delay_seconds = 1 * (
-                        2 ** (retry_count - 1)
-                    )  # Default exponential backoff
-                logger.warning(f'Request failed: {e}. Retrying in {retry_delay_seconds}s...')
-                time.sleep(retry_delay_seconds)
-            else:
-                logger.error(f'Request failed after {max_retries} attempts: {e}')
-                raise
+                    logger.error(f'Request failed after {max_retries} attempts: {e}')
+                    raise
 
-    if last_exception:
-        raise last_exception
-    return None
+        if last_exception:
+            raise last_exception
+        return None
 
 
-async def test_prometheus_connection():
-    """Test the connection to Prometheus.
+class PrometheusConnection:
+    """Handles Prometheus connection testing."""
 
-    Returns:
-        bool: True if connection is successful, False otherwise
-    """
-    logger.info('Testing Prometheus connection...')
-    try:
-        await make_prometheus_request('label/__name__/values', params={})
-        logger.info('Successfully connected to Prometheus!')
-        return True
-    except ClientError as e:
-        error_code = e.response.get('Error', {}).get('Code', 'Unknown')
-        if error_code == 'AccessDeniedException':
-            logger.error('ERROR: Access denied when connecting to Prometheus')
-            logger.error('Please check that your AWS credentials have the following permissions:')
-            logger.error('  - aps:QueryMetrics')
-            logger.error('  - aps:GetLabels')
-            logger.error('  - aps:GetMetricMetadata')
-        elif error_code == 'ResourceNotFoundException':
-            logger.error('ERROR: Prometheus workspace not found')
-            prometheus_url = 'Not configured'
-            if config and hasattr(config, 'prometheus_url') and config.prometheus_url:
-                prometheus_url = config.prometheus_url
-            logger.error(
-                f'Please verify the workspace ID in your Prometheus URL: {prometheus_url}'
+    @staticmethod
+    async def test_connection(
+        prometheus_url: str, region: str = DEFAULT_AWS_REGION, profile: Optional[str] = None
+    ) -> bool:
+        """Test the connection to Prometheus.
+
+        Args:
+            prometheus_url: The Prometheus URL to test
+            region: AWS region to use
+            profile: AWS profile to use
+
+        Returns:
+            bool: True if connection is successful, False otherwise
+        """
+        logger.info('Testing Prometheus connection...')
+        try:
+            # Use the PrometheusClient.make_request method
+            await PrometheusClient.make_request(
+                prometheus_url=prometheus_url,
+                endpoint='label/__name__/values',
+                params={},
+                region=region,
+                profile=profile,
+                max_retries=DEFAULT_MAX_RETRIES,
+                retry_delay=DEFAULT_RETRY_DELAY,
+                service_name=DEFAULT_SERVICE_NAME,
             )
-        else:
-            logger.error(f'ERROR: AWS API error when connecting to Prometheus: {error_code}')
-            logger.error(f'Details: {str(e)}')
-        return False
-    except requests.RequestException as e:
-        logger.error(f'ERROR: Network error when connecting to Prometheus: {str(e)}')
-        logger.error('Please check your network connection and Prometheus URL')
-        return False
-    except Exception as e:
-        logger.error(f'ERROR: Error connecting to Prometheus: {str(e)}')
-        logger.error('Common issues:')
-        logger.error('1. Incorrect Prometheus URL')
-        logger.error('2. Missing or incorrect AWS region')
-        logger.error('3. Invalid AWS credentials or insufficient permissions')
-        return False
+            logger.info('Successfully connected to Prometheus!')
+            return True
+        except ClientError as e:
+            error_code = e.response.get('Error', {}).get('Code', 'Unknown')
+            if error_code == 'AccessDeniedException':
+                logger.error('ERROR: Access denied when connecting to Prometheus')
+                logger.error(
+                    'Please check that your AWS credentials have the following permissions:'
+                )
+                logger.error('  - aps:QueryMetrics')
+                logger.error('  - aps:GetLabels')
+                logger.error('  - aps:GetMetricMetadata')
+            elif error_code == 'ResourceNotFoundException':
+                logger.error('ERROR: Prometheus workspace not found')
+                logger.error(
+                    f'Please verify the workspace ID in your Prometheus URL: {prometheus_url}'
+                )
+            else:
+                logger.error(f'ERROR: AWS API error when connecting to Prometheus: {error_code}')
+                logger.error(f'Details: {str(e)}')
+            return False
+        except requests.RequestException as e:
+            logger.error(f'ERROR: Network error when connecting to Prometheus: {str(e)}')
+            logger.error('Please check your network connection and Prometheus URL')
+            return False
+        except Exception as e:
+            logger.error(f'ERROR: Error connecting to Prometheus: {str(e)}')
+            logger.error('Common issues:')
+            logger.error('1. Incorrect Prometheus URL')
+            logger.error('2. Missing or incorrect AWS region')
+            logger.error('3. Invalid AWS credentials or insufficient permissions')
+            return False
 
 
 # Initialize MCP
@@ -411,73 +402,137 @@ mcp = FastMCP(
     ],
 )
 
-# Global config object
-config = None  # Will be initialized in main()
+# No global configuration - using environment variables instead
 
 
-@mcp.tool(name='ExecuteQuery')
-def validate_query(query: str) -> bool:
-    """Validate a PromQL query for potential security issues.
+def get_prometheus_client(region_name: Optional[str] = None, profile_name: Optional[str] = None):
+    """Create a boto3 AMP client using credentials from environment variables.
 
     Args:
-        query: The PromQL query to validate
+        region_name: AWS region to use (defaults to environment variable or us-east-1)
+        profile_name: AWS profile to use (defaults to None)
 
     Returns:
-        bool: True if the query is safe, False otherwise
+        boto3 AMP client with fresh credentials
+    """
+    # Use provided region, or get from env, or fall back to default
+    region = region_name or os.getenv('AWS_REGION') or DEFAULT_AWS_REGION
+
+    # Configure custom user agent
+    config = Config(user_agent_extra='prometheus-mcp-server')
+
+    # Create a new session to force credentials to reload
+    session = boto3.Session(profile_name=profile_name, region_name=region)
+
+    # Return AMP client
+    return session.client('amp', config=config)
+
 
-    This function checks for potentially dangerous patterns in PromQL queries.
+async def get_workspace_details(
+    workspace_id: str, region: str = DEFAULT_AWS_REGION, profile: Optional[str] = None
+) -> Dict[str, Any]:
+    """Get details for a specific Prometheus workspace using DescribeWorkspace API.
+
+    Args:
+        workspace_id: The Prometheus workspace ID
+        region: AWS region where the workspace is located
+        profile: AWS profile to use (defaults to None)
+
+    Returns:
+        Dictionary containing workspace details including URL from API
     """
-    # List of dangerous patterns to check for
-    dangerous_patterns = [
-        # Command injection attempts
-        ';',
-        '&&',
-        '||',
-        '`',
-        '$(',
-        '${',
-        # File access attempts
-        'file://',
-        '/etc/',
-        '/var/log',
-        # Network access attempts
-        'http://',
-        'https://',
-    ]
-
-    # Check for dangerous patterns
-    for pattern in dangerous_patterns:
-        if pattern in query:
-            logger.warning(f'Potentially dangerous query pattern detected: {pattern}')
-            return False
+    # Get a fresh client for this request
+    aps_client = get_prometheus_client(region_name=region, profile_name=profile)
 
-    return True
+    try:
+        # Get workspace details directly from DescribeWorkspace API
+        response = aps_client.describe_workspace(workspaceId=workspace_id)
+        workspace = response.get('workspace', {})
+
+        # Get the URL from the API response
+        prometheus_url = workspace.get('prometheusEndpoint')
+        if not prometheus_url:
+            raise ValueError(
+                f'No prometheusEndpoint found in workspace response for {workspace_id}'
+            )
+
+        logger.info(f'Retrieved workspace URL from DescribeWorkspace API: {prometheus_url}')
+
+        return {
+            'workspace_id': workspace_id,
+            'alias': workspace.get('alias', 'No alias'),
+            'status': workspace.get('status', {}).get('statusCode', 'UNKNOWN'),
+            'prometheus_url': prometheus_url,
+            'region': region,
+        }
+    except Exception as e:
+        logger.error(f'Error in DescribeWorkspace API: {str(e)}')
+        raise
 
 
+# validate_query function removed - now part of SecurityValidator class
+
+
+@mcp.tool(name='ExecuteQuery')
 async def execute_query(
     ctx: Context,
+    workspace_id: Optional[str] = Field(
+        None,
+        description='The Prometheus workspace ID to use (e.g., ws-12345678-abcd-1234-efgh-123456789012). Optional if a URL is configured via command line arguments.',
+    ),
     query: str = Field(..., description='The PromQL query to execute'),
     time: Optional[str] = Field(
         None, description='Optional timestamp for query evaluation (RFC3339 or Unix timestamp)'
     ),
+    region: Optional[str] = Field(None, description='AWS region (defaults to current region)'),
+    profile: Optional[str] = Field(None, description='AWS profile to use (defaults to None)'),
 ) -> Dict[str, Any]:
-    """Execute an instant query and return the result.
+    """Execute a PromQL query against Amazon Managed Prometheus.
 
     ## Usage
     - Use this tool to execute a PromQL query at a specific instant in time
     - The query will return the current value of the specified metrics
     - For time series data over a range, use execute_range_query instead
+    - If workspace_id is not known, use GetAvailableWorkspaces tool first to find available workspaces and ASK THE USER to choose one
+    - Uses DescribeWorkspace API to get the exact workspace URL
+    - No manual URL construction is performed
 
-    ## Example queries
+    ## Example
+    Input:
+      workspace_id: "ws-12345678-abcd-1234-efgh-123456789012"
+      query: "up"
+      region: "us-east-1"
+
+    Output:
+      {
+        "resultType": "vector",
+        "result": [
+          {
+            "metric": {"__name__": "up", "instance": "localhost:9090", "job": "prometheus"},
+            "value": [1680307200, "1"]
+          },
+          {
+            "metric": {"__name__": "up", "instance": "localhost:9100", "job": "node"},
+            "value": [1680307200, "1"]
+          }
+        ]
+      }
+
+    Example queries:
     - `up` - Shows which targets are up
     - `rate(node_cpu_seconds_total{mode="system"}[1m])` - CPU usage rate
     - `sum by(instance) (rate(node_network_receive_bytes_total[5m]))` - Network receive rate by instance
     """
     try:
+        # Configure workspace using the provided workspace_id
+        workspace_config = await configure_workspace_for_request(
+            ctx, workspace_id, region, profile
+        )
+
         logger.info(f'Executing instant query: {query}')
 
         # Validate query for security
-        if not validate_query(query):
+        if not SecurityValidator.validate_query(query):
             error_msg = 'Query validation failed: potentially dangerous query pattern detected'
             logger.error(error_msg)
             await ctx.error(error_msg)
@@ -487,11 +542,16 @@ async def execute_query(
         if time:
             params['time'] = time
 
-        max_retries = 3  # Default value
-        if config and hasattr(config, 'max_retries') and config.max_retries is not None:
-            max_retries = config.max_retries
-
-        return await make_prometheus_request('query', params, max_retries)
+        return await PrometheusClient.make_request(
+            prometheus_url=workspace_config['prometheus_url'],
+            endpoint='query',
+            params=params,
+            region=workspace_config['region'],
+            profile=workspace_config['profile'],
+            max_retries=DEFAULT_MAX_RETRIES,
+            retry_delay=DEFAULT_RETRY_DELAY,
+            service_name=DEFAULT_SERVICE_NAME,
+        )
     except Exception as e:
         error_msg = f'Error executing query: {str(e)}'
         logger.error(error_msg)
@@ -502,33 +562,58 @@ async def execute_query(
 @mcp.tool(name='ExecuteRangeQuery')
 async def execute_range_query(
     ctx: Context,
+    workspace_id: Optional[str] = Field(
+        None,
+        description='The Prometheus workspace ID to use (e.g., ws-12345678-abcd-1234-efgh-123456789012). Optional if a URL is configured via command line arguments.',
+    ),
     query: str = Field(..., description='The PromQL query to execute'),
     start: str = Field(..., description='Start timestamp (RFC3339 or Unix timestamp)'),
     end: str = Field(..., description='End timestamp (RFC3339 or Unix timestamp)'),
     step: str = Field(
         ..., description="Query resolution step width (duration format, e.g. '15s', '1m', '1h')"
     ),
+    region: Optional[str] = Field(None, description='AWS region (defaults to current region)'),
+    profile: Optional[str] = Field(None, description='AWS profile to use (defaults to None)'),
 ) -> Dict[str, Any]:
-    """Execute a range query and return the result.
+    r"""Execute a range query and return the result.
 
     ## Usage
     - Use this tool to execute a PromQL query over a time range
     - The query will return a series of values for the specified time range
     - Useful for generating time series data for graphs or trend analysis
+    - If workspace_id is not known, use GetAvailableWorkspaces tool first to find available workspaces and ASK THE USER to choose one
+    - Uses DescribeWorkspace API to get the exact workspace URL
+    - No manual URL construction is performed
 
     ## Example
-    - Query: `rate(node_cpu_seconds_total{mode="system"}[5m])`
-    - Start: `2023-04-01T00:00:00Z`
-    - End: `2023-04-01T01:00:00Z`
-    - Step: `5m`
-
-    This will return CPU usage rate sampled every 5 minutes over a 1-hour period.
+    Input:
+      workspace_id: "ws-12345678-abcd-1234-efgh-123456789012"
+      query: "rate(node_cpu_seconds_total{mode=\"system\"}[5m])"
+      start: "2023-04-01T00:00:00Z"
+      end: "2023-04-01T01:00:00Z"
+      step: "5m"
+
+    Output:
+      {
+        "resultType": "matrix",
+        "result": [
+          {
+            "metric": {"__name__": "rate", "mode": "system", "instance": "localhost:9100"},
+            "values": [[1680307200, "0.01"], [1680307500, "0.012"], ...]
+          }
+        ]
+      }
     """
     try:
+        # Configure workspace using the provided workspace_id
+        workspace_config = await configure_workspace_for_request(
+            ctx, workspace_id, region, profile
+        )
+
         logger.info(f'Executing range query: {query} from {start} to {end} with step {step}')
 
         # Validate query for security
-        if not validate_query(query):
+        if not SecurityValidator.validate_query(query):
             error_msg = 'Query validation failed: potentially dangerous query pattern detected'
             logger.error(error_msg)
             await ctx.error(error_msg)
@@ -536,11 +621,16 @@ async def execute_range_query(
 
         params = {'query': query, 'start': start, 'end': end, 'step': step}
 
-        max_retries = 3  # Default value
-        if config and hasattr(config, 'max_retries') and config.max_retries is not None:
-            max_retries = config.max_retries
-
-        return await make_prometheus_request('query_range', params, max_retries)
+        return await PrometheusClient.make_request(
+            prometheus_url=workspace_config['prometheus_url'],
+            endpoint='query_range',
+            params=params,
+            region=workspace_config['region'],
+            profile=workspace_config['profile'],
+            max_retries=DEFAULT_MAX_RETRIES,
+            retry_delay=DEFAULT_RETRY_DELAY,
+            service_name=DEFAULT_SERVICE_NAME,
+        )
     except Exception as e:
         error_msg = f'Error executing range query: {str(e)}'
         logger.error(error_msg)
@@ -549,28 +639,55 @@ async def execute_range_query(
 
 
 @mcp.tool(name='ListMetrics')
-async def list_metrics(ctx: Context) -> MetricsList:
+async def list_metrics(
+    ctx: Context,
+    workspace_id: Optional[str] = Field(
+        None,
+        description='The Prometheus workspace ID to use (e.g., ws-12345678-abcd-1234-efgh-123456789012). Optional if a URL is configured via command line arguments.',
+    ),
+    region: Optional[str] = Field(None, description='AWS region (defaults to current region)'),
+    profile: Optional[str] = Field(None, description='AWS profile to use (defaults to None)'),
+) -> MetricsList:
     """Get a list of all metric names.
 
     ## Usage
     - Use this tool to discover available metrics in the Prometheus server
     - Returns a sorted list of all metric names
     - Useful for exploration before crafting specific queries
+    - If workspace_id is not known, use GetAvailableWorkspaces tool first to find available workspaces and ASK THE USER to choose one
 
     ## Example
-    ```
-    metrics_response = await list_metrics()
-    print('Available metrics:', metrics_response.metrics[:10])  # Show first 10 metrics
-    ```
+    Input:
+      workspace_id: "ws-12345678-abcd-1234-efgh-123456789012"
+      region: "us-east-1"
+
+    Output:
+      {
+        "metrics": [
+          "go_gc_duration_seconds",
+          "go_goroutines",
+          "http_requests_total",
+          ...
+        ]
+      }
     """
     try:
+        # Configure workspace using the provided workspace_id
+        workspace_config = await configure_workspace_for_request(
+            ctx, workspace_id, region, profile
+        )
+
         logger.info('Listing all available metrics')
-        max_retries = 3  # Default value
-        if config and hasattr(config, 'max_retries') and config.max_retries is not None:
-            max_retries = config.max_retries
 
-        data = await make_prometheus_request(
-            'label/__name__/values', params={}, max_retries=max_retries
+        data = await PrometheusClient.make_request(
+            prometheus_url=workspace_config['prometheus_url'],
+            endpoint='label/__name__/values',
+            params={},
+            region=workspace_config['region'],
+            profile=workspace_config['profile'],
+            max_retries=DEFAULT_MAX_RETRIES,
+            retry_delay=DEFAULT_RETRY_DELAY,
+            service_name=DEFAULT_SERVICE_NAME,
         )
         return MetricsList(metrics=sorted(data))
     except Exception as e:
@@ -581,35 +698,51 @@ async def list_metrics(ctx: Context) -> MetricsList:
 
 
 @mcp.tool(name='GetServerInfo')
-async def get_server_info(ctx: Context) -> ServerInfo:
+async def get_server_info(
+    ctx: Context,
+    workspace_id: Optional[str] = Field(
+        None,
+        description='The Prometheus workspace ID to use (e.g., ws-12345678-abcd-1234-efgh-123456789012). Optional if a URL is configured via command line arguments.',
+    ),
+    region: Optional[str] = Field(None, description='AWS region (defaults to current region)'),
+    profile: Optional[str] = Field(None, description='AWS profile to use (defaults to None)'),
+) -> ServerInfo:
     """Get information about the Prometheus server configuration.
 
     ## Usage
     - Use this tool to retrieve the current server configuration
     - Returns details about the Prometheus URL, AWS region, profile, and service name
     - Useful for debugging connection issues
+    - If workspace_id is not known, use GetAvailableWorkspaces tool first to find available workspaces and ASK THE USER to choose one
+    - Uses DescribeWorkspace API to get the exact workspace URL
+    - No manual URL construction is performed
 
     ## Example
-    ```
-    info = await get_server_info()
-    print(f'Connected to Prometheus at {info.prometheus_url} in region {info.aws_region}')
-    ```
+    Input:
+      workspace_id: "ws-12345678-abcd-1234-efgh-123456789012"
+      region: "us-east-1"
+
+    Output:
+      {
+        "prometheus_url": "https://aps-workspaces.us-east-1.amazonaws.com/workspaces/ws-12345678-abcd-1234-efgh-123456789012",
+        "aws_region": "us-east-1",
+        "aws_profile": "default",
+        "service_name": "aps"
+      }
     """
     try:
+        # Configure workspace using the provided workspace_id
+        workspace_config = await configure_workspace_for_request(
+            ctx, workspace_id, region, profile
+        )
+
         logger.info('Retrieving server configuration information')
-        if not config:
-            return ServerInfo(
-                prometheus_url='Not configured',
-                aws_region='Not configured',
-                aws_profile='Not configured',
-                service_name='Not configured',
-            )
 
         return ServerInfo(
-            prometheus_url=config.prometheus_url or 'Not configured',
-            aws_region=config.aws_region or 'Not configured',
-            aws_profile=config.aws_profile or 'default',
-            service_name=config.service_name or DEFAULT_SERVICE_NAME,
+            prometheus_url=workspace_config['prometheus_url'],
+            aws_region=workspace_config['region'],
+            aws_profile=workspace_config['profile'] or 'default',
+            service_name=DEFAULT_SERVICE_NAME,
         )
     except Exception as e:
         error_msg = f'Error retrieving server info: {str(e)}'
@@ -618,14 +751,302 @@ async def get_server_info(ctx: Context) -> ServerInfo:
         raise
 
 
+@mcp.tool(name='GetAvailableWorkspaces')
+async def get_available_workspaces(
+    ctx: Context,
+    region: Optional[str] = Field(None, description='AWS region (defaults to current region)'),
+    profile: Optional[str] = Field(None, description='AWS profile to use (defaults to None)'),
+) -> Dict[str, Any]:
+    """List all available Prometheus workspaces in the specified region.
+
+    ## Usage
+    - Use this tool to see all available Prometheus workspaces
+    - Shows workspace ID, alias, status, and URL for active workspaces
+    - IMPORTANT: When multiple workspaces are available, present them to the user and ask them to choose one
+    - DO NOT automatically select a workspace; always ask the user to choose when multiple options exist
+    - Uses DescribeWorkspace API to get the exact URL for each workspace
+    - No manual URL construction is performed
+
+    ## Example
+    Input:
+      region: "us-east-1"
+
+    Output:
+      {
+        "workspaces": [
+          {
+            "workspace_id": "ws-12345678-abcd-1234-efgh-123456789012",
+            "alias": "production",
+            "status": "ACTIVE",
+            "prometheus_url": "https://aps-workspaces.us-east-1.amazonaws.com/workspaces/ws-12345678-abcd-1234-efgh-123456789012",
+            "is_configured": true
+          },
+          {
+            "workspace_id": "ws-87654321-dcba-4321-hgfe-210987654321",
+            "alias": "development",
+            "status": "ACTIVE",
+            "prometheus_url": "https://aps-workspaces.us-east-1.amazonaws.com/workspaces/ws-87654321-dcba-4321-hgfe-210987654321",
+            "is_configured": false
+          }
+        ],
+        "count": 2,
+        "region": "us-east-1",
+        "requires_user_selection": true,
+        "configured_workspace_id": "ws-12345678-abcd-1234-efgh-123456789012"
+      }
+    """
+    try:
+        # Use provided region or default from environment
+        aws_region = region or os.getenv('AWS_REGION') or DEFAULT_AWS_REGION
+        aws_profile = profile or os.getenv(ENV_AWS_PROFILE)
+
+        # Check if we already have a URL configured and if it contains a workspace ID
+        prometheus_url = os.getenv('PROMETHEUS_URL')
+        configured_workspace_id = None
+        if prometheus_url:
+            configured_workspace_id = extract_workspace_id_from_url(prometheus_url)
+            if configured_workspace_id:
+                logger.info(f'Found configured workspace ID in URL: {configured_workspace_id}')
+
+        logger.info(f'Listing available Prometheus workspaces in region {aws_region}')
+
+        # Get a fresh client for this request
+        aps_client = get_prometheus_client(region_name=aws_region, profile_name=aws_profile)
+        response = aps_client.list_workspaces()
+
+        workspaces = []
+        configured_workspace_details = None
+
+        for ws in response.get('workspaces', []):
+            workspace_id = ws['workspaceId']
+
+            # Only get details for active workspaces
+            if ws['status']['statusCode'] == 'ACTIVE':
+                try:
+                    # Get full details including URL from DescribeWorkspace API
+                    details = await get_workspace_details(workspace_id, aws_region, aws_profile)
+
+                    # If this is the configured workspace, mark it
+                    if configured_workspace_id and workspace_id == configured_workspace_id:
+                        details['is_configured'] = True
+                        details['note'] = '(Detected from URL - will be used automatically)'
+                        configured_workspace_details = details
+                    else:
+                        details['is_configured'] = False
+
+                    workspaces.append(details)
+                except Exception as e:
+                    logger.warning(f'Could not get details for workspace {workspace_id}: {str(e)}')
+                    # Skip this workspace if we can't get its details
+                    continue
+            else:
+                # For non-active workspaces, just include basic info without URL
+                workspaces.append(
+                    {
+                        'workspace_id': workspace_id,
+                        'alias': ws.get('alias', 'No alias'),
+                        'status': ws['status']['statusCode'],
+                        'region': aws_region,
+                        'is_configured': configured_workspace_id
+                        and workspace_id == configured_workspace_id,
+                    }
+                )
+
+        # If we have a configured workspace but it wasn't found in the list,
+        # it might be in a different region. Add it to the list if we have details.
+        if configured_workspace_id and not configured_workspace_details and prometheus_url:
+            try:
+                # Create a basic entry for the configured workspace
+                workspaces.append(
+                    {
+                        'workspace_id': configured_workspace_id,
+                        'alias': 'Configured Workspace',
+                        'status': 'ACTIVE',  # Assume active since we have a URL
+                        'prometheus_url': prometheus_url,
+                        'region': aws_region,
+                        'is_configured': True,
+                        'note': '(Detected from URL - will be used automatically)',
+                    }
+                )
+                logger.info(f'Added configured workspace {configured_workspace_id} from URL')
+            except Exception as e:
+                logger.warning(f'Could not add configured workspace: {str(e)}')
+
+        # Sort workspaces to put configured workspace first
+        workspaces.sort(key=lambda ws: 0 if ws.get('is_configured') else 1)
+
+        logger.info(f'Found {len(workspaces)} workspaces in region {aws_region}')
+
+        # If we have a configured workspace ID from the URL, we don't need user selection
+        requires_selection = not configured_workspace_id and len(workspaces) > 1
+
+        message = ''
+        if configured_workspace_id:
+            message = f'A workspace ID ({configured_workspace_id}) was detected in the URL and will be used automatically. You can override it by explicitly providing a workspace_id parameter.'
+        elif len(workspaces) > 1:
+            message = 'Please choose a workspace ID to use with your queries.'
+        else:
+            message = 'Only one workspace is available. You can use it by specifying its workspace_id in your queries.'
+
+        return {
+            'workspaces': workspaces,
+            'count': len(workspaces),
+            'region': aws_region,
+            'requires_user_selection': requires_selection,
+            'configured_workspace_id': configured_workspace_id,
+            'message': message,
+        }
+    except Exception as e:
+        error_msg = f'Error listing workspaces: {str(e)}'
+        logger.error(error_msg)
+        await ctx.error(error_msg)
+        raise
+
+
+def extract_workspace_id_from_url(url: str) -> Optional[str]:
+    """Extract workspace ID from a Prometheus URL.
+
+    Args:
+        url: The Prometheus URL that may contain a workspace ID
+
+    Returns:
+        The extracted workspace ID or None if not found
+    """
+    if not url:
+        return None
+
+    # Look for the pattern /workspaces/ws-XXXX in the URL
+    import re
+
+    match = re.search(r'/workspaces/(ws-[\w-]+)', url)
+    if match:
+        workspace_id = match.group(1)
+        logger.info(f'Extracted workspace ID from URL: {workspace_id}')
+        return workspace_id
+    return None
+
+
+async def configure_workspace_for_request(
+    ctx: Context,
+    workspace_id: Optional[str] = None,
+    region: Optional[str] = None,
+    profile: Optional[str] = None,
+) -> Dict[str, Any]:
+    """Configure the workspace for the current request.
+
+    If a URL is provided via environment variable, it will be used directly.
+    If a workspace ID is provided, it will be used to fetch the URL from AWS API.
+    If no workspace ID is provided but the URL contains one, it will be extracted and used.
+
+    Args:
+        ctx: The MCP context
+        workspace_id: The Prometheus workspace ID to use (optional if URL contains workspace ID)
+        region: Optional AWS region (defaults to current region)
+        profile: Optional AWS profile to use
+
+    Returns:
+        Dictionary with workspace configuration including the URL
+    """
+    try:
+        # Use provided region or default from environment
+        aws_region = region or os.getenv('AWS_REGION') or DEFAULT_AWS_REGION
+        aws_profile = profile or os.getenv(ENV_AWS_PROFILE)
+
+        # Check if we have a URL from environment
+        prometheus_url = os.getenv('PROMETHEUS_URL')
+
+        # If no workspace_id is provided, extract it from the URL if possible
+        if not workspace_id and prometheus_url:
+            extracted_workspace_id = extract_workspace_id_from_url(prometheus_url)
+            if extracted_workspace_id:
+                workspace_id = extracted_workspace_id
+                logger.info(f'Using workspace ID extracted from URL: {workspace_id}')
+
+        # If we have a URL but no workspace_id could be extracted, use the URL directly
+        if prometheus_url:
+            logger.info(f'Using Prometheus URL from environment: {prometheus_url}')
+
+            # Test connection with the URL
+            if not await PrometheusConnection.test_connection(
+                prometheus_url, aws_region, aws_profile
+            ):
+                error_msg = f'Failed to connect to Prometheus with configured URL {prometheus_url}'
+                logger.error(error_msg)
+                await ctx.error(error_msg)
+                raise RuntimeError(error_msg)
+
+            return {
+                'prometheus_url': prometheus_url,
+                'region': aws_region,
+                'profile': aws_profile,
+                'workspace_id': workspace_id,
+            }
+
+        # If no URL is configured, require workspace_id
+        if not workspace_id:
+            error_msg = 'Workspace ID is required when no Prometheus URL is configured. Please use GetAvailableWorkspaces to list available workspaces and choose one.'
+            logger.error(error_msg)
+            await ctx.error(error_msg)
+            raise ValueError(error_msg)
+
+        logger.info(f'Configuring workspace ID for request: {workspace_id}')
+
+        # Validate workspace ID format
+        if not workspace_id.startswith('ws-'):
+            logger.warning(
+                f'Workspace ID "{workspace_id}" does not start with "ws-", which is unusual'
+            )
+
+        # Get workspace details from DescribeWorkspace API
+        workspace_details = await get_workspace_details(workspace_id, aws_region, aws_profile)
+        prometheus_url = workspace_details['prometheus_url']
+        logger.info(f'Using Prometheus URL from DescribeWorkspace API: {prometheus_url}')
+
+        # Test connection with the URL
+        if not await PrometheusConnection.test_connection(prometheus_url, aws_region, aws_profile):
+            error_msg = f'Failed to connect to Prometheus with workspace ID {workspace_id}'
+            logger.error(error_msg)
+            await ctx.error(error_msg)
+            raise RuntimeError(error_msg)
+
+        logger.info(f'Successfully configured workspace {workspace_id} for request')
+
+        # Return workspace configuration
+        return {
+            'prometheus_url': prometheus_url,
+            'region': aws_region,
+            'profile': aws_profile,
+            'workspace_id': workspace_id,
+        }
+    except Exception as e:
+        error_msg = f'Error configuring workspace: {str(e)}'
+        logger.error(error_msg)
+        await ctx.error(error_msg)
+        raise
+
+
 async def async_main():
     """Run the async initialization tasks."""
-    # Test connection
-    if not await test_prometheus_connection():
-        logger.error('Prometheus connection test failed')
-        sys.exit(1)
+    # Check if URL is configured in environment
+    prometheus_url = os.getenv('PROMETHEUS_URL')
+    if prometheus_url:
+        logger.info(f'Using Prometheus URL from environment: {prometheus_url}')
+
+        # Check if the URL contains a workspace ID
+        workspace_id = extract_workspace_id_from_url(prometheus_url)
+        if workspace_id:
+            logger.info(f'Detected workspace ID in URL: {workspace_id}')
+            logger.info(
+                'This workspace ID can be used with queries, but must be explicitly provided'
+            )
+        else:
+            logger.info('No workspace ID detected in URL')
 
-    logger.info('Prometheus connection successful')
+        logger.info('Workspace ID will be required for each tool invocation')
+    else:
+        logger.info(
+            'Initializing Prometheus MCP Server - workspace ID will be required for each tool invocation'
+        )
 
 
 def main():
@@ -633,25 +1054,36 @@ def main():
     logger.info('Starting Prometheus MCP Server...')
 
     # Parse arguments
-    args = parse_arguments()
-
-    # Load configuration
-    config_data = load_config(args)
-
-    # Create config object
-    global config
-    config = PrometheusConfig(
-        prometheus_url=config_data['prometheus_url'],
-        aws_region=config_data['aws_region'],
-        aws_profile=config_data['aws_profile'],
-        service_name=config_data['service_name'],
-        retry_delay=config_data['retry_delay'],
-        max_retries=config_data['max_retries'],
-    )
-
-    # Setup environment
-    if not setup_environment(config_data):
-        logger.error('Environment setup failed')
+    args = ConfigManager.parse_arguments()
+
+    # Setup basic configuration
+    config = ConfigManager.setup_basic_config(args)
+
+    # Set as environment variables for other functions to use
+    if config['url']:
+        os.environ['PROMETHEUS_URL'] = config['url']
+    if config['region']:
+        os.environ['AWS_REGION'] = config['region']
+    if config['profile']:
+        os.environ[ENV_AWS_PROFILE] = config['profile']
+
+    if config['url']:
+        logger.info(f'Using configured Prometheus URL: {config["url"]}')
+
+        # Check if the URL contains a workspace ID
+        workspace_id = extract_workspace_id_from_url(config['url'])
+        if workspace_id:
+            logger.info(f'Detected workspace ID in URL: {workspace_id}')
+            logger.info(
+                'This workspace will be used automatically when no workspace ID is provided'
+            )
+        else:
+            logger.info('No workspace ID detected in URL')
+            logger.info('Workspace ID will be required for each tool invocation')
+
+    # Validate AWS credentials
+    if not AWSCredentials.validate(config['region'], config['profile']):
+        logger.error('AWS credentials validation failed')
         sys.exit(1)
 
     # Run async initialization in an event loop
@@ -660,8 +1092,14 @@ def main():
     asyncio.run(async_main())
 
     logger.info('Starting server...')
+
     # Run with stdio transport
-    mcp.run(transport='stdio')
+    try:
+        logger.info('Starting with stdio transport...')
+        mcp.run(transport='stdio')
+    except Exception as e:
+        logger.error(f'Error starting server with stdio transport: {e}')
+        sys.exit(1)
 
 
 if __name__ == '__main__':