mcp-mesh 0.6.3__py3-none-any.whl → 0.7.0__py3-none-any.whl

_mcp_mesh/pipeline/api_heartbeat/api_heartbeat_send.py

@@ -8,6 +8,7 @@ for FastAPI applications using @mesh.route decorators.
 import logging
 from typing import Any
 
+from ...engine.decorator_registry import DecoratorRegistry
 from ..shared.base_step import PipelineStep
 from ..shared.pipeline_types import PipelineResult
 
@@ -17,7 +18,7 @@ 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.
     """
@@ -49,12 +50,11 @@ class APIHeartbeatSendStep(PipelineStep):
             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
+                    status=PipelineStatus.FAILED, message=error_msg, context=context
                 )
                 result.add_error(error_msg)
                 return result
@@ -62,12 +62,11 @@ class APIHeartbeatSendStep(PipelineStep):
             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
+                    status=PipelineStatus.FAILED, message=error_msg, context=context
                 )
                 result.add_error(error_msg)
                 return result
@@ -81,75 +80,111 @@ class APIHeartbeatSendStep(PipelineStep):
 
             # Send heartbeat to registry using the same format as test_api_service.json
             # Import json at the beginning
-            import aiohttp
             import json
-            
+
+            import aiohttp
+
             try:
                 # For API services, send directly to registry using the format that works
-                # Get registry URL 
+                # 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)
-                
+
+                # Build per-route tool entries using METHOD:path as unique identifiers
+                # This allows dependency resolution to map back to specific route wrappers
+                route_wrappers = DecoratorRegistry.get_all_route_wrappers()
+                tools_list = []
+
+                for route_id, route_info in route_wrappers.items():
+                    dependencies = route_info.get("dependencies", [])
+                    if dependencies:  # Only include routes with dependencies
+                        tools_list.append(
+                            {
+                                "function_name": route_id,  # e.g., "GET:/api/v1/benchmark-services"
+                                "dependencies": [
+                                    {"capability": dep, "tags": []}
+                                    for dep in dependencies
+                                ],
+                            }
+                        )
+
+                # Fallback to old behavior if no route wrappers registered yet
+                if not tools_list:
+                    all_route_dependencies = self._extract_all_route_dependencies(
+                        context
+                    )
+                    if all_route_dependencies:
+                        tools_list.append(
+                            {
+                                "function_name": "api_endpoint_handler",
+                                "dependencies": all_route_dependencies,
+                            }
+                        )
+
+                self.logger.debug(
+                    f"Route wrappers registered: {list(route_wrappers.keys())}"
+                )
+
                 api_service_payload = {
                     "agent_id": service_id,
                     "agent_type": "api",
-                    "tools": [
-                        {
-                            "function_name": "api_endpoint_handler",
-                            "dependencies": all_route_dependencies
-                        }
-                    ],
+                    "tools": tools_list,
                     "http_host": display_config.get("display_host", "127.0.0.1"),
-                    "http_port": display_config.get("display_port", 8080)
+                    "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)}")
-                
+
+                self.logger.debug(
+                    f"Sending API service payload to {registry_url}/heartbeat"
+                )
+
                 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)
+                            data=json.dumps(api_service_payload),
                         ) as response:
-                            self.logger.debug(f"{registry_url} \"POST /heartbeat HTTP/1.1\" {response.status}")
+                            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}")
-                
+                                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}'")
-                    
+                    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
+                        status=PipelineStatus.FAILED, message=error_msg, context=context
                     )
                     result.add_error(error_msg)
                     return result
@@ -157,12 +192,11 @@ class APIHeartbeatSendStep(PipelineStep):
             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
+                    status=PipelineStatus.FAILED, message=error_msg, context=context
                 )
                 result.add_error(str(e))
                 return result
@@ -172,10 +206,9 @@ class APIHeartbeatSendStep(PipelineStep):
             self.logger.error(f"❌ {error_msg}")
 
             from ..shared.pipeline_types import PipelineStatus
+
             result = PipelineResult(
-                status=PipelineStatus.FAILED,
-                message=error_msg,
-                context=context
+                status=PipelineStatus.FAILED, message=error_msg, context=context
             )
             result.add_error(str(e))
             return result
@@ -190,7 +223,7 @@ class APIHeartbeatSendStep(PipelineStep):
             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")
@@ -207,12 +240,26 @@ class APIHeartbeatSendStep(PipelineStep):
                     "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", {}),
-                }
+                "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
@@ -222,19 +269,21 @@ class APIHeartbeatSendStep(PipelineStep):
             return {
                 "service_id": service_id,
                 "service_type": "api",
-                "error": f"Failed to prepare heartbeat data: {e}"
+                "error": f"Failed to prepare heartbeat data: {e}",
             }
 
-    def _extract_all_route_dependencies(self, context: dict[str, Any]) -> list[dict[str, Any]]:
+    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
         """
@@ -242,15 +291,15 @@ class APIHeartbeatSendStep(PipelineStep):
             # 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}")
-            
+
+            self.logger.debug(
+                f"api_service_metadata keys: {list(api_service_metadata.keys())}"
+            )
+
             # 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:
@@ -258,75 +307,80 @@ class APIHeartbeatSendStep(PipelineStep):
                     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
-                        })
-            
+                        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: "
+                self.logger.debug(
+                    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")
+            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]]:
+
+    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", [])
-                                })
+                                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: "
+                            all_dependencies.append({"capability": dep, "tags": []})
+
+            self.logger.debug(
+                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 []
+            return []