mcp-mesh 0.7.21__py3-none-any.whl → 0.8.0b1__py3-none-any.whl

_mcp_mesh/pipeline/mcp_heartbeat/heartbeat_orchestrator.py

@@ -1,311 +0,0 @@
-"""
-Heartbeat orchestrator for managing periodic pipeline execution.
-
-Provides a high-level interface for executing heartbeat pipelines with proper
-context management and error handling.
-"""
-
-import json
-import logging
-from datetime import UTC, datetime
-from typing import Any, Optional
-
-from ...shared.support_types import HealthStatus, HealthStatusType
-from .heartbeat_pipeline import HeartbeatPipeline
-
-logger = logging.getLogger(__name__)
-
-
-class HeartbeatOrchestrator:
-    """
-    Orchestrates heartbeat pipeline execution for periodic registry communication.
-
-    Manages the context preparation, pipeline execution, and result processing
-    for the 30-second heartbeat cycle.
-    """
-
-    def __init__(self):
-        self.logger = logging.getLogger(f"{__name__}.HeartbeatOrchestrator")
-        self.pipeline = HeartbeatPipeline()
-        self._heartbeat_count = 0
-
-    async def execute_heartbeat(self, agent_id: str, context: dict[str, Any]) -> bool:
-        """
-        Execute a complete heartbeat cycle with comprehensive error handling.
-
-        Args:
-            agent_id: Agent identifier
-            context: Full pipeline context from startup
-
-        Returns:
-            bool: True if heartbeat succeeded, False if failed
-        """
-        self._heartbeat_count += 1
-
-        try:
-
-            # Prepare heartbeat context with validation
-            heartbeat_context = await self._prepare_heartbeat_context(agent_id, context)
-
-            # Validate required context before proceeding
-            if not self._validate_heartbeat_context(heartbeat_context):
-                self.logger.error(
-                    f"❌ Heartbeat #{self._heartbeat_count} failed: invalid context"
-                )
-                return False
-
-            # Check if health status is unhealthy - skip heartbeat if so
-            health_status = heartbeat_context.get("health_status")
-            if health_status and health_status.status == HealthStatusType.UNHEALTHY:
-                self.logger.warning(
-                    f"⚠️ Heartbeat #{self._heartbeat_count} skipped for agent '{agent_id}': Health status is UNHEALTHY"
-                )
-                self.logger.warning(f"   Health checks failed: {health_status.checks}")
-                self.logger.warning(f"   Errors: {health_status.errors}")
-                return False
-
-            # Log heartbeat request details for debugging
-            self._log_heartbeat_request(heartbeat_context, self._heartbeat_count)
-
-            # Execute heartbeat pipeline with timeout protection
-            self.logger.trace(
-                f"💓 Executing heartbeat #{self._heartbeat_count} for agent '{agent_id}'"
-            )
-
-            # Add timeout to prevent hanging heartbeats (30 seconds max)
-            import asyncio
-
-            try:
-
-                result = await asyncio.wait_for(
-                    self.pipeline.execute_heartbeat_cycle(heartbeat_context),
-                    timeout=30.0,
-                )
-
-            except TimeoutError:
-                self.logger.error(
-                    f"❌ Heartbeat #{self._heartbeat_count} timed out after 30 seconds"
-                )
-                return False
-
-            # Process results
-            success = self._process_heartbeat_result(
-                result, agent_id, self._heartbeat_count
-            )
-
-            # Log periodic status updates
-            if self._heartbeat_count % 10 == 0:
-                elapsed_time = self._heartbeat_count * 30  # Assuming 30s interval
-                self.logger.info(
-                    f"💓 Heartbeat #{self._heartbeat_count} for agent '{agent_id}' - "
-                    f"running for {elapsed_time} seconds"
-                )
-
-            return success
-
-        except Exception as e:
-            # Log detailed error information for debugging
-            import traceback
-
-            self.logger.error(
-                f"❌ Heartbeat #{self._heartbeat_count} failed for agent '{agent_id}': {e}\n"
-                f"Traceback: {traceback.format_exc()}"
-            )
-            return False
-
-    async def _prepare_heartbeat_context(
-        self, agent_id: str, startup_context: dict[str, Any]
-    ) -> dict[str, Any]:
-        """Prepare context for heartbeat pipeline execution."""
-
-        # Build health status from startup context
-        health_status = await self._build_health_status_from_context(
-            startup_context, agent_id
-        )
-
-        # Prepare heartbeat-specific context - registry_wrapper will be created in pipeline
-        heartbeat_context = {
-            "agent_id": agent_id,
-            "health_status": health_status,
-            "agent_config": startup_context.get("agent_config", {}),
-            "registration_data": startup_context.get("registration_data", {}),
-            # Include other relevant context from startup
-            "fastmcp_servers": startup_context.get("fastmcp_servers", {}),
-            "capabilities": startup_context.get("capabilities", []),
-        }
-
-        return heartbeat_context
-
-    def _validate_heartbeat_context(self, heartbeat_context: dict[str, Any]) -> bool:
-        """Validate that heartbeat context has all required components."""
-
-        required_fields = ["agent_id", "health_status"]
-
-        for field in required_fields:
-            if field not in heartbeat_context or heartbeat_context[field] is None:
-                self.logger.error(
-                    f"❌ Heartbeat context validation failed: missing '{field}'"
-                )
-                return False
-
-        # Additional validation for health_status
-        health_status = heartbeat_context.get("health_status")
-        if not hasattr(health_status, "agent_name") or not hasattr(
-            health_status, "status"
-        ):
-            self.logger.error(
-                "❌ Heartbeat context validation failed: invalid health_status object"
-            )
-            return False
-
-        return True
-
-    async def _build_health_status_from_context(
-        self, startup_context: dict[str, Any], agent_id: str
-    ) -> HealthStatus:
-        """Build health status object from startup context with optional user health check."""
-
-        agent_config = startup_context.get("agent_config", {})
-
-        # Check if user provided a health_check function
-        health_check_fn = agent_config.get("health_check")
-        health_check_ttl = agent_config.get("health_check_ttl", 15)
-
-        # If health check is configured, use the cache
-        if health_check_fn:
-            from ...shared.health_check_manager import get_health_status_with_cache
-
-            return await get_health_status_with_cache(
-                agent_id=agent_id,
-                health_check_fn=health_check_fn,
-                agent_config=agent_config,
-                startup_context=startup_context,
-                ttl=health_check_ttl,
-            )
-
-        # No health check configured - use existing logic
-        existing_health_status = startup_context.get("health_status")
-
-        if existing_health_status:
-            # Update timestamp to current time for fresh heartbeat
-            if hasattr(existing_health_status, "timestamp"):
-                existing_health_status.timestamp = datetime.now(UTC)
-            return existing_health_status
-
-        # Build minimal health status from context if none exists
-        return HealthStatus(
-            agent_name=agent_id,
-            status=HealthStatusType.HEALTHY,
-            capabilities=agent_config.get("capabilities", []),
-            timestamp=datetime.now(UTC),
-            version=agent_config.get("version", "1.0.0"),
-            metadata=agent_config,
-        )
-
-    def _log_heartbeat_request(
-        self, heartbeat_context: dict[str, Any], heartbeat_count: int
-    ) -> None:
-        """Log heartbeat request details for debugging."""
-
-        health_status = heartbeat_context.get("health_status")
-        if not health_status:
-            return
-
-        # Convert health status to dict for logging
-        if hasattr(health_status, "__dict__"):
-            health_dict = {
-                "agent_name": getattr(health_status, "agent_name", "unknown"),
-                "status": (
-                    getattr(health_status, "status", "healthy").value
-                    if hasattr(getattr(health_status, "status", "healthy"), "value")
-                    else str(getattr(health_status, "status", "healthy"))
-                ),
-                "capabilities": getattr(health_status, "capabilities", []),
-                "timestamp": (
-                    getattr(health_status, "timestamp", "").isoformat()
-                    if hasattr(getattr(health_status, "timestamp", ""), "isoformat")
-                    else str(getattr(health_status, "timestamp", ""))
-                ),
-                "version": getattr(health_status, "version", "1.0.0"),
-                "metadata": getattr(health_status, "metadata", {}),
-            }
-        else:
-            health_dict = health_status
-
-        request_json = json.dumps(health_dict, indent=2, default=str)
-
-    def _process_heartbeat_result(
-        self, result: Any, agent_id: str, heartbeat_count: int
-    ) -> bool:
-        """Process heartbeat pipeline result and log appropriately."""
-
-        if result.is_success():
-            # Check fast heartbeat status to understand what happened
-            fast_heartbeat_status = result.context.get("fast_heartbeat_status")
-            heartbeat_response = result.context.get("heartbeat_response")
-
-            # Import FastHeartbeatStatus for comparison
-            from ...shared.fast_heartbeat_status import (
-                FastHeartbeatStatus,
-                FastHeartbeatStatusUtil,
-            )
-
-            if not fast_heartbeat_status:
-                self.logger.error(
-                    f"💔 Heartbeat #{heartbeat_count} failed for agent '{agent_id}' - missing fast_heartbeat_status"
-                )
-                return False
-
-            if FastHeartbeatStatusUtil.should_skip_for_optimization(
-                fast_heartbeat_status
-            ):
-                # Fast heartbeat optimization - no changes detected
-                self.logger.debug(
-                    f"🚀 Heartbeat #{heartbeat_count} optimized for agent '{agent_id}' - no changes detected"
-                )
-                return True
-            elif FastHeartbeatStatusUtil.should_skip_for_resilience(
-                fast_heartbeat_status
-            ):
-                # Fast heartbeat resilience - registry/network error
-                self.logger.info(
-                    f"⚠️ Heartbeat #{heartbeat_count} skipped for agent '{agent_id}' - resilience mode ({fast_heartbeat_status.value})"
-                )
-                return True
-            elif FastHeartbeatStatusUtil.requires_full_heartbeat(fast_heartbeat_status):
-                # Full heartbeat was executed - check for response
-                if heartbeat_response:
-                    # Log response details for debugging
-                    response_json = json.dumps(
-                        heartbeat_response, indent=2, default=str
-                    )
-                    self.logger.debug(
-                        f"🔍 Heartbeat response #{heartbeat_count}:\n{response_json}"
-                    )
-
-                    # Log dependency resolution info if available
-                    deps_resolved = heartbeat_response.get("dependencies_resolved", {})
-                    if deps_resolved:
-                        self.logger.info(
-                            f"🔗 Dependencies resolved: {len(deps_resolved)} items"
-                        )
-
-                    self.logger.info(
-                        f"💚 Heartbeat #{heartbeat_count} sent successfully for agent '{agent_id}' (full refresh: {fast_heartbeat_status.value})"
-                    )
-                    return True
-                else:
-                    self.logger.warning(
-                        f"💔 Heartbeat #{heartbeat_count} failed for agent '{agent_id}' - full heartbeat expected but no response"
-                    )
-                    return False
-            else:
-                self.logger.warning(
-                    f"💔 Heartbeat #{heartbeat_count} unknown status for agent '{agent_id}': {fast_heartbeat_status}"
-                )
-                return False
-        else:
-            self.logger.warning(
-                f"💔 Heartbeat #{heartbeat_count} pipeline failed for agent '{agent_id}': {result.message}"
-            )
-            return False

_mcp_mesh/pipeline/mcp_heartbeat/heartbeat_pipeline.py

@@ -1,282 +0,0 @@
-"""
-Heartbeat pipeline for MCP Mesh periodic operations.
-
-Provides structured execution of heartbeat operations with proper error handling
-and logging. Runs every 30 seconds to maintain registry communication and
-dependency resolution.
-"""
-
-import logging
-from typing import Any
-
-from ...shared.fast_heartbeat_status import FastHeartbeatStatus, FastHeartbeatStatusUtil
-from ..shared import PipelineResult, PipelineStatus
-from ..shared.mesh_pipeline import MeshPipeline
-from .dependency_resolution import DependencyResolutionStep
-from .fast_heartbeat_check import FastHeartbeatStep
-from .heartbeat_send import HeartbeatSendStep
-from .llm_tools_resolution import LLMToolsResolutionStep
-from .registry_connection import RegistryConnectionStep
-
-logger = logging.getLogger(__name__)
-
-
-class HeartbeatPipeline(MeshPipeline):
-    """
-    Specialized pipeline for heartbeat operations with fast optimization.
-
-    Executes the five core heartbeat steps in sequence:
-    1. Registry connection preparation
-    2. Fast heartbeat check (HEAD request)
-    3. Heartbeat sending (conditional POST request)
-    4. Dependency resolution (conditional)
-    5. LLM tools resolution (conditional)
-
-    Steps 3, 4, and 5 only run if fast heartbeat indicates changes are needed.
-    Provides optimization for NO_CHANGES and resilience for error conditions.
-    """
-
-    def __init__(self):
-        super().__init__(name="heartbeat-pipeline")
-        self._setup_heartbeat_steps()
-
-    def _setup_heartbeat_steps(self) -> None:
-        """Setup the heartbeat pipeline steps."""
-        steps = [
-            RegistryConnectionStep(),
-            FastHeartbeatStep(),
-            HeartbeatSendStep(required=True),
-            DependencyResolutionStep(),
-            LLMToolsResolutionStep(),
-        ]
-
-        self.add_steps(steps)
-        self.logger.trace(f"Heartbeat pipeline configured with {len(steps)} steps")
-
-    async def execute_heartbeat_cycle(
-        self, heartbeat_context: dict[str, Any]
-    ) -> PipelineResult:
-        """
-        Execute a complete heartbeat cycle with enhanced error handling.
-
-        Args:
-            heartbeat_context: Context containing registry_wrapper, agent_id, health_status, etc.
-
-        Returns:
-            PipelineResult with execution status and any context updates
-        """
-        self.logger.trace("Starting heartbeat pipeline execution")
-
-        # Initialize pipeline context with heartbeat-specific data
-        self.context.clear()
-        self.context.update(heartbeat_context)
-
-        try:
-            # Execute the pipeline with conditional logic for fast optimization
-            result = await self._execute_with_conditional_logic()
-
-            if result.is_success():
-                self.logger.trace("✅ Heartbeat pipeline completed successfully")
-            elif result.status == PipelineStatus.PARTIAL:
-                self.logger.warning(
-                    f"⚠️ Heartbeat pipeline completed partially: {result.message}"
-                )
-                # Log which steps failed
-                if result.errors:
-                    for error in result.errors:
-                        self.logger.warning(f"  - Step error: {error}")
-            else:
-                self.logger.error(f"❌ Heartbeat pipeline failed: {result.message}")
-                # Log detailed error information
-                if result.errors:
-                    for error in result.errors:
-                        self.logger.error(f"  - Pipeline error: {error}")
-
-            return result
-
-        except Exception as e:
-            # Log detailed error information for debugging
-            import traceback
-
-            self.logger.error(
-                f"❌ Heartbeat pipeline failed with exception: {e}\n"
-                f"Context keys: {list(self.context.keys())}\n"
-                f"Traceback: {traceback.format_exc()}"
-            )
-
-            # Create failure result with detailed context
-            failure_result = PipelineResult(
-                status=PipelineStatus.FAILED,
-                message=f"Heartbeat pipeline exception: {str(e)[:200]}...",  # Truncate long error messages
-                context=self.context,
-            )
-            failure_result.add_error(str(e))
-
-            return failure_result
-
-    async def _execute_with_conditional_logic(self) -> PipelineResult:
-        """
-        Execute pipeline with conditional logic based on fast heartbeat status.
-
-        Always executes:
-        - RegistryConnectionStep
-        - FastHeartbeatStep
-
-        Conditionally executes based on fast heartbeat status:
-        - NO_CHANGES: Skip remaining steps (optimization)
-        - TOPOLOGY_CHANGED, AGENT_UNKNOWN: Execute all remaining steps
-        - REGISTRY_ERROR, NETWORK_ERROR: Skip remaining steps (resilience)
-
-        Returns:
-            PipelineResult with execution status and context
-        """
-        overall_result = PipelineResult(
-            message="Heartbeat pipeline execution completed"
-        )
-
-        # Track which steps were executed for logging
-        executed_steps = []
-        skipped_steps = []
-
-        try:
-            # Always execute registry connection and fast heartbeat steps
-            mandatory_steps = self.steps[
-                :2
-            ]  # RegistryConnectionStep, FastHeartbeatStep
-            conditional_steps = self.steps[
-                2:
-            ]  # HeartbeatSendStep, DependencyResolutionStep, LLMToolsResolutionStep
-
-            # Execute mandatory steps
-            for step in mandatory_steps:
-                self.logger.trace(f"Executing mandatory step: {step.name}")
-
-                step_result = await step.execute(self.context)
-                executed_steps.append(step.name)
-
-                # Merge step context into pipeline context
-                self.context.update(step_result.context)
-
-                # If step fails, handle accordingly
-                if not step_result.is_success():
-                    overall_result.status = PipelineStatus.FAILED
-                    overall_result.message = (
-                        f"Mandatory step '{step.name}' failed: {step_result.message}"
-                    )
-                    overall_result.add_error(
-                        f"Step '{step.name}': {step_result.message}"
-                    )
-
-                    if step.required:
-                        # Stop execution if required step fails
-                        for key, value in self.context.items():
-                            overall_result.add_context(key, value)
-                        return overall_result
-
-            # Check fast heartbeat status for conditional execution
-            fast_heartbeat_status = self.context.get("fast_heartbeat_status")
-
-            if fast_heartbeat_status is None:
-                # Fast heartbeat step failed to set status - fallback to full execution
-                self.logger.warning(
-                    "⚠️ Fast heartbeat status not found - falling back to full execution"
-                )
-                should_execute_remaining = True
-                reason = "fallback (missing status)"
-            elif FastHeartbeatStatusUtil.should_skip_for_optimization(
-                fast_heartbeat_status
-            ):
-                # NO_CHANGES - skip for optimization
-                should_execute_remaining = False
-                reason = "optimization (no changes detected)"
-                self.logger.trace(
-                    f"🚀 Skipping remaining steps for optimization: {reason}"
-                )
-            elif FastHeartbeatStatusUtil.should_skip_for_resilience(
-                fast_heartbeat_status
-            ):
-                # REGISTRY_ERROR, NETWORK_ERROR - skip for resilience
-                should_execute_remaining = False
-                reason = "resilience (preserve existing state)"
-                self.logger.warning(
-                    f"⚠️ Skipping remaining steps for resilience: {reason}"
-                )
-            elif FastHeartbeatStatusUtil.requires_full_heartbeat(fast_heartbeat_status):
-                # TOPOLOGY_CHANGED, AGENT_UNKNOWN - execute full pipeline
-                should_execute_remaining = True
-                reason = "changes detected or re-registration needed"
-                self.logger.info(f"🔄 Executing remaining steps: {reason}")
-            else:
-                # Unknown status - fallback to full execution
-                self.logger.warning(
-                    f"⚠️ Unknown fast heartbeat status '{fast_heartbeat_status}' - falling back to full execution"
-                )
-                should_execute_remaining = True
-                reason = "fallback (unknown status)"
-
-            # Execute or skip conditional steps based on decision
-            if should_execute_remaining:
-                for step in conditional_steps:
-                    self.logger.trace(f"Executing conditional step: {step.name}")
-
-                    step_result = await step.execute(self.context)
-                    executed_steps.append(step.name)
-
-                    # Merge step context into pipeline context
-                    self.context.update(step_result.context)
-
-                    # Handle step failure
-                    if not step_result.is_success():
-                        if step.required:
-                            overall_result.status = PipelineStatus.FAILED
-                            overall_result.message = f"Required step '{step.name}' failed: {step_result.message}"
-                            overall_result.add_error(
-                                f"Step '{step.name}': {step_result.message}"
-                            )
-                            break
-                        else:
-                            # Optional step failed - mark as partial success
-                            if overall_result.status == PipelineStatus.SUCCESS:
-                                overall_result.status = PipelineStatus.PARTIAL
-                            overall_result.add_error(
-                                f"Optional step '{step.name}': {step_result.message}"
-                            )
-                            self.logger.warning(
-                                f"⚠️ Optional step '{step.name}' failed: {step_result.message}"
-                            )
-            else:
-                # Mark skipped steps
-                for step in conditional_steps:
-                    skipped_steps.append(step.name)
-
-            # Set final result message
-            if executed_steps and skipped_steps:
-                overall_result.message = (
-                    f"Pipeline completed with conditional execution - "
-                    f"executed: {executed_steps}, skipped: {skipped_steps} ({reason})"
-                )
-            elif executed_steps:
-                overall_result.message = (
-                    f"Pipeline completed - executed: {executed_steps} ({reason})"
-                )
-            else:
-                overall_result.message = (
-                    f"Pipeline completed - all steps skipped ({reason})"
-                )
-
-            # Add final context
-            for key, value in self.context.items():
-                overall_result.add_context(key, value)
-
-            return overall_result
-
-        except Exception as e:
-            # Handle unexpected exceptions
-            overall_result.status = PipelineStatus.FAILED
-            overall_result.message = f"Pipeline execution failed with exception: {e}"
-            overall_result.add_error(str(e))
-            for key, value in self.context.items():
-                overall_result.add_context(key, value)
-
-            self.logger.error(f"❌ Conditional pipeline execution failed: {e}")
-            return overall_result

_mcp_mesh/pipeline/mcp_heartbeat/heartbeat_send.py

@@ -1,98 +0,0 @@
-"""
-Heartbeat sending step for MCP Mesh pipeline.
-
-Handles sending heartbeat to the mesh registry service.
-"""
-
-import logging
-from typing import Any
-
-from ..shared import PipelineResult, PipelineStatus, PipelineStep
-
-logger = logging.getLogger(__name__)
-
-
-class HeartbeatSendStep(PipelineStep):
-    """
-    Sends heartbeat to the mesh registry.
-
-    Performs the actual registry communication using the prepared
-    heartbeat data from previous steps.
-    """
-
-    def __init__(self, required: bool = True):
-        super().__init__(
-            name="heartbeat-send",
-            required=required,
-            description="Send heartbeat to mesh registry",
-        )
-
-    async def execute(self, context: dict[str, Any]) -> PipelineResult:
-        """Send heartbeat to registry or print JSON in debug mode."""
-        result = PipelineResult(message="Heartbeat processed successfully")
-
-        try:
-            # Get required context
-            health_status = context.get("health_status")
-            agent_id = context.get("agent_id", "unknown-agent")
-            registration_data = context.get("registration_data")
-
-            if not health_status:
-                raise ValueError("Health status not available in context")
-
-            # Prepare heartbeat for registry
-            self.logger.trace(f"🔍 Preparing heartbeat for agent '{agent_id}'")
-
-            # Send actual HTTP request to registry
-            registry_wrapper = context.get("registry_wrapper")
-
-            if not registry_wrapper:
-                # If no registry wrapper, just log the payload and mark as successful
-                self.logger.info(
-                    f"⚠️ No registry connection - would send heartbeat for agent '{agent_id}'"
-                )
-                result.add_context(
-                    "heartbeat_response", {"status": "no_registry", "logged": True}
-                )
-                result.add_context("dependencies_resolved", {})
-                result.message = (
-                    f"Heartbeat logged for agent '{agent_id}' (no registry)"
-                )
-                return result
-
-            self.logger.info(f"💓 Sending heartbeat for agent '{agent_id}'...")
-
-            response = await registry_wrapper.send_heartbeat_with_dependency_resolution(
-                health_status
-            )
-
-            if response:
-                # Store response data
-                result.add_context("heartbeat_response", response)
-                result.add_context(
-                    "dependencies_resolved",
-                    response.get("dependencies_resolved", {}),
-                )
-
-                result.message = f"Heartbeat sent successfully for agent '{agent_id}'"
-                self.logger.info(f"💚 Heartbeat successful for agent '{agent_id}'")
-
-                # Log dependency resolution info
-                deps_resolved = response.get("dependencies_resolved", {})
-                if deps_resolved:
-                    self.logger.info(
-                        f"🔗 Dependencies resolved: {len(deps_resolved)} items"
-                    )
-
-            else:
-                result.status = PipelineStatus.FAILED
-                result.message = "Heartbeat failed - no response from registry"
-                self.logger.error("💔 Heartbeat failed - no response")
-
-        except Exception as e:
-            result.status = PipelineStatus.FAILED
-            result.message = f"Heartbeat processing failed: {e}"
-            result.add_error(str(e))
-            self.logger.error(f"❌ Heartbeat processing failed: {e}")
-
-        return result