Service Discovery in K8s¶

# agent_discovery.py
import os
import requests
import mesh

class ServiceDiscovery:
    def __init__(self):
        # Service name auto-detected from K8s labels
        self.service_name = os.getenv('SERVICE_NAME')  # e.g., "mcp-mesh-hello-world"
        self.namespace = os.getenv('NAMESPACE', 'mcp-mesh')
        self.registry_host = os.getenv('MCP_MESH_REGISTRY_HOST', 'mcp-mesh-registry')
        self.registry_port = os.getenv('MCP_MESH_REGISTRY_PORT', '8000')
        self.registry_url = f"http://{self.registry_host}:{self.registry_port}"

    def discover_by_capability(self, capability):
        """Discover agents by capability through registry"""
        response = requests.get(
            f"{self.registry_url}/agents",
            params={"capability": capability}
        )
        return response.json()

    def discover_by_k8s_service(self, service_name):
        """Direct K8s service discovery using service name"""
        # Use Kubernetes DNS with actual port 8080
        k8s_endpoint = f"http://{service_name}.{self.namespace}:8080"
        return k8s_endpoint

    def get_my_service_url(self):
        """Get this agent's own service URL"""
        return f"http://{self.service_name}.{self.namespace}:8080"

@mesh.agent(name="discovery-aware")
class DiscoveryAwareAgent:
    pass

@mesh.tool(
    capability="discovery_aware",
    dependencies=["weather_service", "data_processor"]
)
def discovery_example(weather_service=None, data_processor=None):
    discovery = ServiceDiscovery()

    # Method 1: Use injected dependencies (recommended)
    if weather_service:
        weather_data = weather_service("London")

    # Method 2: Manual discovery through registry
    weather_agents = discovery.discover_by_capability("weather")

    # Method 3: Direct K8s service call using known service names
    hello_world_url = discovery.discover_by_k8s_service("mcp-mesh-hello-world")
    system_agent_url = discovery.discover_by_k8s_service("mcp-mesh-system-agent")

    # Method 4: Get own service URL
    my_url = discovery.get_my_service_url()

    return {
        "injected": weather_data if weather_service else None,
        "registry": weather_agents,
        "hello_world_endpoint": hello_world_url,
        "system_agent_endpoint": system_agent_url,
        "my_service_url": my_url
    }

# Direct pod access
def access_specific_processor(processor_id):
    """Access specific data processor pod"""
    # Pod DNS format for StatefulSet
    pod_dns = f"data-processor-{processor_id}.data-processor-headless.mcp-mesh.svc.cluster.local"

    response = requests.get(f"http://{pod_dns}:8080/status")
    return response.json()

# Access all pods
def access_all_processors():
    """Access all data processor pods"""
    results = []
    for i in range(3):  # Assuming 3 replicas
        pod_dns = f"data-processor-{i}.data-processor-headless.mcp-mesh.svc.cluster.local"
        try:
            response = requests.get(f"http://{pod_dns}:8080/status")
            results.append(response.json())
        except:
            results.append({"error": f"Pod {i} unreachable"})
    return results

# real_k8s_discovery.py - Based on actual examples/k8s/base/agents/
import os
import requests
import socket
from functools import lru_cache
from kubernetes import client, config

class K8sServiceDiscovery:
    def __init__(self):
        # Load K8s config
        if os.path.exists('/var/run/secrets/kubernetes.io'):
            config.load_incluster_config()
        else:
            config.load_kube_config()

        self.v1 = client.CoreV1Api()
        # 🎯 Auto-detected from environment (matches real examples)
        self.service_name = os.getenv('SERVICE_NAME')  # e.g., "mcp-mesh-hello-world"
        self.namespace = os.getenv('NAMESPACE', 'mcp-mesh')
        self.registry_host = os.getenv('MCP_MESH_REGISTRY_HOST', 'mcp-mesh-registry')
        self.registry_port = os.getenv('MCP_MESH_REGISTRY_PORT', '8000')
        self.registry_url = f"http://{self.registry_host}:{self.registry_port}"

    def get_known_service_endpoints(self):
        """Get all known MCP Mesh agent service endpoints"""
        # 🎯 Based on actual service names from K8s examples
        known_services = {
            'hello-world': 'mcp-mesh-hello-world',
            'system-agent': 'mcp-mesh-system-agent',
            'registry': 'mcp-mesh-registry'
        }

        endpoints = {}
        for logical_name, service_name in known_services.items():
            endpoints[logical_name] = {
                'url': f"http://{service_name}.{self.namespace}:8080",
                'service_name': service_name,
                'dns': f"{service_name}.{self.namespace}.svc.cluster.local"
            }
        return endpoints

    @lru_cache(maxsize=128)
    def discover_service(self, service_name, capability=None):
        """Discover service using actual K8s patterns"""
        endpoints = []

        # Method 1: Direct service name (most reliable)
        try:
            svc = self.v1.read_namespaced_service(service_name, self.namespace)
            if svc.spec.cluster_ip:
                endpoints.append({
                    'url': f"http://{svc.spec.cluster_ip}:8080",
                    'source': 'kubernetes-service',
                    'type': 'service',
                    'service_name': service_name
                })
        except:
            pass

        # Method 2: Registry discovery for capabilities
        if capability and self.registry_url:
            try:
                resp = requests.get(
                    f"{self.registry_url}/agents",
                    params={'capability': capability}
                )
                for agent in resp.json():
                    endpoints.append({
                        'url': agent['endpoint'],
                        'source': 'registry',
                        'type': 'agent',
                        'metadata': agent.get('metadata', {})
                    })
            except:
                pass

        # Method 3: DNS resolution (fallback)
        if not endpoints:
            try:
                ip = socket.gethostbyname(f"{service_name}.{self.namespace}")
                endpoints.append({
                    'url': f"http://{ip}:8080",
                    'source': 'dns',
                    'type': 'resolved',
                    'service_name': service_name
                })
            except:
                pass

        return endpoints

    def call_hello_world_agent(self):
        """Example: Call hello-world agent using service name"""
        url = "http://mcp-mesh-hello-world.mcp-mesh:8080/health"
        response = requests.get(url)
        return response.json()

    def call_system_agent(self):
        """Example: Call system agent using service name"""
        url = "http://mcp-mesh-system-agent.mcp-mesh:8080/health"
        response = requests.get(url)
        return response.json()

# test_service_discovery.py
import pytest
import requests
from kubernetes import client, config

class TestServiceDiscovery:
    @classmethod
    def setup_class(cls):
        config.load_kube_config()
        cls.v1 = client.CoreV1Api()
        cls.namespace = "mcp-mesh"

    def test_dns_resolution(self):
        """Test K8s DNS works for services"""
        # Deploy test service
        test_svc = client.V1Service(
            metadata=client.V1ObjectMeta(name="test-dns"),
            spec=client.V1ServiceSpec(
                selector={"app": "test"},
                ports=[client.V1ServicePort(port=8080)]
            )
        )

        self.v1.create_namespaced_service(self.namespace, test_svc)

        # Test DNS resolution from pod
        pod_manifest = {
            "apiVersion": "v1",
            "kind": "Pod",
            "metadata": {"name": "dns-test"},
            "spec": {
                "containers": [{
                    "name": "test",
                    "image": "busybox",
                    "command": ["nslookup", "test-dns"]
                }],
                "restartPolicy": "Never"
            }
        }

        # Cleanup
        self.v1.delete_namespaced_service("test-dns", self.namespace)

    def test_registry_discovery(self):
        """Test registry-based discovery"""
        registry_url = "http://mcp-mesh-registry:8080"

        # Register test agent
        agent_data = {
            "id": "test-agent-1",
            "capabilities": ["test-capability"],
            "endpoint": "http://test-agent:8080"
        }

        response = requests.post(
            f"{registry_url}/api/v1/agents",
            json=agent_data
        )
        assert response.status_code == 200

        # Discover by capability
        response = requests.get(
            f"{registry_url}/api/v1/agents",
            params={"capability": "test-capability"}
        )
        agents = response.json()
        assert len(agents) > 0
        assert agents[0]["id"] == "test-agent-1"