mcp-mesh 0.4.1__py3-none-any.whl → 0.5.0__py3-none-any.whl

_mcp_mesh/pipeline/api_heartbeat/api_heartbeat_pipeline.py

@@ -0,0 +1,309 @@
+"""
+API heartbeat pipeline for FastAPI service health monitoring.
+
+Provides structured execution of API service heartbeat operations with proper 
+error handling and logging. Runs periodically to maintain registry communication 
+and service health status for FastAPI applications using @mesh.route decorators.
+"""
+
+import logging
+from typing import Any
+
+from ...shared.fast_heartbeat_status import FastHeartbeatStatus, FastHeartbeatStatusUtil
+from ..shared.mesh_pipeline import MeshPipeline
+from ..shared.pipeline_types import PipelineStatus
+from .api_registry_connection import APIRegistryConnectionStep
+from .api_health_check import APIHealthCheckStep  
+from .api_fast_heartbeat_check import APIFastHeartbeatStep
+from .api_heartbeat_send import APIHeartbeatSendStep
+from .api_dependency_resolution import APIDependencyResolutionStep
+
+logger = logging.getLogger(__name__)
+
+
+class APIHeartbeatPipeline(MeshPipeline):
+    """
+    Specialized pipeline for API service heartbeat operations with fast optimization.
+
+    Executes the core API heartbeat steps in sequence:
+    1. Registry connection preparation
+    2. API health check (validate FastAPI app status)
+    3. Fast heartbeat check (HEAD request)
+    4. Heartbeat sending (conditional POST request)
+    5. Dependency resolution (conditional)
+
+    Steps 4 and 5 only run if fast heartbeat indicates changes are needed.
+    Provides optimization for NO_CHANGES and resilience for error conditions.
+
+    API services now support:
+    - Service availability and health status monitoring
+    - Efficient HEAD->conditional POST pattern (like MCP agents)
+    - Dynamic dependency resolution and injection
+    - Route handler dependency updates from registry responses
+    """
+
+    def __init__(self, name: str = "api-heartbeat-pipeline"):
+        super().__init__(name=name)
+        self._setup_api_heartbeat_steps()
+
+    def _setup_api_heartbeat_steps(self) -> None:
+        """Setup the API heartbeat pipeline steps with fast optimization."""
+        # API heartbeat steps with fast optimization pattern
+        steps = [
+            APIRegistryConnectionStep(),    # Prepare registry communication
+            APIHealthCheckStep(),           # Check FastAPI app health status
+            APIFastHeartbeatStep(),         # Fast heartbeat check (HEAD request)
+            APIHeartbeatSendStep(),         # Conditional heartbeat send (POST request)
+            APIDependencyResolutionStep(),  # Conditional dependency resolution
+        ]
+
+        self.add_steps(steps)
+        self.logger.debug(f"API heartbeat pipeline configured with {len(steps)} steps")
+        
+        # Log the pipeline strategy
+        self.logger.info(
+            f"🌐 API Heartbeat Pipeline initialized: fast optimization for FastAPI apps"
+        )
+        self.logger.debug(
+            f"📋 Pipeline steps: {[step.name for step in steps]}"
+        )
+
+    async def execute_api_heartbeat_cycle(
+        self, heartbeat_context: dict[str, Any]
+    ) -> Any:
+        """
+        Execute a complete API heartbeat cycle with fast optimization and enhanced error handling.
+
+        Args:
+            heartbeat_context: Context containing registry_wrapper, service_id, 
+                              health_status, fastapi_app, etc.
+
+        Returns:
+            PipelineResult with execution status and any context updates
+        """
+        self.logger.debug("Starting API 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.debug("✅ API heartbeat pipeline completed successfully")
+            elif result.status == PipelineStatus.PARTIAL:
+                self.logger.warning(
+                    f"⚠️ API 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"❌ API 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"❌ API 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
+            from ..shared.pipeline_types import PipelineResult
+            
+            failure_result = PipelineResult(
+                status=PipelineStatus.FAILED,
+                message=f"API heartbeat pipeline exception: {str(e)[:200]}...",
+                context=self.context,
+            )
+            failure_result.add_error(str(e))
+
+            return failure_result
+
+    async def _execute_with_conditional_logic(self) -> "PipelineResult":
+        """
+        Execute API pipeline with conditional logic based on fast heartbeat status.
+
+        Always executes:
+        - APIRegistryConnectionStep
+        - APIHealthCheckStep  
+        - APIFastHeartbeatStep
+
+        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
+        """
+        from ..shared.pipeline_types import PipelineResult
+        
+        overall_result = PipelineResult(
+            message="API heartbeat pipeline execution completed"
+        )
+
+        # Track which steps were executed for logging
+        executed_steps = []
+        skipped_steps = []
+
+        try:
+            # Always execute registry connection, health check, and fast heartbeat steps
+            mandatory_steps = self.steps[
+                :3
+            ]  # APIRegistryConnectionStep, APIHealthCheckStep, APIFastHeartbeatStep
+            conditional_steps = self.steps[
+                3:
+            ]  # APIHeartbeatSendStep, APIDependencyResolutionStep
+
+            # Execute mandatory steps
+            for step in mandatory_steps:
+                self.logger.debug(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(
+                    "⚠️ API 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.info(
+                    f"🚀 API heartbeat: 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"⚠️ API heartbeat: 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"🔄 API heartbeat: Executing remaining steps: {reason}")
+            else:
+                # Unknown status - fallback to full execution
+                self.logger.warning(
+                    f"⚠️ Unknown API 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.debug(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)
+                
+                # For skipped heartbeat due to NO_CHANGES, set success context
+                if fast_heartbeat_status == FastHeartbeatStatus.NO_CHANGES:
+                    overall_result.add_context("heartbeat_success", True)
+                    overall_result.add_context("heartbeat_skipped", True)
+                    overall_result.add_context("skip_reason", "no_changes_optimization")
+
+            # Set final result message
+            if executed_steps and skipped_steps:
+                overall_result.message = (
+                    f"API pipeline completed with conditional execution - "
+                    f"executed: {executed_steps}, skipped: {skipped_steps} ({reason})"
+                )
+            elif executed_steps:
+                overall_result.message = (
+                    f"API pipeline completed - executed: {executed_steps} ({reason})"
+                )
+            else:
+                overall_result.message = (
+                    f"API 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"API 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"❌ API conditional pipeline execution failed: {e}")
+            return overall_result

_mcp_mesh/pipeline/api_heartbeat/api_heartbeat_send.py

@@ -0,0 +1,332 @@
+"""
+API heartbeat send step for API heartbeat pipeline.
+
+Sends service health status and registration data to the registry
+for FastAPI applications using @mesh.route decorators.
+"""
+
+import logging
+from typing import Any
+
+from ..shared.base_step import PipelineStep
+from ..shared.pipeline_types import PipelineResult
+
+logger = logging.getLogger(__name__)
+
+
+class APIHeartbeatSendStep(PipelineStep):
+    """
+    Send API service heartbeat to registry.
+    
+    Communicates service health status and registration information
+    to the registry for monitoring and discovery purposes.
+    """
+
+    def __init__(self, required: bool = True):
+        super().__init__(
+            name="api-heartbeat-send",
+            required=required,
+        )
+
+    async def execute(self, context: dict[str, Any]) -> PipelineResult:
+        """
+        Send API service heartbeat to registry.
+
+        Args:
+            context: Pipeline context containing registry_wrapper, health_status, service_id
+
+        Returns:
+            PipelineResult with heartbeat_response in context
+        """
+        self.logger.debug("Sending API service heartbeat to registry")
+
+        try:
+            # Get required components from context
+            registry_wrapper = context.get("registry_wrapper")
+            health_status = context.get("health_status")
+            service_id = context.get("service_id") or context.get("agent_id", "unknown")
+
+            if not registry_wrapper:
+                error_msg = "No registry wrapper available for heartbeat"
+                self.logger.error(f"❌ {error_msg}")
+                
+                from ..shared.pipeline_types import PipelineStatus
+                result = PipelineResult(
+                    status=PipelineStatus.FAILED,
+                    message=error_msg,
+                    context=context
+                )
+                result.add_error(error_msg)
+                return result
+
+            if not health_status:
+                error_msg = "No health status available for heartbeat"
+                self.logger.error(f"❌ {error_msg}")
+                
+                from ..shared.pipeline_types import PipelineStatus
+                result = PipelineResult(
+                    status=PipelineStatus.FAILED,
+                    message=error_msg,
+                    context=context
+                )
+                result.add_error(error_msg)
+                return result
+
+            # Prepare heartbeat data for API service
+            heartbeat_data = self._prepare_api_heartbeat_data(
+                health_status, service_id, context
+            )
+
+            self.logger.debug(f"📡 Sending heartbeat for API service '{service_id}'")
+
+            # Send heartbeat to registry using the same format as test_api_service.json
+            # Import json at the beginning
+            import aiohttp
+            import json
+            
+            try:
+                # For API services, send directly to registry using the format that works
+                # Get registry URL 
+                registry_url = context.get("registry_url", "http://localhost:8000")
+                
+                # Build the API service payload using actual dependencies from @mesh.route decorators
+                display_config = context.get("display_config", {})
+                
+                # Extract all dependencies from registered @mesh.route decorators
+                all_route_dependencies = self._extract_all_route_dependencies(context)
+                
+                api_service_payload = {
+                    "agent_id": service_id,
+                    "agent_type": "api",
+                    "tools": [
+                        {
+                            "function_name": "api_endpoint_handler",
+                            "dependencies": all_route_dependencies
+                        }
+                    ],
+                    "http_host": display_config.get("display_host", "127.0.0.1"),
+                    "http_port": display_config.get("display_port", 8080)
+                }
+                
+                self.logger.debug(f"📡 Sending API service payload to {registry_url}/heartbeat")
+                self.logger.debug(f"🔍 POST payload: {json.dumps(api_service_payload, indent=2)}")
+                
+                try:
+                    async with aiohttp.ClientSession() as session:
+                        async with session.post(
+                            f"{registry_url}/heartbeat",
+                            headers={"Content-Type": "application/json"},
+                            data=json.dumps(api_service_payload)
+                        ) as response:
+                            self.logger.debug(f"{registry_url} \"POST /heartbeat HTTP/1.1\" {response.status}")
+                            if response.status == 200:
+                                heartbeat_response = await response.json()
+                            else:
+                                response_text = await response.text()
+                                self.logger.error(f"❌ Registry error {response.status}: {response_text}")
+                                raise Exception(f"Registry returned {response.status}: {response_text}")
+                
+                except Exception as http_error:
+                    self.logger.error(f"❌ HTTP request failed: {http_error}")
+                    raise http_error
+                
+                if heartbeat_response:
+                    self.logger.info(f"💚 API heartbeat successful for service '{service_id}'")
+                    
+                    return PipelineResult(
+                        message=f"API heartbeat sent for service {service_id}",
+                        context={
+                            "heartbeat_response": heartbeat_response,
+                            "heartbeat_success": True,
+                            "heartbeat_data": heartbeat_data,
+                        }
+                    )
+                else:
+                    error_msg = f"Registry heartbeat failed for service {service_id}"
+                    self.logger.warning(f"⚠️ {error_msg}")
+                    
+                    from ..shared.pipeline_types import PipelineStatus
+                    result = PipelineResult(
+                        status=PipelineStatus.FAILED,
+                        message=error_msg,
+                        context=context
+                    )
+                    result.add_error(error_msg)
+                    return result
+
+            except Exception as e:
+                error_msg = f"Registry communication failed: {e}"
+                self.logger.error(f"❌ {error_msg}")
+                
+                from ..shared.pipeline_types import PipelineStatus
+                result = PipelineResult(
+                    status=PipelineStatus.FAILED,
+                    message=error_msg,
+                    context=context
+                )
+                result.add_error(str(e))
+                return result
+
+        except Exception as e:
+            error_msg = f"API heartbeat send failed: {e}"
+            self.logger.error(f"❌ {error_msg}")
+
+            from ..shared.pipeline_types import PipelineStatus
+            result = PipelineResult(
+                status=PipelineStatus.FAILED,
+                message=error_msg,
+                context=context
+            )
+            result.add_error(str(e))
+            return result
+
+    def _prepare_api_heartbeat_data(
+        self, health_status: Any, service_id: str, context: dict[str, Any]
+    ) -> dict[str, Any]:
+        """Prepare heartbeat data specific to API services."""
+        try:
+            # Extract FastAPI-specific information
+            app_title = context.get("app_title", "Unknown API")
+            app_version = context.get("app_version", "1.0.0")
+            routes_total = context.get("routes_total", 0)
+            routes_with_mesh = context.get("routes_with_mesh", 0)
+            
+            # Get display configuration
+            display_config = context.get("display_config", {})
+            host = display_config.get("host", "0.0.0.0")
+            port = display_config.get("port", 8080)
+
+            heartbeat_data = {
+                "service_id": service_id,
+                "service_type": "api",
+                "app_title": app_title,
+                "app_version": app_version,
+                "host": host,
+                "port": port,
+                "routes": {
+                    "total": routes_total,
+                    "with_mesh": routes_with_mesh,
+                },
+                "health_status": health_status if isinstance(health_status, dict) else {
+                    "status": health_status.status.value if hasattr(health_status, "status") and hasattr(health_status.status, "value") else str(getattr(health_status, "status", "healthy")),
+                    "timestamp": health_status.timestamp.isoformat() if hasattr(health_status, "timestamp") and hasattr(health_status.timestamp, "isoformat") else str(getattr(health_status, "timestamp", "")),
+                    "version": getattr(health_status, "version", "1.0.0"),
+                    "metadata": getattr(health_status, "metadata", {}),
+                }
+            }
+
+            return heartbeat_data
+
+        except Exception as e:
+            self.logger.warning(f"⚠️ Could not prepare heartbeat data: {e}")
+            return {
+                "service_id": service_id,
+                "service_type": "api",
+                "error": f"Failed to prepare heartbeat data: {e}"
+            }
+
+    def _extract_all_route_dependencies(self, context: dict[str, Any]) -> list[dict[str, Any]]:
+        """
+        Extract all unique dependencies from @mesh.route decorators in the FastAPI app.
+        
+        This method looks at the actual route dependencies that were discovered during
+        the API startup pipeline and extracts them for registry registration.
+        
+        Args:
+            context: Pipeline context containing FastAPI app and route information
+            
+        Returns:
+            List of unique dependency objects in the format expected by registry
+        """
+        try:
+            # Try to get dependencies from startup context (preferred method)
+            api_service_metadata = context.get("api_service_metadata", {})
+            route_capabilities = api_service_metadata.get("capabilities", [])
+            
+            self.logger.debug(f"🔍 api_service_metadata keys: {list(api_service_metadata.keys())}")
+            self.logger.debug(f"🔍 route_capabilities count: {len(route_capabilities)}")
+            self.logger.debug(f"🔍 route_capabilities: {route_capabilities}")
+            
+            # Extract dependencies from route capabilities
+            all_dependencies = []
+            seen_capabilities = set()
+            
+            for route_capability in route_capabilities:
+                route_deps = route_capability.get("dependencies", [])
+                for dep in route_deps:
+                    # dep should already be a string (capability name)
+                    if isinstance(dep, str) and dep not in seen_capabilities:
+                        seen_capabilities.add(dep)
+                        # Convert to object format for registry
+                        all_dependencies.append({
+                            "capability": dep,
+                            "tags": []  # No tags info available at this level
+                        })
+            
+            # If we found dependencies from startup context, use them
+            if all_dependencies:
+                self.logger.info(
+                    f"🔍 Extracted {len(all_dependencies)} unique dependencies from API startup: "
+                    f"{[dep['capability'] for dep in all_dependencies]}"
+                )
+                return all_dependencies
+            
+            # Fallback: try to extract directly from FastAPI app routes
+            fastapi_app = context.get("fastapi_app")
+            if fastapi_app:
+                return self._extract_dependencies_from_routes(fastapi_app)
+                
+            # Final fallback: empty dependencies
+            self.logger.warning("⚠️ No route dependencies found in context or FastAPI app")
+            return []
+            
+        except Exception as e:
+            self.logger.error(f"❌ Failed to extract route dependencies: {e}")
+            return []
+            
+    def _extract_dependencies_from_routes(self, fastapi_app: Any) -> list[dict[str, Any]]:
+        """
+        Fallback method to extract dependencies directly from FastAPI route metadata.
+        
+        Args:
+            fastapi_app: FastAPI application instance
+            
+        Returns:
+            List of unique dependency objects
+        """
+        try:
+            all_dependencies = []
+            seen_capabilities = set()
+            
+            routes = getattr(fastapi_app, "routes", [])
+            for route in routes:
+                endpoint = getattr(route, "endpoint", None)
+                if endpoint and hasattr(endpoint, "_mesh_route_metadata"):
+                    metadata = endpoint._mesh_route_metadata
+                    route_deps = metadata.get("dependencies", [])
+                    
+                    for dep in route_deps:
+                        if isinstance(dep, dict):
+                            capability = dep.get("capability")
+                            if capability and capability not in seen_capabilities:
+                                seen_capabilities.add(capability)
+                                all_dependencies.append({
+                                    "capability": capability,
+                                    "tags": dep.get("tags", [])
+                                })
+                        elif isinstance(dep, str) and dep not in seen_capabilities:
+                            seen_capabilities.add(dep)
+                            all_dependencies.append({
+                                "capability": dep,
+                                "tags": []
+                            })
+            
+            self.logger.info(
+                f"🔍 Extracted {len(all_dependencies)} unique dependencies from FastAPI routes: "
+                f"{[dep['capability'] for dep in all_dependencies]}"
+            )
+            return all_dependencies
+            
+        except Exception as e:
+            self.logger.error(f"❌ Failed to extract dependencies from routes: {e}")
+            return []