From 03d9fde50a8d534d974bfd1b3e5589fb56aea4ea Mon Sep 17 00:00:00 2001 From: falamarcao Date: Fri, 8 Aug 2025 23:32:25 -0300 Subject: [PATCH 1/3] =?UTF-8?q?=F0=9F=8E=89=20refactor(logger):=20centrali?= =?UTF-8?q?ze=20logging=20with=20Rich=20for=20better=20formatting=20and=20?= =?UTF-8?q?highlighting?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/app/logger.py | 19 +++++ src/app/routers/browsers/routes.py | 5 +- src/app/routers/selenium_proxy.py | 4 +- .../services/selenium_hub/common/logger.py | 55 +++++++++++++++ .../services/selenium_hub/common/pidfile.py | 47 +++++++------ .../selenium_hub/core/docker_backend.py | 70 +++++++++---------- .../services/selenium_hub/core/hub_backend.py | 12 ++-- .../selenium_hub/core/kubernetes/backend.py | 52 +++++++------- .../core/kubernetes/common/decorators.py | 16 ++--- .../core/kubernetes/k8s_config.py | 20 +++--- .../core/kubernetes/k8s_port_forwarder.py | 43 ++++++------ .../core/kubernetes/k8s_resource_manager.py | 28 ++++---- .../core/kubernetes/k8s_url_resolver.py | 20 +++--- 13 files changed, 232 insertions(+), 159 deletions(-) create mode 100644 src/app/logger.py create mode 100644 src/app/services/selenium_hub/common/logger.py diff --git a/src/app/logger.py b/src/app/logger.py new file mode 100644 index 0000000..f491bd8 --- /dev/null +++ b/src/app/logger.py @@ -0,0 +1,19 @@ +import logging +from os import getenv + +from rich.logging import RichHandler + +LOG_LEVEL = getenv("LOG_LEVEL", "INFO") + +# Create +logger = logging.getLogger(f"MCP Selenium Grid:{__name__}") +logger.setLevel(LOG_LEVEL) + +rich_handler = RichHandler( + level=LOG_LEVEL, + markup=True, + rich_tracebacks=True, + tracebacks_show_locals=True, +) + +logger.addHandler(rich_handler) diff --git a/src/app/routers/browsers/routes.py b/src/app/routers/browsers/routes.py index bd6d153..8fb4291 100644 --- a/src/app/routers/browsers/routes.py +++ b/src/app/routers/browsers/routes.py @@ -7,6 +7,7 @@ from app.core.settings import Settings from app.dependencies import get_settings, verify_token +from app.logger import logger from app.services.selenium_hub import SeleniumHub from app.services.selenium_hub.models.browser import BrowserConfig, BrowserInstance @@ -66,9 +67,7 @@ async def create_browsers( app_state.browsers_instances[browser.id] = browser except Exception as e: # Log the error and current browser configs for diagnostics - import logging # noqa: PLC0415 - - logging.error( + logger.error( f"Exception in create_browsers: {e}. BROWSER_CONFIGS: {settings.selenium_grid.BROWSER_CONFIGS}" ) raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e)) diff --git a/src/app/routers/selenium_proxy.py b/src/app/routers/selenium_proxy.py index 6482da2..1cded5c 100644 --- a/src/app/routers/selenium_proxy.py +++ b/src/app/routers/selenium_proxy.py @@ -1,7 +1,6 @@ """Proxy router to securely expose Selenium Hub via FastAPI, supporting both Docker and Kubernetes deployments. All routes require HTTP Basic Auth matching the Selenium Hub configuration.""" import base64 -import logging from typing import Annotated from urllib.parse import urljoin @@ -12,6 +11,7 @@ from app.core.settings import Settings from app.dependencies import get_settings, verify_basic_auth +from app.logger import logger from app.services.selenium_hub import SeleniumHub # Constants @@ -34,9 +34,7 @@ "cache-control", } - router = APIRouter(prefix=SELENIUM_HUB_PREFIX, tags=["Selenium Hub"]) -logger = logging.getLogger(__name__) # --- Utility Functions --- diff --git a/src/app/services/selenium_hub/common/logger.py b/src/app/services/selenium_hub/common/logger.py new file mode 100644 index 0000000..3fbafc0 --- /dev/null +++ b/src/app/services/selenium_hub/common/logger.py @@ -0,0 +1,55 @@ +import logging +from os import getenv + +from rich.highlighter import Highlighter, ReprHighlighter +from rich.logging import RichHandler +from rich.text import Text + +LOG_LEVEL: str = getenv("LOG_LEVEL", "INFO") + + +class CustomKeywordHighlighter(Highlighter): + keywords: dict[str, str] + + def __init__(self) -> None: + super().__init__() + self.keywords = { + "running": "bold green", + "failed": "red", + "failling": "red", + "fail": "red", + "success": "bold bright_green", + "SUCCEED!": "bold bright_green", + "selenium-hub": "bold", + "selenium-grid": "bold", + "DockerHubBackend": "bold", + "KubernetesHubBackend": "bold", + "Docker": "bold", + "Kubernetes": "bold", + "health": "bold", + "Health check": "bold blue", + "Port-forward": "bold", + "port-forward": "bold", + } + + self.repr_highlighter = ReprHighlighter() + + def highlight(self, text: Text) -> None: + for keyword, style in self.keywords.items(): + text.highlight_words([keyword], style=style) + + self.repr_highlighter.highlight(text) + + +logger: logging.Logger = logging.getLogger(f"Selenium Grid:{__name__}") +logger.setLevel(LOG_LEVEL) + +rich_handler: RichHandler = RichHandler( + level=LOG_LEVEL, + markup=True, + rich_tracebacks=True, + tracebacks_show_locals=True, +) +rich_handler.highlighter = CustomKeywordHighlighter() + +logger.addHandler(rich_handler) diff --git a/src/app/services/selenium_hub/common/pidfile.py b/src/app/services/selenium_hub/common/pidfile.py index 14b8d14..da75d41 100644 --- a/src/app/services/selenium_hub/common/pidfile.py +++ b/src/app/services/selenium_hub/common/pidfile.py @@ -1,21 +1,22 @@ -import logging from pathlib import Path from typing import Iterable from psutil import AccessDenied, NoSuchProcess, Process, TimeoutExpired +from .logger import logger + class PidFile: def __init__(self, path: Path) -> None: self.path = path try: self.path.parent.mkdir(parents=True, exist_ok=True) - logging.debug(f"Created directory {self.path.parent} successfully.") + logger.debug(f"Created directory {self.path.parent} successfully.") except PermissionError: - logging.error(f"Permission denied: Cannot create directory {self.path.parent}") + logger.error(f"Permission denied: Cannot create directory {self.path.parent}") raise except OSError as e: - logging.error(f"OS error while creating directory {self.path.parent}: {e}") + logger.error(f"OS error while creating directory {self.path.parent}: {e}") raise def exists_and_alive(self) -> bool: @@ -23,47 +24,47 @@ def exists_and_alive(self) -> bool: if pid is None: return False alive = self._pid_alive(pid) - logging.debug(f"PID file {self.path} exists and process {pid} alive: {alive}") + logger.debug(f"PID file {self.path} exists and process {pid} alive: {alive}") return alive def read(self) -> int | None: if not self.path.exists(): - logging.debug(f"PID file {self.path} does not exist.") + logger.debug(f"PID file {self.path} does not exist.") return None try: content = self.path.read_text().strip() if not content: - logging.debug(f"PID file {self.path} is empty.") + logger.debug(f"PID file {self.path} is empty.") return None return int(content) except ValueError as e: - logging.debug(f"Invalid PID in {self.path}: {e}") + logger.debug(f"Invalid PID in {self.path}: {e}") self.remove() return None except Exception as e: # Only log unexpected I/O issues at higher level - logging.warning(f"Unexpected error reading PID file {self.path}: {e}") + logger.warning(f"Unexpected error reading PID file {self.path}: {e}") self.remove() return None def write(self, pid: int) -> None: try: self.path.write_text(str(pid)) - logging.debug(f"Wrote PID {pid} to {self.path}") + logger.debug(f"Wrote PID {pid} to {self.path}") except Exception as e: # This one should be ERROR — we expect this to work - logging.error(f"Failed to write PID {pid} to {self.path}: {e}") + logger.error(f"Failed to write PID {pid} to {self.path}: {e}") def remove(self) -> None: try: if self.path.exists(): self.path.unlink() - logging.debug(f"Removed PID file {self.path}") + logger.debug(f"Removed PID file {self.path}") else: - logging.debug(f"PID file {self.path} already removed.") + logger.debug(f"PID file {self.path} already removed.") except Exception as e: # Rare, but possible (e.g. permission change) - logging.warning(f"Failed to remove PID file {self.path}: {e}") + logger.warning(f"Failed to remove PID file {self.path}: {e}") def _pid_alive(self, pid: int) -> bool: try: @@ -89,12 +90,12 @@ def is_process_running_with_cmdline(pid: int, expected_cmd_parts: Iterable[str]) cmdline = " ".join(proc.cmdline()) for part in expected_cmd_parts: if part not in cmdline: - logging.debug(f"Cmdline check failed: '{part}' not in '{cmdline}'") + logger.debug(f"Cmdline check failed: '{part}' not in '{cmdline}'") return False - logging.debug(f"Process PID {pid} cmdline contains all expected parts.") + logger.debug(f"Process PID {pid} cmdline contains all expected parts.") return True except (NoSuchProcess, AccessDenied) as e: - logging.warning(f"Cannot access process PID {pid}: {e}") + logger.warning(f"Cannot access process PID {pid}: {e}") return False @@ -102,17 +103,17 @@ def terminate_pid(pid: int) -> None: """Attempt to terminate (and possibly kill) a process by PID.""" try: proc = Process(pid) - logging.info(f"Attempting to terminate process PID {pid}...") + logger.info(f"Attempting to terminate process PID {pid}...") proc.terminate() proc.wait(timeout=5) - logging.info(f"Terminated process PID {pid}.") + logger.info(f"Terminated process PID {pid}.") except (NoSuchProcess, AccessDenied) as e: - logging.warning(f"Could not terminate process PID {pid}: {e}") + logger.warning(f"Could not terminate process PID {pid}: {e}") except TimeoutExpired: - logging.warning(f"Process PID {pid} did not terminate in time, killing...") + logger.warning(f"Process PID {pid} did not terminate in time, killing...") try: proc.kill() except Exception as e: - logging.error(f"Failed to kill process PID {pid}: {e}") + logger.error(f"Failed to kill process PID {pid}: {e}") except Exception as e: - logging.error(f"Unexpected error terminating PID {pid}: {e}") + logger.error(f"Unexpected error terminating PID {pid}: {e}") diff --git a/src/app/services/selenium_hub/core/docker_backend.py b/src/app/services/selenium_hub/core/docker_backend.py index 23e2523..3d5ff28 100644 --- a/src/app/services/selenium_hub/core/docker_backend.py +++ b/src/app/services/selenium_hub/core/docker_backend.py @@ -1,9 +1,9 @@ -import logging from typing import override import docker from docker.errors import APIError, NotFound +from ..common.logger import logger from ..models.browser import BrowserConfig, BrowserConfigs, BrowserType from ..models.general_settings import SeleniumHubGeneralSettings from .hub_backend import HubBackend @@ -23,30 +23,30 @@ def URL(self) -> str: def _remove_container(self, container_name: str) -> None: """Helper method to remove a container by name.""" try: - logging.info(f"Attempting to remove container {container_name}.") + logger.info(f"Attempting to remove container {container_name}.") container = self.client.containers.get(container_name) container.remove(force=True) - logging.info(f"Removed container {container_name}.") + logger.info(f"Removed container {container_name}.") except NotFound: - logging.info(f"Container {container_name} not found for removal.") + logger.info(f"Container {container_name} not found for removal.") except APIError as e: - logging.error(f"Docker API error removing container {container_name}: {e}") + logger.error(f"Docker API error removing container {container_name}: {e}") except Exception as e: - logging.exception(f"Unexpected error removing container {container_name}: {e}") + logger.exception(f"Unexpected error removing container {container_name}: {e}") def _remove_network(self, network_name: str) -> None: """Helper method to remove a network by name.""" try: - logging.info(f"Attempting to remove network {network_name}.") + logger.info(f"Attempting to remove network {network_name}.") net = self.client.networks.get(network_name) net.remove() - logging.info(f"Removed network {network_name}.") + logger.info(f"Removed network {network_name}.") except NotFound: - logging.info(f"Network {network_name} not found for removal.") + logger.info(f"Network {network_name} not found for removal.") except APIError as e: - logging.error(f"Docker API error removing network {network_name}: {e}") + logger.error(f"Docker API error removing network {network_name}: {e}") except Exception as e: - logging.exception(f"Unexpected error removing network {network_name}: {e}") + logger.exception(f"Unexpected error removing network {network_name}: {e}") @override def cleanup_hub(self) -> None: @@ -63,9 +63,9 @@ def cleanup_browsers(self) -> None: for container in containers: self._remove_container(container.name) except APIError as e: - logging.error(f"Docker API error listing browser containers: {e}") + logger.error(f"Docker API error listing browser containers: {e}") except Exception as e: - logging.exception(f"Unexpected error cleaning up browser containers: {e}") + logger.exception(f"Unexpected error cleaning up browser containers: {e}") @override def cleanup(self) -> None: @@ -78,22 +78,22 @@ async def ensure_hub_running(self) -> bool: # Ensure network exists try: self.client.networks.get(self.settings.docker.DOCKER_NETWORK_NAME) - logging.info( + logger.info( f"Docker network '{self.settings.docker.DOCKER_NETWORK_NAME}' already exists." ) except NotFound: - logging.info( + logger.info( f"Docker network '{self.settings.docker.DOCKER_NETWORK_NAME}' not found, creating." ) self.client.networks.create(self.settings.docker.DOCKER_NETWORK_NAME, driver="bridge") - logging.info(f"Docker network '{self.settings.docker.DOCKER_NETWORK_NAME}' created.") + logger.info(f"Docker network '{self.settings.docker.DOCKER_NETWORK_NAME}' created.") except APIError as e: - logging.error( + logger.error( f"Docker API error ensuring network '{self.settings.docker.DOCKER_NETWORK_NAME}': {e}" ) return False except Exception as e: - logging.exception( + logger.exception( f"Unexpected error ensuring network '{self.settings.docker.DOCKER_NETWORK_NAME}': {e}" ) return False @@ -102,15 +102,15 @@ async def ensure_hub_running(self) -> bool: try: hub = self.client.containers.get(self.settings.HUB_NAME) if hub.status != "running": - logging.info( + logger.info( f"{self.settings.HUB_NAME} container found but not running, restarting." ) hub.restart() - logging.info(f"{self.settings.HUB_NAME} container restarted.") + logger.info(f"{self.settings.HUB_NAME} container restarted.") else: - logging.info(f"{self.settings.HUB_NAME} container is already running.") + logger.info(f"{self.settings.HUB_NAME} container is already running.") except NotFound: - logging.info(f"{self.settings.HUB_NAME} container not found, creating.") + logger.info(f"{self.settings.HUB_NAME} container not found, creating.") self.client.containers.run( self.settings.selenium_grid.HUB_IMAGE, name=self.settings.HUB_NAME, @@ -137,12 +137,12 @@ async def ensure_hub_running(self) -> bool: cpu_quota=int(0.5 * 100000), # Convert to microseconds cpu_period=100000, # 100ms period ) - logging.info(f"{self.settings.HUB_NAME} container created and started.") + logger.info(f"{self.settings.HUB_NAME} container created and started.") except APIError as e: - logging.error(f"Docker API error ensuring {self.settings.HUB_NAME} container: {e}") + logger.error(f"Docker API error ensuring {self.settings.HUB_NAME} container: {e}") return False except Exception as e: - logging.exception(f"Unexpected error ensuring {self.settings.HUB_NAME} container: {e}") + logger.exception(f"Unexpected error ensuring {self.settings.HUB_NAME} container: {e}") return False return True @@ -161,21 +161,21 @@ async def create_browsers( # Ensure image exists, pull if necessary try: self.client.images.get(config.image) - logging.info(f"Docker image {config.image} already exists.") + logger.info(f"Docker image {config.image} already exists.") except NotFound: - logging.info(f"Docker image {config.image} not found, pulling.") + logger.info(f"Docker image {config.image} not found, pulling.") self.client.images.pull(config.image) - logging.info(f"Docker image {config.image} pulled.") + logger.info(f"Docker image {config.image} pulled.") except APIError as e: - logging.error(f"Docker API error ensuring image {config.image}: {e}") + logger.error(f"Docker API error ensuring image {config.image}: {e}") continue except Exception as e: - logging.exception(f"Unexpected error ensuring image {config.image}: {e}") + logger.exception(f"Unexpected error ensuring image {config.image}: {e}") continue # Create and run container try: - logging.info(f"Creating container for browser type {browser_type}.") + logger.info(f"Creating container for browser type {browser_type}.") container = self.client.containers.run( config.image, detach=True, @@ -201,15 +201,15 @@ async def create_browsers( ) cid = getattr(container, "id", None) if not cid: - logging.error("Failed to start browser container or retrieve container ID.") + logger.error("Failed to start browser container or retrieve container ID.") continue browser_ids.append(cid[:12]) - logging.info(f"Created container with ID: {cid[:12]}") + logger.info(f"Created container with ID: {cid[:12]}") except APIError as e: - logging.error(f"Docker API error creating container for {browser_type}: {e}") + logger.error(f"Docker API error creating container for {browser_type}: {e}") continue except Exception as e: - logging.exception(f"Unexpected error creating container for {browser_type}: {e}") + logger.exception(f"Unexpected error creating container for {browser_type}: {e}") continue return browser_ids diff --git a/src/app/services/selenium_hub/core/hub_backend.py b/src/app/services/selenium_hub/core/hub_backend.py index 1e1e52d..971e3b7 100644 --- a/src/app/services/selenium_hub/core/hub_backend.py +++ b/src/app/services/selenium_hub/core/hub_backend.py @@ -1,11 +1,11 @@ import asyncio -import logging from abc import ABC, abstractmethod from typing import Any from urllib.parse import urljoin import httpx +from ..common.logger import logger from ..models.browser import BrowserConfigs, BrowserType @@ -65,26 +65,26 @@ async def check_hub_health(self, username: str, password: str) -> bool: Returns True if the hub responds with 200 OK, False otherwise. """ url = urljoin(self.URL, "status") - logging.info(f"{self.__class__.__name__}: Checking health for {url}") + logger.info(f"{self.__class__.__name__}: Checking health for {url}") auth = httpx.BasicAuth(username, password) try: # Use a longer timeout for health checks to allow for startup time async with httpx.AsyncClient(timeout=httpx.Timeout(10.0), auth=auth) as client: response = await client.get(url) if response.status_code == httpx.codes.OK: - logging.info("Health check sucess.") + logger.info("Health check SUCCEED!") return True else: try: response_body = response.text - logging.warning( + logger.warning( f"Health check failed with status code: {response.status_code}, response body: {response_body}" ) except Exception: - logging.warning( + logger.warning( f"Health check failed with status code: {response.status_code}, could not read response body" ) return False except httpx.RequestError as e: - logging.warning(f"Health check request failed: {e}") + logger.warning(f"Health check request failed: {e or 'No error message.'}") return False diff --git a/src/app/services/selenium_hub/core/kubernetes/backend.py b/src/app/services/selenium_hub/core/kubernetes/backend.py index 734a713..e8b5e41 100644 --- a/src/app/services/selenium_hub/core/kubernetes/backend.py +++ b/src/app/services/selenium_hub/core/kubernetes/backend.py @@ -1,4 +1,3 @@ -import logging import uuid from os import environ from typing import Any, Callable, override @@ -32,6 +31,7 @@ V1ServiceSpec, ) +from ...common.logger import logger from ...models.browser import BrowserConfig, BrowserConfigs, BrowserType from ...models.general_settings import SeleniumHubGeneralSettings from ..hub_backend import HubBackend @@ -91,26 +91,26 @@ def cleanup_hub(self) -> None: ResourceType.DEPLOYMENT, self.settings.kubernetes.SELENIUM_GRID_SERVICE_NAME ) except Exception as e: - logging.exception(f"Exception during deletion of deployment: {e}") + logger.exception(f"Exception during deletion of deployment: {e}") try: self.resource_manager.delete_resource( ResourceType.SERVICE, self.settings.kubernetes.SELENIUM_GRID_SERVICE_NAME ) except Exception as e: - logging.exception(f"Exception during deletion of service: {e}") + logger.exception(f"Exception during deletion of service: {e}") @handle_kubernetes_exceptions(ErrorStrategy.GRACEFUL) def cleanup_browsers(self) -> None: """Clean up all browser pods.""" - logging.info( + logger.info( f"Deleting {self.settings.NODE_LABEL} pods in namespace {self.settings.kubernetes.NAMESPACE}..." ) self.k8s_core.delete_collection_namespaced_pod( namespace=self.settings.kubernetes.NAMESPACE, label_selector=f"app={self.settings.NODE_LABEL}", ) - logging.info(f"{self.settings.NODE_LABEL} pods delete request sent.") + logger.info(f"{self.settings.NODE_LABEL} pods delete request sent.") @override def cleanup(self) -> None: @@ -133,12 +133,12 @@ def _validate_deployment_config(self, deployment: V1Deployment) -> bool: or deployment.spec.template.spec is None or not deployment.spec.template.spec.security_context ): - logging.warning("Deployment missing security context") + logger.warning("Deployment missing security context") return False return True except Exception as e: - logging.error(f"Error validating deployment: {e}") + logger.error(f"Error validating deployment: {e}") return False def _has_valid_spec_structure(self, deployment: V1Deployment) -> bool: @@ -148,7 +148,7 @@ def _has_valid_spec_structure(self, deployment: V1Deployment) -> bool: or deployment.spec.template is None or deployment.spec.template.spec is None ): - logging.warning("Invalid deployment spec structure") + logger.warning("Invalid deployment spec structure") return False return True @@ -160,16 +160,16 @@ def _has_valid_resource_limits(self, deployment: V1Deployment) -> bool: or deployment.spec.template.spec is None or deployment.spec.template.spec.containers is None ): - logging.warning("Invalid deployment spec structure for resource limits") + logger.warning("Invalid deployment spec structure for resource limits") return False for container in deployment.spec.template.spec.containers: if not container.resources or not container.resources.limits: - logging.warning("Deployment missing resource limits") + logger.warning("Deployment missing resource limits") return False required_limits = ["cpu", "memory"] if not all(key in container.resources.limits for key in required_limits): - logging.warning("Deployment missing required resource limits") + logger.warning("Deployment missing required resource limits") return False return True @@ -177,26 +177,26 @@ def _validate_service_config(self, service: V1Service) -> bool: """Validate service configuration.""" try: if not service.spec: - logging.warning("Invalid service spec structure") + logger.warning("Invalid service spec structure") return False if service.spec.type not in ["ClusterIP", "NodePort"]: - logging.warning("Invalid service type") + logger.warning("Invalid service type") return False if not service.spec.ports: - logging.warning("Service missing port configuration") + logger.warning("Service missing port configuration") return False required_attrs = ["port", "target_port"] for port in service.spec.ports: if not all(hasattr(port, attr) for attr in required_attrs): - logging.warning("Service port missing required attributes") + logger.warning("Service port missing required attributes") return False return True except Exception as e: - logging.error(f"Error validating service: {e}") + logger.error(f"Error validating service: {e}") return False @handle_kubernetes_exceptions(ErrorStrategy.STRICT) @@ -210,10 +210,10 @@ def ensure_resource_exists( """Generic method to ensure a resource exists.""" try: self.resource_manager.read_resource(resource_type, name) - logging.info(f"{name} {resource_type.value} already exists.") + logger.info(f"{name} {resource_type.value} already exists.") except ApiException as e: if e.status == HTTP_NOT_FOUND: - logging.info(f"{name} {resource_type.value} not found, creating...") + logger.info(f"{name} {resource_type.value} not found, creating...") resource = create_func() if validate_func and not validate_func(resource): @@ -230,7 +230,7 @@ def ensure_resource_exists( elif resource_type == ResourceType.NAMESPACE: self.k8s_core.create_namespace(body=resource) - logging.info(f"{name} {resource_type.value} created.") + logger.info(f"{name} {resource_type.value} created.") else: raise @@ -299,11 +299,11 @@ async def ensure_hub_running(self) -> bool: return True except Exception as e: - logging.exception(f"Attempt {i + 1} to ensure K8s hub failed: {e}") + logger.exception(f"Attempt {i + 1} to ensure K8s hub failed: {e}") if i < self.settings.kubernetes.MAX_RETRIES - 1: await self.resource_manager.sleep(i) else: - logging.exception("Max retries reached for ensuring K8s hub.") + logger.exception("Max retries reached for ensuring K8s hub.") return False return False @@ -368,15 +368,15 @@ async def create_browsers( browser_ids.append(pod_name) break except Exception as e: - logging.exception(f"Unexpected error creating browser pod: {e}") + logger.exception(f"Unexpected error creating browser pod: {e}") if i < self.settings.kubernetes.MAX_RETRIES - 1: await self.resource_manager.sleep(i) else: - logging.exception( + logger.exception( "Max retries reached for creating browser pod due to unexpected error." ) else: - logging.error("Failed to create browser pod after all retries.") + logger.error("Failed to create browser pod after all retries.") return browser_ids @@ -384,7 +384,7 @@ async def create_browsers( async def _create_browser_pod_with_retry(self, pod_name: str, pod: V1Pod, attempt: int) -> None: """Create a browser pod with retry logic.""" self.k8s_core.create_namespaced_pod(namespace=self.settings.kubernetes.NAMESPACE, body=pod) - logging.info(f"Pod {pod_name} created.") + logger.info(f"Pod {pod_name} created.") def _create_browser_pod( self, pod_name: str, browser_type: BrowserType, config: BrowserConfig @@ -611,4 +611,4 @@ def _stop_service_port_forward(self) -> None: if self.port_forward_manager: self.port_forward_manager.stop() self.port_forward_manager = None - logging.info("Stopped kubectl service port-forward.") + logger.info("Stopped kubectl service port-forward.") diff --git a/src/app/services/selenium_hub/core/kubernetes/common/decorators.py b/src/app/services/selenium_hub/core/kubernetes/common/decorators.py index 1fc0532..fab842e 100644 --- a/src/app/services/selenium_hub/core/kubernetes/common/decorators.py +++ b/src/app/services/selenium_hub/core/kubernetes/common/decorators.py @@ -1,13 +1,13 @@ """Decorators for Kubernetes operations.""" import asyncio -import logging from enum import Enum from functools import wraps from typing import Any, Callable, ParamSpec, TypeVar from kubernetes.client.exceptions import ApiException +from ....common.logger import logger from .constants import HTTP_NOT_FOUND # Type variables for generic decorator @@ -29,29 +29,29 @@ def _handle_exceptions(func_name: str, strategy: ErrorStrategy, e: Exception) -> if e.status == HTTP_NOT_FOUND: match strategy: case ErrorStrategy.GRACEFUL: - logging.info(f"Resource not found in {func_name}: {e}") + logger.info(f"Resource not found in {func_name}: {e}") return None case ErrorStrategy.RETURN_FALSE: - logging.info(f"Resource not found in {func_name}: {e}") + logger.info(f"Resource not found in {func_name}: {e}") return False case _: # STRICT - logging.info(f"Resource not found in {func_name}: {e}") + logger.info(f"Resource not found in {func_name}: {e}") raise e else: match strategy: case ErrorStrategy.RETURN_FALSE: - logging.error(f"Kubernetes API error in {func_name}: {e}") + logger.error(f"Kubernetes API error in {func_name}: {e}") return False case _: # STRICT or GRACEFUL - logging.error(f"Kubernetes API error in {func_name}: {e}") + logger.error(f"Kubernetes API error in {func_name}: {e}") raise e else: match strategy: case ErrorStrategy.RETURN_FALSE: - logging.exception(f"Error in {func_name}: {e}") + logger.exception(f"Error in {func_name}: {e}") return False case _: # STRICT or GRACEFUL - logging.exception(f"Unexpected error in {func_name}: {e}") + logger.exception(f"Unexpected error in {func_name}: {e}") raise e diff --git a/src/app/services/selenium_hub/core/kubernetes/k8s_config.py b/src/app/services/selenium_hub/core/kubernetes/k8s_config.py index cd427a5..0c0f672 100644 --- a/src/app/services/selenium_hub/core/kubernetes/k8s_config.py +++ b/src/app/services/selenium_hub/core/kubernetes/k8s_config.py @@ -1,4 +1,3 @@ -import logging import os from kubernetes.client import CoreV1Api, V1ObjectMeta @@ -6,6 +5,7 @@ from kubernetes.config.incluster_config import load_incluster_config from kubernetes.config.kube_config import load_kube_config +from ...common.logger import logger from ...models.kubernetes_settings import KubernetesSettings @@ -21,24 +21,24 @@ def __init__(self, k8s_settings: KubernetesSettings) -> None: def _load_config(self) -> None: try: try: - logging.info("Trying in-cluster config...") + logger.info("Trying in-cluster config...") load_incluster_config() - logging.info("Loaded in-cluster config.") + logger.info("Loaded in-cluster config.") except ConfigException: kubeconfig_path = self.k8s_settings.KUBECONFIG - logging.info(f"In-cluster config failed, trying kubeconfig: '{kubeconfig_path}'") + logger.info(f"In-cluster config failed, trying kubeconfig: '{kubeconfig_path}'") if kubeconfig_path: - logging.info(f"KUBECONFIG path: {kubeconfig_path}") - logging.info(f"KUBECONFIG exists: {os.path.exists(kubeconfig_path)}") + logger.info(f"KUBECONFIG path: {kubeconfig_path}") + logger.info(f"KUBECONFIG exists: {os.path.exists(kubeconfig_path)}") else: - logging.info("KUBECONFIG is empty, will use default kubeconfig resolution.") + logger.info("KUBECONFIG is empty, will use default kubeconfig resolution.") load_kube_config( config_file=kubeconfig_path, context=self.k8s_settings.CONTEXT, ) - logging.info(f"Loaded kubeconfig from {kubeconfig_path}") + logger.info(f"Loaded kubeconfig from {kubeconfig_path}") except Exception as e: - logging.exception(f"Failed to load Kubernetes configuration: {e}") + logger.exception(f"Failed to load Kubernetes configuration: {e}") raise def _detect_kind_cluster(self) -> None: @@ -53,7 +53,7 @@ def _detect_kind_cluster(self) -> None: self._is_kind = True break if self._is_kind: - logging.info("KinD cluster detected via node name suffix '-control-plane'.") + logger.info("KinD cluster detected via node name suffix '-control-plane'.") except Exception: self._is_kind = False diff --git a/src/app/services/selenium_hub/core/kubernetes/k8s_port_forwarder.py b/src/app/services/selenium_hub/core/kubernetes/k8s_port_forwarder.py index fa90ef3..6a153a0 100644 --- a/src/app/services/selenium_hub/core/kubernetes/k8s_port_forwarder.py +++ b/src/app/services/selenium_hub/core/kubernetes/k8s_port_forwarder.py @@ -6,6 +6,7 @@ from threading import Thread from typing import Awaitable, Callable +from ...common.logger import logger from ...common.pidfile import PidFile, is_process_running_with_cmdline, terminate_pid @@ -51,9 +52,9 @@ def __init__( # noqa: PLR0913 # Consider refactoring to use a config object or self.pidfile = PidFile(pid_dir / f"{service_name}-{local_port}.pid") @staticmethod - def _start_logging_thread(process: Popen[str]) -> None: + def _start_logger_thread(process: Popen[str]) -> None: def log_message(level: int, line: str) -> None: - logging.log(level, f"kubectl port-forward: {line.strip()}") + logger.log(level, f"kubectl port-forward: {line.strip()}") # Drain stdout in a thread to avoid blocking def log_output() -> None: @@ -84,7 +85,7 @@ def _build_cmd_args(self) -> list[str]: def _kubectl_port_foward(self) -> Popen[str]: cmd = self._build_cmd_args() - logging.info(f"Executing: {' '.join(cmd)}") + logger.info(f"Executing: {' '.join(cmd)}") # Start subprocess with live stdout capturing process = Popen( # noqa: S603 @@ -96,7 +97,7 @@ def _kubectl_port_foward(self) -> Popen[str]: shell=False, ) - self._start_logging_thread(process) + self._start_logger_thread(process) return process @@ -110,7 +111,7 @@ def _is_existing_port_forward_alive(self) -> bool: is_running = is_process_running_with_cmdline(pid, self._build_cmd_args()) if not is_running: - logging.warning( + logger.warning( f"PID {pid} exists but process does not match expected command line. " "It may be stale or from another context." ) @@ -121,7 +122,7 @@ def _is_existing_port_forward_alive(self) -> bool: def _start_port_forward(self) -> Popen[str] | None: if self._is_existing_port_forward_alive(): - logging.info( + logger.info( f"Port-forward for {self.service_name} on port {self.local_port} already running." ) return None @@ -130,12 +131,12 @@ def _start_port_forward(self) -> Popen[str] | None: process = self._kubectl_port_foward() # Write PID self.pidfile.write(process.pid) - logging.info(f"Started kubectl port-forward process (PID: {process.pid})") + logger.info(f"Started kubectl port-forward process (PID: {process.pid})") except FileNotFoundError: - logging.error("kubectl not found! Please install and add to PATH.") + logger.error("kubectl not found! Please install and add to PATH.") except Exception as e: - logging.error(f"Error starting kubectl port-forward: {e}") + logger.error(f"Error starting kubectl port-forward: {e}") else: return process return None @@ -143,16 +144,16 @@ def _start_port_forward(self) -> Popen[str] | None: async def start(self) -> bool: if self._is_existing_port_forward_alive(): if await self.check_health(): - logging.info("Port-forward already running and healthy.") + logger.info("Port-forward already running and healthy.") return True else: - logging.warning( + logger.warning( "Existing port-forward is running but health check failed, cleaning up." ) self.stop() for attempt in range(1, self.max_retries + 1): - logging.info(f"Attempt {attempt} to start port-forward...") + logger.info(f"Attempt {attempt} to start port-forward...") self.process = self._start_port_forward() if not self.process: @@ -163,13 +164,13 @@ async def start(self) -> bool: try: is_alive = self.process.poll() is None exit_code = self.process.returncode - logging.debug(f"Process returned: {is_alive}, exit_code: {exit_code}") + logger.debug(f"Process returned: {is_alive}, exit_code: {exit_code}") except Exception as exc: - logging.error(f"Error checking is_alive(): {exc}") + logger.error(f"Error checking is_alive(): {exc}") is_alive = False exit_code = None if not is_alive: - logging.error( + logger.error( "kubectl port-forward exited immediately. exit_code: %r (type %s)", exit_code, type(exit_code).__name__ if exit_code is not None else "NoneType", @@ -179,21 +180,21 @@ async def start(self) -> bool: await asyncio.sleep(2) continue - logging.info("Process still alive, checking health.") + logger.info("Process still alive, checking health.") if await self.check_health(): - logging.info("Port-forward started and health check passed.") + logger.info("Port-forward started and health check passed.") return True - logging.warning("Health check failed, stopping port-forward and retrying.") + logger.warning("Health check failed, stopping port-forward and retrying.") self.stop() await asyncio.sleep(2) - logging.error("Failed to start port-forward after retries.") + logger.error("Failed to start port-forward after retries.") return False def stop(self) -> None: if self.process: - logging.info("Terminating port-forward process...") + logger.info("Terminating port-forward process...") self.process.terminate() try: self.process.wait(timeout=5) @@ -205,4 +206,4 @@ def stop(self) -> None: if pid is not None: terminate_pid(pid) self.pidfile.remove() - logging.info("Port-forward stopped.") + logger.info("Port-forward stopped.") diff --git a/src/app/services/selenium_hub/core/kubernetes/k8s_resource_manager.py b/src/app/services/selenium_hub/core/kubernetes/k8s_resource_manager.py index e3f509f..4c3b80c 100644 --- a/src/app/services/selenium_hub/core/kubernetes/k8s_resource_manager.py +++ b/src/app/services/selenium_hub/core/kubernetes/k8s_resource_manager.py @@ -1,5 +1,4 @@ import asyncio -import logging import time from typing import Callable @@ -12,6 +11,7 @@ V1Service, ) +from ...common.logger import logger from ...models.kubernetes_settings import KubernetesSettings from .common.constants import ( DEFAULT_POLL_INTERVAL, @@ -69,7 +69,7 @@ def is_service_ready( return False except Exception as e: - logging.debug(f"Error checking endpoints for service {name}: {e}") + logger.debug(f"Error checking endpoints for service {name}: {e}") return False @@ -148,20 +148,20 @@ def delete_resource(self, resource_type: ResourceType, name: str) -> None: def _wait_for_deletion(self, resource_type: ResourceType, name: str) -> None: """Wait for resource deletion to complete.""" - logging.info(f"Waiting for {resource_type.value} {name} to be deleted...") + logger.info(f"Waiting for {resource_type.value} {name} to be deleted...") for _ in range(self.max_retries): try: self.read_resource(resource_type, name) time.sleep(self.retry_delay) except ApiException as e: if e.status == HTTP_NOT_FOUND: - logging.info(f"{resource_type.value} {name} deleted successfully.") + logger.info(f"{resource_type.value} {name} deleted successfully.") return raise except Exception as e: - logging.exception(f"Error waiting for {resource_type.value} {name} deletion: {e}") + logger.exception(f"Error waiting for {resource_type.value} {name} deletion: {e}") raise - logging.warning(f"Timeout waiting for {resource_type.value} {name} to be deleted.") + logger.warning(f"Timeout waiting for {resource_type.value} {name} to be deleted.") async def wait_for_resource_ready( self, @@ -174,17 +174,17 @@ async def wait_for_resource_ready( if config is None: config = WaitConfig() - logging.info(f"Waiting for {resource_type.value} {name} to be ready...") + logger.info(f"Waiting for {resource_type.value} {name} to be ready...") try: async with asyncio.timeout(config.timeout_seconds): while True: if await self._check_resource_ready(resource_type, name, check_ready_func): - logging.info(f"{resource_type.value} {name} is ready.") + logger.info(f"{resource_type.value} {name} is ready.") return await asyncio.sleep(config.poll_interval or DEFAULT_POLL_INTERVAL) except asyncio.TimeoutError: - logging.error(f"Timeout waiting for {resource_type.value} {name} to be ready.") + logger.error(f"Timeout waiting for {resource_type.value} {name} to be ready.") raise async def _check_resource_ready( @@ -202,12 +202,12 @@ async def _check_resource_ready( return self._is_resource_ready_by_type(resource_type, resource, name) except ApiException as e: if e.status == HTTP_NOT_FOUND: - logging.info(f"{resource_type.value} {name} not found yet, waiting...") + logger.info(f"{resource_type.value} {name} not found yet, waiting...") return False - logging.error(f"Error checking {resource_type.value} {name} readiness: {e}") + logger.error(f"Error checking {resource_type.value} {name} readiness: {e}") return False except Exception as e: - logging.exception( + logger.exception( f"Unexpected error during polling for {resource_type.value} {name}: {e}" ) return False @@ -226,11 +226,11 @@ def _is_resource_ready_by_type( case ResourceType.NAMESPACE: return is_namespace_ready(resource, name) case _: - logging.warning(f"Unknown resource type {resource_type.value} for {name}") + logger.warning(f"Unknown resource type {resource_type.value} for {name}") return False async def sleep(self, attempt: int) -> None: """Sleep with exponential backoff.""" delay = self.retry_delay * (2**attempt) - logging.info(f"Retrying in {delay} seconds...", stacklevel=2) + logger.info(f"Retrying in {delay} seconds...", stacklevel=2) await asyncio.sleep(delay) diff --git a/src/app/services/selenium_hub/core/kubernetes/k8s_url_resolver.py b/src/app/services/selenium_hub/core/kubernetes/k8s_url_resolver.py index 80fdb11..b388209 100644 --- a/src/app/services/selenium_hub/core/kubernetes/k8s_url_resolver.py +++ b/src/app/services/selenium_hub/core/kubernetes/k8s_url_resolver.py @@ -1,9 +1,9 @@ -import logging import time from os import environ from kubernetes.client import CoreV1Api +from ...common.logger import logger from ...models.general_settings import SeleniumHubGeneralSettings from .common.decorators import ErrorStrategy, handle_kubernetes_exceptions @@ -27,14 +27,14 @@ def get_hub_url(self) -> str: if self._is_kind: # For KinD environments, use port-forwarded URL FALLBACK_URL = f"http://localhost:{self.settings.kubernetes.PORT_FORWARD_LOCAL_PORT}" - logging.info(f"Using port-forwarded URL for KinD: {FALLBACK_URL}") + logger.info(f"Using port-forwarded URL for KinD: {FALLBACK_URL}") return FALLBACK_URL return self._get_nodeport_url(FALLBACK_URL) def _get_in_cluster_url(self) -> str: url = f"http://{self.settings.kubernetes.SELENIUM_GRID_SERVICE_NAME}.{self.settings.kubernetes.NAMESPACE}.svc.cluster.local:{self.settings.selenium_grid.SELENIUM_HUB_PORT}" - logging.info(f"Using in-cluster DNS for Selenium Hub URL: {url}") + logger.info(f"Using in-cluster DNS for Selenium Hub URL: {url}") return url def _get_nodeport_url(self, fallback_url: str) -> str: @@ -46,11 +46,11 @@ def _get_nodeport_url(self, fallback_url: str) -> str: if url: return url if attempt < max_retries - 1: - logging.info(f"NodePort not yet assigned, retrying in {retry_delay} seconds...") + logger.info(f"NodePort not yet assigned, retrying in {retry_delay} seconds...") time.sleep(retry_delay) retry_delay *= 2 except Exception as e: - logging.error( + logger.error( f"Error getting NodePort URL (attempt {attempt + 1}/{max_retries}): {e}. Fallback: {fallback_url}" ) if attempt < max_retries - 1: @@ -62,25 +62,25 @@ def _get_nodeport_url(self, fallback_url: str) -> str: @handle_kubernetes_exceptions(ErrorStrategy.STRICT) def _try_get_nodeport_url(self, attempt: int, max_retries: int) -> str | None: - logging.info( + logger.info( f"Attempt {attempt + 1}/{max_retries}: Getting NodePort for service {self.settings.kubernetes.SELENIUM_GRID_SERVICE_NAME} in namespace {self.settings.kubernetes.NAMESPACE}" ) service = self.k8s_core.read_namespaced_service( name=self.settings.kubernetes.SELENIUM_GRID_SERVICE_NAME, namespace=self.settings.kubernetes.NAMESPACE, ) - logging.info(f"Service type: {service.spec.type if service.spec else 'None'}") - logging.info( + logger.info(f"Service type: {service.spec.type if service.spec else 'None'}") + logger.info( f"Service ports: {service.spec.ports if service.spec and service.spec.ports else 'None'}" ) if not service.spec or not service.spec.ports: return None for port in service.spec.ports: - logging.info( + logger.info( f"Checking port: port={port.port}, target_port={port.target_port}, node_port={port.node_port}" ) if port.node_port: url = f"http://localhost:{port.node_port}" - logging.info(f"Resolved NodePort URL: {url}") + logger.info(f"Resolved NodePort URL: {url}") return url return None From a15767b582ebcd393f30a3a45a33bc30c246c493 Mon Sep 17 00:00:00 2001 From: falamarcao Date: Fri, 8 Aug 2025 23:59:33 -0300 Subject: [PATCH 2/3] =?UTF-8?q?=F0=9F=93=A1=20feat:=20redirect=20/=20to=20?= =?UTF-8?q?/mcp=20or=20/sse=20by=20Accept=20header=20&=20load=20settings?= =?UTF-8?q?=20from=20TOML?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Added root route redirecting to MCP or SSE endpoints based on Accept header (fixed MCP Clients usage). - Centralized pyproject.toml value loading for VERSION, and description. ✅ Local Tests: all 116 tests passed successfully. --- config.yaml | 3 +-- pyproject.toml | 6 ++--- src/app/common/toml.py | 43 ++++++++++++++++++++++++++++++++++++ src/app/core/settings.py | 17 ++++++++++----- src/app/main.py | 47 ++++++++++++++++++++++++++++++++++------ uv.lock | 6 ++--- 6 files changed, 102 insertions(+), 20 deletions(-) create mode 100644 src/app/common/toml.py diff --git a/config.yaml b/config.yaml index 2e4e6a1..2afe914 100644 --- a/config.yaml +++ b/config.yaml @@ -1,5 +1,4 @@ -project_name: MCP Selenium Server -version: 0.1.0 +project_name: MCP Selenium Grid deployment_mode: docker # one of: docker, kubernetes (DeploymentMode enum values) api_v1_str: /api/v1 api_token: CHANGE_ME diff --git a/pyproject.toml b/pyproject.toml index 8292052..3bea913 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,8 +13,8 @@ [project] name = "mcp-selenium-grid" -version = "0.1.0.dev2" -description = "MCP Server for managing Selenium Grid instances" +version = "0.1.0.dev3" +description = "MCP Server for managing Selenium Grid" readme = "README.md" license = { file = "LICENSE" } requires-python = ">=3.12" @@ -25,7 +25,7 @@ dependencies = [ "fastapi[standard]>=0.115.14", # Web framework "fastapi-cli[standard-no-fastapi-cloud-cli]>=0.0.8", "fastapi-mcp>=0.3.4", # MCP integration for FastAPI - "pydantic[email]>=2.11.7", # Data validation, email support + "pydantic>=2.11.7", # Data validation "pydantic-settings>=2.10.1", # Settings management (latest is 2.2.1) "docker>=7.1.0", # Docker API client "kubernetes>=33.1.0", # Kubernetes API client3 diff --git a/src/app/common/toml.py b/src/app/common/toml.py new file mode 100644 index 0000000..9f78e9f --- /dev/null +++ b/src/app/common/toml.py @@ -0,0 +1,43 @@ +from os import getcwd +from pathlib import Path +from tomllib import load +from typing import Any + +ROOT_DIR = Path(getcwd()).resolve() + + +def load_value_from_toml( + keys: list[str], file_path: Path = ROOT_DIR / "pyproject.toml", default: Any = None +) -> Any: + """ + Load a nested value from a TOML file. + + Args: + keys: List of nested keys to traverse. + file_path: Path to the TOML file. + default: Value to return if keys not found. + + Returns: + The value from the TOML file. + + Raises: + FileNotFoundError: If the file doesn't exist and no default is provided. + ValueError: If the keys are missing and no default is provided. + """ + if not file_path.exists(): + if default is not None: + return default + raise FileNotFoundError(f"{file_path} not found") + + try: + with file_path.open("rb") as f: + data = load(f) + for key in keys: + data = data[key] + return data + except KeyError: + if default is not None: + return default + raise ValueError(f"Keys {'.'.join(keys)} not found in {file_path}") + except Exception as e: + raise ValueError(f"Error reading {file_path}: {e}") from e diff --git a/src/app/core/settings.py b/src/app/core/settings.py index efc1d96..428b15e 100644 --- a/src/app/core/settings.py +++ b/src/app/core/settings.py @@ -1,7 +1,8 @@ """Core settings for MCP Server.""" -from pydantic import Field, SecretStr +from pydantic import Field, SecretStr, field_validator +from app.common.toml import load_value_from_toml from app.services.selenium_hub.models.general_settings import SeleniumHubGeneralSettings @@ -9,12 +10,18 @@ class Settings(SeleniumHubGeneralSettings): """MCP Server settings.""" # API Settings - PROJECT_NAME: str = Field(default="MCP Selenium Server") - VERSION: str = Field(default="0.1.0") - API_V1_STR: str = Field(default="/api/v1") + PROJECT_NAME: str = "MCP Selenium Grid" + VERSION: str = "" + + @field_validator("VERSION", mode="before") + @classmethod + def load_version_from_pyproject(cls, v: str) -> str: + return v or load_value_from_toml(["project", "version"]) + + API_V1_STR: str = "/api/v1" # API Token - API_TOKEN: SecretStr = Field(default=SecretStr("CHANGE_ME")) + API_TOKEN: SecretStr = SecretStr("CHANGE_ME") # Security Settings BACKEND_CORS_ORIGINS: list[str] = Field( diff --git a/src/app/main.py b/src/app/main.py index 2cb7111..1a51f3a 100644 --- a/src/app/main.py +++ b/src/app/main.py @@ -1,4 +1,4 @@ -"""MCP Server for managing Selenium Grid instances.""" +"""MCP Server for managing Selenium Grid.""" import asyncio from contextlib import asynccontextmanager @@ -6,13 +6,16 @@ from fastapi import Depends, FastAPI, HTTPException, Request, status from fastapi.middleware.cors import CORSMiddleware +from fastapi.responses import JSONResponse, RedirectResponse from fastapi.security import HTTPAuthorizationCredentials from fastapi_mcp import FastApiMCP from prometheus_client import generate_latest from prometheus_fastapi_instrumentator import Instrumentator from starlette.responses import Response +from app.common.toml import load_value_from_toml from app.dependencies import get_settings, verify_token +from app.logger import logger from app.models import HealthCheckResponse, HealthStatus, HubStatusResponse from app.routers.browsers import router as browsers_router from app.routers.selenium_proxy import router as selenium_proxy_router @@ -23,6 +26,7 @@ def create_application() -> FastAPI: """Create FastAPI application for MCP.""" # Initialize settings once at the start settings = get_settings() + DESCRIPTION = load_value_from_toml(["project", "description"]) @asynccontextmanager async def lifespan(app: FastAPI) -> AsyncGenerator[None, Any]: @@ -52,14 +56,12 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, Any]: yield # --- Server shutdown: remove Selenium Hub resources (Docker or Kubernetes) --- - # manager = SeleniumHubManager(settings) - # manager.cleanup() hub.cleanup() app = FastAPI( title=settings.PROJECT_NAME, version=settings.VERSION, - description="MCP Server for managing Selenium Grid instances", + description=DESCRIPTION, lifespan=lifespan, ) @@ -77,7 +79,9 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, Any]: # Prometheus metrics endpoint @app.get("/metrics") - def metrics(credentials: HTTPAuthorizationCredentials = Depends(verify_token)) -> Response: + async def metrics( + credentials: HTTPAuthorizationCredentials = Depends(verify_token), + ) -> Response: return Response(generate_latest(), media_type="text/plain") # Health check endpoint @@ -127,8 +131,37 @@ async def get_hub_stats( app.include_router(selenium_proxy_router) # --- MCP Integration --- - mcp = FastApiMCP(app) - mcp.mount() # Mounts at /mcp by default + mcp = FastApiMCP( + app, + name="MCP Selenium Grid", + description=DESCRIPTION, + describe_full_response_schema=True, + describe_all_responses=True, + ) + MCP_HTTP_PATH = "/mcp" + MCP_SSE_PATH = "/sse" + mcp.mount_http(mount_path=MCP_HTTP_PATH) + mcp.mount_sse(mount_path=MCP_SSE_PATH) + + @app.api_route("/", methods=["GET", "POST"], include_in_schema=False) + async def root_redirect(request: Request) -> Response: + accept: str = request.headers.get("accept", "").lower() + method: str = request.method.upper() + + logger.info(f"Received {method=} with Accept: {accept}") + + if "text/event-stream" in accept: + # MCP allows POST or GET here + logger.info(f"Redirecting to SSE endpoint /sse (method={method})") + return RedirectResponse(url="/sse") + elif "application/json" in accept: + # JSON RPC endpoint (usually POST) + logger.info(f"Redirecting to HTTP JSON RPC endpoint /mcp (method={method})") + return RedirectResponse(url="/mcp") + else: + logger.warning(f"Unsupported Accept header or method: method={method}, accept={accept}") + return JSONResponse({"detail": "Unsupported Accept header or method"}, status_code=405) + # ---------------------- return app diff --git a/uv.lock b/uv.lock index 43c8084..ab3e614 100644 --- a/uv.lock +++ b/uv.lock @@ -583,7 +583,7 @@ wheels = [ [[package]] name = "mcp-selenium-grid" -version = "0.1.0.dev2" +version = "0.1.0.dev3" source = { editable = "." } dependencies = [ { name = "docker" }, @@ -595,7 +595,7 @@ dependencies = [ { name = "prometheus-client" }, { name = "prometheus-fastapi-instrumentator" }, { name = "psutil" }, - { name = "pydantic", extra = ["email"] }, + { name = "pydantic" }, { name = "pydantic-settings" }, { name = "typer" }, { name = "uvicorn" }, @@ -633,7 +633,7 @@ requires-dist = [ { name = "prometheus-client", specifier = ">=0.22.1" }, { name = "prometheus-fastapi-instrumentator", specifier = ">=7.1.0" }, { name = "psutil", specifier = ">=7.0.0" }, - { name = "pydantic", extras = ["email"], specifier = ">=2.11.7" }, + { name = "pydantic", specifier = ">=2.11.7" }, { name = "pydantic-settings", specifier = ">=2.10.1" }, { name = "pytest", marker = "extra == 'test'", specifier = ">=8.4.1" }, { name = "pytest-asyncio", marker = "extra == 'test'", specifier = ">=1.0.0" }, From 19e88abb13f6a372235f640f932469508d554842 Mon Sep 17 00:00:00 2001 From: falamarcao Date: Sat, 9 Aug 2025 04:01:56 -0300 Subject: [PATCH 3/3] =?UTF-8?q?=E2=9C=A8=20fix:=20improve=20MCP=20SSE/HTTP?= =?UTF-8?q?=20proxy=20handling=20with=20better=20redirects?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit FastApiMCP doesn't allow mounting at root `/`. MCP clients need to connect at root when using uvx, so we proxy requests dynamically to SSE and HTTP endpoints, handling methods and headers properly. --- README.md | 3 +- pyproject.toml | 2 +- src/app/common/fastapi_mcp.py | 25 ++++++++ src/app/{ => common}/logger.py | 0 src/app/main.py | 92 ++++++++++++++++++++---------- src/app/routers/browsers/routes.py | 2 +- src/app/routers/selenium_proxy.py | 2 +- 7 files changed, 93 insertions(+), 33 deletions(-) create mode 100644 src/app/common/fastapi_mcp.py rename src/app/{ => common}/logger.py (100%) diff --git a/README.md b/README.md index 6988332..20e178b 100644 --- a/README.md +++ b/README.md @@ -30,7 +30,8 @@ The MCP Selenium Grid provides a MCP Server for creating and managing browser in ### 📖 Usage -The MCP Selenium Grid provides a Web API for creating and managing browser instances. The server runs on `localhost:8000` and exposes MCP endpoints at `/mcp`. +The MCP Selenium Grid provides a Web API for creating and managing browser instances. The server runs on `localhost:8000` and exposes MCP endpoints at `/mcp` (Http Transport) and `/sse` (Server Sent Events). +> Note: All requests to the server root `http://localhost:8000` will be redirected to either `/mcp` or `/sse` endpoints, depending on the request. ### MCP Client Configuration diff --git a/pyproject.toml b/pyproject.toml index 3bea913..bd07a6b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -93,7 +93,7 @@ disallow_untyped_defs = true check_untyped_defs = true [[tool.mypy.overrides]] -module = ["fastapi_mcp"] +module = ["fastapi_mcp.*"] ignore_missing_imports = true [tool.pytest.ini_options] diff --git a/src/app/common/fastapi_mcp.py b/src/app/common/fastapi_mcp.py new file mode 100644 index 0000000..7eb9328 --- /dev/null +++ b/src/app/common/fastapi_mcp.py @@ -0,0 +1,25 @@ +from fastapi.requests import Request +from fastapi.responses import Response +from fastapi_mcp.transport.http import FastApiHttpSessionManager + +from .logger import logger + + +async def handle_fastapi_request( + name: str, + request: Request, + target_path: str, + method: str, + session_manager: FastApiHttpSessionManager, +) -> Response: + scope = dict(request.scope) + scope["path"] = target_path + scope["raw_path"] = target_path.encode("utf-8") + scope["query_string"] = request.scope.get("query_string", b"") + scope["method"] = method + + new_request = Request(scope, request.receive) + + logger.debug(f"Proxying internally to {name} at {target_path}") + response: Response = await session_manager.handle_fastapi_request(new_request) + return response diff --git a/src/app/logger.py b/src/app/common/logger.py similarity index 100% rename from src/app/logger.py rename to src/app/common/logger.py diff --git a/src/app/main.py b/src/app/main.py index 1a51f3a..007438f 100644 --- a/src/app/main.py +++ b/src/app/main.py @@ -3,30 +3,34 @@ import asyncio from contextlib import asynccontextmanager from typing import Any, AsyncGenerator +from urllib.parse import urljoin from fastapi import Depends, FastAPI, HTTPException, Request, status from fastapi.middleware.cors import CORSMiddleware -from fastapi.responses import JSONResponse, RedirectResponse +from fastapi.responses import JSONResponse, Response from fastapi.security import HTTPAuthorizationCredentials -from fastapi_mcp import FastApiMCP +from fastapi_mcp import AuthConfig, FastApiMCP from prometheus_client import generate_latest from prometheus_fastapi_instrumentator import Instrumentator -from starlette.responses import Response +from app.common.fastapi_mcp import handle_fastapi_request +from app.common.logger import logger from app.common.toml import load_value_from_toml from app.dependencies import get_settings, verify_token -from app.logger import logger from app.models import HealthCheckResponse, HealthStatus, HubStatusResponse from app.routers.browsers import router as browsers_router from app.routers.selenium_proxy import router as selenium_proxy_router from app.services.selenium_hub import SeleniumHub +DESCRIPTION = load_value_from_toml(["project", "description"]) +SETTINGS = get_settings() +MCP_HTTP_PATH = "/mcp" +MCP_SSE_PATH = "/sse" + def create_application() -> FastAPI: """Create FastAPI application for MCP.""" # Initialize settings once at the start - settings = get_settings() - DESCRIPTION = load_value_from_toml(["project", "description"]) @asynccontextmanager async def lifespan(app: FastAPI) -> AsyncGenerator[None, Any]: @@ -35,7 +39,7 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, Any]: app.state.browsers_instances_lock = asyncio.Lock() # Initialize Selenium Hub singleton - hub = SeleniumHub(settings) # This will create or return the singleton instance + hub = SeleniumHub(SETTINGS) # This will create or return the singleton instance # Ensure hub is running and healthy before starting the application try: @@ -59,8 +63,8 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, Any]: hub.cleanup() app = FastAPI( - title=settings.PROJECT_NAME, - version=settings.VERSION, + title=SETTINGS.PROJECT_NAME, + version=SETTINGS.VERSION, description=DESCRIPTION, lifespan=lifespan, ) @@ -68,10 +72,10 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, Any]: Instrumentator().instrument(app) # CORS middleware - if settings.BACKEND_CORS_ORIGINS: + if SETTINGS.BACKEND_CORS_ORIGINS: app.add_middleware( CORSMiddleware, - allow_origins=[str(origin) for origin in settings.BACKEND_CORS_ORIGINS], + allow_origins=[str(origin) for origin in SETTINGS.BACKEND_CORS_ORIGINS], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], @@ -94,7 +98,7 @@ async def health_check( is_healthy = await hub.check_hub_health() return HealthCheckResponse( status=HealthStatus.HEALTHY if is_healthy else HealthStatus.UNHEALTHY, - deployment_mode=settings.DEPLOYMENT_MODE, + deployment_mode=SETTINGS.DEPLOYMENT_MODE, ) # Stats endpoint @@ -120,13 +124,13 @@ async def get_hub_stats( return HubStatusResponse( hub_running=is_running, hub_healthy=is_healthy, - deployment_mode=settings.DEPLOYMENT_MODE, - max_instances=settings.selenium_grid.MAX_BROWSER_INSTANCES, + deployment_mode=SETTINGS.DEPLOYMENT_MODE, + max_instances=SETTINGS.selenium_grid.MAX_BROWSER_INSTANCES, browsers=browsers, ) # Include browser management endpoints - app.include_router(browsers_router, prefix=settings.API_V1_STR) + app.include_router(browsers_router, prefix=SETTINGS.API_V1_STR) # Include Selenium Hub proxy endpoints app.include_router(selenium_proxy_router) @@ -137,27 +141,57 @@ async def get_hub_stats( description=DESCRIPTION, describe_full_response_schema=True, describe_all_responses=True, + auth_config=AuthConfig( + dependencies=[Depends(verify_token)], + ), ) - MCP_HTTP_PATH = "/mcp" - MCP_SSE_PATH = "/sse" mcp.mount_http(mount_path=MCP_HTTP_PATH) mcp.mount_sse(mount_path=MCP_SSE_PATH) - @app.api_route("/", methods=["GET", "POST"], include_in_schema=False) - async def root_redirect(request: Request) -> Response: - accept: str = request.headers.get("accept", "").lower() - method: str = request.method.upper() - - logger.info(f"Received {method=} with Accept: {accept}") + @app.api_route("/", methods=["GET", "POST"]) + async def root_proxy( + request: Request, + credentials: HTTPAuthorizationCredentials = Depends(verify_token), + ) -> Response: + """ + FastApiMCP does not allow mounting directly on the root path `/`. + However, MCP clients (especially when using uvx) expect to connect on `/`. + This proxy handles requests on `/` and internally routes them to the proper MCP endpoints. + For SSE (Server-Sent Events) and HTTP transports, it redirects or proxies requests accordingly, + ensuring compatibility with client expectations without violating FastApiMCP mounting rules. + """ + + accept = request.headers.get("accept", "").lower() + method = request.method.upper() + session_manager = mcp._http_transport # noqa: SLF001 if "text/event-stream" in accept: - # MCP allows POST or GET here - logger.info(f"Redirecting to SSE endpoint /sse (method={method})") - return RedirectResponse(url="/sse") + if method == "GET": + return await handle_fastapi_request( + name="SSE", + request=request, + target_path=MCP_SSE_PATH, + method=method, + session_manager=session_manager, + ) + elif method == "POST": + return await handle_fastapi_request( + name="SSE messages", + request=request, + target_path=urljoin(MCP_SSE_PATH, "/messages"), + method=method, + session_manager=session_manager, + ) + else: + return JSONResponse({"detail": "Unsupported method"}, status_code=405) elif "application/json" in accept: - # JSON RPC endpoint (usually POST) - logger.info(f"Redirecting to HTTP JSON RPC endpoint /mcp (method={method})") - return RedirectResponse(url="/mcp") + return await handle_fastapi_request( + name="HTTP", + request=request, + target_path=MCP_HTTP_PATH, + method=method, + session_manager=session_manager, + ) else: logger.warning(f"Unsupported Accept header or method: method={method}, accept={accept}") return JSONResponse({"detail": "Unsupported Accept header or method"}, status_code=405) diff --git a/src/app/routers/browsers/routes.py b/src/app/routers/browsers/routes.py index 8fb4291..4744296 100644 --- a/src/app/routers/browsers/routes.py +++ b/src/app/routers/browsers/routes.py @@ -5,9 +5,9 @@ from fastapi import APIRouter, Depends, HTTPException, Request, status from fastapi.security import HTTPAuthorizationCredentials +from app.common.logger import logger from app.core.settings import Settings from app.dependencies import get_settings, verify_token -from app.logger import logger from app.services.selenium_hub import SeleniumHub from app.services.selenium_hub.models.browser import BrowserConfig, BrowserInstance diff --git a/src/app/routers/selenium_proxy.py b/src/app/routers/selenium_proxy.py index 1cded5c..ac040e9 100644 --- a/src/app/routers/selenium_proxy.py +++ b/src/app/routers/selenium_proxy.py @@ -9,9 +9,9 @@ from fastapi.responses import RedirectResponse from fastapi.security import HTTPBasicCredentials +from app.common.logger import logger from app.core.settings import Settings from app.dependencies import get_settings, verify_basic_auth -from app.logger import logger from app.services.selenium_hub import SeleniumHub # Constants