agentscope-runtime 1.0.2__py3-none-any.whl → 1.0.3__py3-none-any.whl

agentscope_runtime/engine/deployers/adapter/a2a/nacos_a2a_registry.py

@@ -0,0 +1,640 @@
+# -*- coding: utf-8 -*-
+"""
+Nacos Registry Implementation for A2A Protocol
+
+This module provides a Nacos-based registry implementation for A2A protocol
+agent service registration. It serves as the default registry implementation
+for A2A protocol adapters.
+"""
+import asyncio
+import logging
+from enum import Enum
+from typing import Optional, TYPE_CHECKING, List
+import threading
+
+from a2a.types import AgentCard
+
+# Make the v2.nacos imports optional to avoid hard dependency at
+# module import time.
+if TYPE_CHECKING:
+    from v2.nacos import ClientConfig, ClientConfigBuilder
+    from v2.nacos.ai.model.ai_param import (
+        RegisterAgentEndpointParam,
+        ReleaseAgentCardParam,
+    )
+    from v2.nacos.ai.nacos_ai_service import NacosAIService
+
+    _NACOS_SDK_AVAILABLE = True
+else:
+    try:
+        from v2.nacos import ClientConfig, ClientConfigBuilder
+        from v2.nacos.ai.model.ai_param import (
+            RegisterAgentEndpointParam,
+            ReleaseAgentCardParam,
+        )
+        from v2.nacos.ai.nacos_ai_service import NacosAIService
+
+        _NACOS_SDK_AVAILABLE = True
+    except Exception:
+
+        class ClientConfig:
+            pass
+
+        class ClientConfigBuilder:
+            pass
+
+        class RegisterAgentEndpointParam:
+            pass
+
+        class ReleaseAgentCardParam:
+            pass
+
+        class NacosAIService:
+            pass
+
+        _NACOS_SDK_AVAILABLE = False
+
+# Import after conditional imports to avoid circular dependencies
+from .a2a_registry import (  # pylint: disable=wrong-import-position
+    A2ARegistry,
+    A2ATransportsProperties,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class RegistrationStatus(Enum):
+    """Registration task status."""
+
+    PENDING = "pending"
+    IN_PROGRESS = "in_progress"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    CANCELLED = "cancelled"
+
+
+class NacosRegistry(A2ARegistry):
+    """Nacos-based registry implementation for A2A protocol.
+
+    This registry registers A2A agent services to Nacos service
+    discovery system. It performs two-step registration:
+    1. Publishes agent card to Nacos (agent metadata)
+    2. Registers agent endpoint with host/port information (service location)
+
+    Args:
+        nacos_client_config: Optional Nacos client configuration.
+            If None, creates default config from environment variables.
+    """
+
+    def __init__(
+        self,
+        nacos_client_config: Optional[ClientConfig] = None,
+    ):
+        self._nacos_client_config = nacos_client_config
+        self._nacos_ai_service: Optional[NacosAIService] = None
+        self._register_task: Optional[asyncio.Task] = None
+        self._register_thread: Optional[threading.Thread] = None
+        self._registration_status: RegistrationStatus = (
+            RegistrationStatus.PENDING
+        )
+        self._registration_lock = threading.Lock()
+        self._shutdown_event = threading.Event()
+
+    def registry_name(self) -> str:
+        """Get the name of this registry implementation.
+
+        Returns:
+            Registry name: "nacos"
+        """
+        return "nacos"
+
+    def register(
+        self,
+        agent_card: AgentCard,
+        a2a_transports_properties: Optional[
+            List[A2ATransportsProperties]
+        ] = None,
+    ) -> None:
+        """Register an A2A agent service to Nacos.
+
+        This method initiates registration asynchronously in the
+        background. Exceptions during registration are caught and
+        logged, but not raised. Use `get_registration_status()` or
+        `wait_for_registration()` to check the result.
+
+        Args:
+            agent_card: The complete A2A agent card generated by runtime
+            a2a_transports_properties: List of transport configurations.
+                Each transport will be registered separately.
+        """
+        # If Nacos SDK is not available, log and return
+        if not _NACOS_SDK_AVAILABLE:
+            logger.debug(
+                "[NacosRegistry] Nacos SDK is not available; "
+                "skipping registration",
+            )
+            return
+
+        self._start_register_task(
+            agent_card=agent_card,
+            a2a_transports_properties=a2a_transports_properties,
+        )
+
+    def _start_register_task(
+        self,
+        agent_card: AgentCard,
+        a2a_transports_properties: Optional[
+            List[A2ATransportsProperties]
+        ] = None,
+    ) -> None:
+        """Start background Nacos registration task.
+
+        If there is an active asyncio event loop, schedule the
+        registration as a task on that loop. Otherwise, spawn a daemon
+        thread and run the registration using asyncio.run so
+        registration still occurs in synchronous contexts.
+        """
+        # All status checks and updates must be within the lock to
+        # ensure atomicity
+        with self._registration_lock:
+            # Check if shutdown was already requested (inside lock
+            # for atomicity)
+            if self._shutdown_event.is_set():
+                logger.info(
+                    "[NacosRegistry] Shutdown already requested, "
+                    "skipping registration",
+                )
+                self._registration_status = RegistrationStatus.CANCELLED
+                return
+
+            # Check if registration is already in progress or completed
+            if self._registration_status in (
+                RegistrationStatus.IN_PROGRESS,
+                RegistrationStatus.COMPLETED,
+            ):
+                logger.warning(
+                    "[NacosRegistry] Registration already in "
+                    "progress or completed, skipping",
+                )
+                return
+            self._registration_status = RegistrationStatus.PENDING
+
+        try:
+            loop = None
+            try:
+                loop = asyncio.get_running_loop()
+            except RuntimeError:
+                # No running loop in this thread; we'll fall back to a
+                # background thread
+                loop = None
+
+            if loop is not None:
+                with self._registration_lock:
+                    self._registration_status = RegistrationStatus.IN_PROGRESS
+                self._register_task = loop.create_task(
+                    self._register_to_nacos(
+                        agent_card=agent_card,
+                        a2a_transports_properties=a2a_transports_properties,
+                    ),
+                )
+                logger.info(
+                    "[NacosRegistry] Registration task scheduled on "
+                    "running event loop",
+                )
+                return
+
+            # No running loop: use a background thread to run asyncio.run
+            def _thread_runner():
+                try:
+                    with self._registration_lock:
+                        if self._shutdown_event.is_set():
+                            logger.info(
+                                "[NacosRegistry] Shutdown requested "
+                                "before registration started",
+                            )
+                            self._registration_status = (
+                                RegistrationStatus.CANCELLED
+                            )
+                            return
+                        self._registration_status = (
+                            RegistrationStatus.IN_PROGRESS
+                        )
+
+                    asyncio.run(
+                        self._register_to_nacos(
+                            agent_card=agent_card,
+                            a2a_transports_properties=(
+                                a2a_transports_properties
+                            ),
+                        ),
+                    )
+                    with self._registration_lock:
+                        if (
+                            self._registration_status
+                            == RegistrationStatus.IN_PROGRESS
+                        ):
+                            self._registration_status = (
+                                RegistrationStatus.COMPLETED
+                            )
+                except asyncio.CancelledError:
+                    with self._registration_lock:
+                        self._registration_status = (
+                            RegistrationStatus.CANCELLED
+                        )
+                    logger.info("[NacosRegistry] Registration cancelled")
+                except Exception:
+                    with self._registration_lock:
+                        self._registration_status = RegistrationStatus.FAILED
+                    logger.error(
+                        "[NacosRegistry] Registration failed in "
+                        "background thread",
+                        exc_info=True,
+                    )
+
+            thread = threading.Thread(
+                target=_thread_runner,
+                name="nacos-registry-register",
+                daemon=True,
+            )
+            thread.start()
+            # Store thread reference after successful start for
+            # proper tracking and cleanup
+            with self._registration_lock:
+                self._register_thread = thread
+            logger.info(
+                "[NacosRegistry] Registration task started in "
+                "background thread",
+            )
+        except Exception:
+            with self._registration_lock:
+                self._registration_status = RegistrationStatus.FAILED
+            logger.warning(
+                "[NacosRegistry] Error starting registration task",
+                exc_info=True,
+            )
+
+    def _get_client_config(self) -> ClientConfig:
+        """Get Nacos client configuration.
+
+        Returns:
+            ClientConfig: Nacos client configuration
+        """
+        if self._nacos_client_config is not None:
+            return self._nacos_client_config
+
+        # Use centralized config builder from a2a_registry module
+        # This ensures consistent behavior with env-based registry creation
+        from .a2a_registry import (
+            get_registry_settings,
+            _build_nacos_client_config,
+        )
+
+        settings = get_registry_settings()
+        return _build_nacos_client_config(settings)
+
+    # pylint: disable=too-many-branches,too-many-statements
+    async def _register_to_nacos(
+        self,
+        agent_card: AgentCard,
+        a2a_transports_properties: Optional[
+            List[A2ATransportsProperties]
+        ] = None,
+    ) -> None:
+        """Register agent to Nacos.
+
+        Performs two-step registration:
+        1. Release agent card to Nacos
+        2. Register agent endpoint with host/port information
+
+        On successful registration, the NacosAIService connection is kept
+        alive to maintain heartbeat and health checks. The service will be
+        closed in cleanup() during application shutdown.
+
+        Most exceptions are caught and logged without re-raising to prevent
+        background tasks from crashing the host. Only `asyncio.CancelledError`
+        is re-raised.
+
+        Args:
+            agent_card: The A2A agent card to register
+            a2a_transports_properties: Optional list of transport
+                configurations
+
+        Raises:
+            asyncio.CancelledError: If the registration task is cancelled
+        """
+        # Check if shutdown was requested
+        if self._shutdown_event.is_set():
+            with self._registration_lock:
+                self._registration_status = RegistrationStatus.CANCELLED
+            logger.info(
+                "[NacosRegistry] Registration cancelled due to shutdown",
+            )
+            return
+
+        client_config = self._get_client_config()
+        try:
+            self._nacos_ai_service = await NacosAIService.create_ai_service(
+                client_config,
+            )
+
+            # Check again after service creation
+            if self._shutdown_event.is_set():
+                with self._registration_lock:
+                    self._registration_status = RegistrationStatus.CANCELLED
+                logger.info(
+                    "[NacosRegistry] Registration cancelled after "
+                    "service creation",
+                )
+                return
+
+            # Publish agent card
+            await self._nacos_ai_service.release_agent_card(
+                ReleaseAgentCardParam(agent_card=agent_card),
+            )
+            logger.info(
+                "[NacosRegistry] Agent card published: agent=%s",
+                agent_card.name,
+            )
+
+            # Check again before endpoint registration
+            if self._shutdown_event.is_set():
+                with self._registration_lock:
+                    self._registration_status = RegistrationStatus.CANCELLED
+                logger.warning(
+                    "[NacosRegistry] Registration cancelled after "
+                    "card publish, endpoint may be partially "
+                    "registered",
+                )
+                return
+
+            # Register agent endpoint
+            if a2a_transports_properties:
+                for transport in a2a_transports_properties:
+                    host = transport.host
+                    port = transport.port
+                    endpoint_param = RegisterAgentEndpointParam(
+                        agent_name=agent_card.name,
+                        version=agent_card.version,
+                        address=host,
+                        port=port,
+                        path=transport.path,
+                        support_tls=transport.support_tls,
+                        transport=transport.transport_type,
+                    )
+
+                    await self._nacos_ai_service.register_agent_endpoint(
+                        endpoint_param,
+                    )
+                    logger.info(
+                        "[NacosRegistry] Agent endpoint registered: "
+                        "agent=%s, address=%s:%s, path=%s",
+                        agent_card.name,
+                        host,
+                        port,
+                        transport.path,
+                    )
+
+            with self._registration_lock:
+                if self._registration_status == RegistrationStatus.IN_PROGRESS:
+                    self._registration_status = RegistrationStatus.COMPLETED
+            # Service remains alive to maintain heartbeat and health checks.
+            # It will be closed in cleanup() during application shutdown.
+
+        except asyncio.CancelledError:
+            with self._registration_lock:
+                self._registration_status = RegistrationStatus.CANCELLED
+            logger.info("[NacosRegistry] Registration task cancelled")
+            try:
+                svc = self._nacos_ai_service
+                if svc is not None:
+                    await svc.shutdown()
+                    logger.debug(
+                        "[NacosRegistry] NacosAIService closed due to "
+                        "cancellation",
+                    )
+            except Exception:
+                logger.debug(
+                    "[NacosRegistry] Error closing NacosAIService on "
+                    "cancellation",
+                    exc_info=True,
+                )
+            raise
+        except Exception as e:
+            with self._registration_lock:
+                self._registration_status = RegistrationStatus.FAILED
+            # Log errors but don't re-raise; background tasks should
+            # not crash the host
+            logger.error(
+                "[NacosRegistry] Failed to register agent=%s: %s",
+                agent_card.name,
+                str(e),
+                exc_info=True,
+            )
+            try:
+                svc = self._nacos_ai_service
+                if svc is not None:
+                    await svc.shutdown()
+                    logger.debug(
+                        "[NacosRegistry] NacosAIService closed due to "
+                        "registration failure",
+                    )
+            except Exception:
+                logger.debug(
+                    "[NacosRegistry] Error closing NacosAIService on failure",
+                    exc_info=True,
+                )
+
+    def get_registration_status(self) -> RegistrationStatus:
+        """Get the current registration status.
+
+        Returns:
+            RegistrationStatus: Current status of the registration task
+        """
+        with self._registration_lock:
+            return self._registration_status
+
+    async def wait_for_registration(
+        self,
+        timeout: Optional[float] = None,
+    ) -> RegistrationStatus:
+        """Wait for registration to complete, fail, or be cancelled.
+
+        Args:
+            timeout: Optional timeout in seconds. If None, waits indefinitely.
+
+        Returns:
+            RegistrationStatus: Final status of the registration
+        """
+        if self._register_task is not None:
+            # Task-based registration: wait for the task
+            try:
+                if timeout is not None:
+                    await asyncio.wait_for(
+                        self._register_task,
+                        timeout=timeout,
+                    )
+                else:
+                    await self._register_task
+            except asyncio.TimeoutError:
+                logger.warning(
+                    "[NacosRegistry] Wait for registration timed "
+                    "out after %s seconds",
+                    timeout,
+                )
+            except (asyncio.CancelledError, Exception):
+                pass  # Task may have been cancelled or failed
+        elif self._register_thread is not None:
+            # Thread-based registration: wait for the thread (matching
+            # LocalDeployManager pattern)
+            self._register_thread.join(timeout=timeout)
+            if self._register_thread.is_alive():
+                logger.warning(
+                    "[NacosRegistry] Registration thread did not "
+                    "complete within timeout",
+                )
+
+        return self.get_registration_status()
+
+    # pylint: disable=too-many-branches,too-many-statements
+    async def cleanup(
+        self,
+        wait_for_completion: bool = True,
+        timeout: Optional[float] = 5.0,
+    ) -> None:
+        """Clean up registration resources and cancel ongoing
+        registration if needed.
+
+        This method should be called during shutdown to ensure proper cleanup.
+        It will:
+        1. Signal shutdown to prevent new registration attempts
+        2. Optionally wait for the registration to complete or timeout
+        3. Cancel the asyncio task if running
+        4. Wait for background thread if running
+        5. Clean up service connections
+
+        Args:
+            wait_for_completion: If True, wait for registration to complete
+                before cancelling. If False, cancel immediately.
+            timeout: Maximum time to wait for completion (if
+                wait_for_completion=True). Default is 5 seconds.
+        """
+        logger.info("[NacosRegistry] Starting cleanup")
+
+        # Signal shutdown first to prevent new registration attempts
+        self._shutdown_event.set()
+
+        with self._registration_lock:
+            current_status = self._registration_status
+
+        # If registration is in progress or pending, handle cancellation
+        if current_status in (
+            RegistrationStatus.IN_PROGRESS,
+            RegistrationStatus.PENDING,
+        ):
+            if (
+                wait_for_completion
+                and current_status == RegistrationStatus.IN_PROGRESS
+            ):
+                logger.info(
+                    "[NacosRegistry] Waiting for registration to "
+                    "complete (timeout=%s)",
+                    timeout,
+                )
+                try:
+                    await self.wait_for_registration(timeout=timeout)
+                    current_status = self.get_registration_status()
+                    if current_status == RegistrationStatus.COMPLETED:
+                        logger.info(
+                            "[NacosRegistry] Registration completed "
+                            "before shutdown",
+                        )
+                    elif current_status == RegistrationStatus.FAILED:
+                        logger.warning(
+                            "[NacosRegistry] Registration failed "
+                            "before shutdown",
+                        )
+                    elif current_status == RegistrationStatus.CANCELLED:
+                        logger.info(
+                            "[NacosRegistry] Registration was cancelled",
+                        )
+                except Exception as e:
+                    logger.warning(
+                        "[NacosRegistry] Error waiting for registration: %s",
+                        e,
+                        exc_info=True,
+                    )
+            else:
+                if current_status == RegistrationStatus.PENDING:
+                    logger.info(
+                        "[NacosRegistry] Registration was pending, "
+                        "marking as cancelled",
+                    )
+                    with self._registration_lock:
+                        self._registration_status = (
+                            RegistrationStatus.CANCELLED
+                        )
+                else:
+                    logger.info(
+                        "[NacosRegistry] Skipping wait, cancelling "
+                        "immediately",
+                    )
+
+            # Cancel the task if it's still running (matching
+            # env_service pattern)
+            if (
+                self._register_task is not None
+                and not self._register_task.done()
+            ):
+                logger.info("[NacosRegistry] Cancelling registration task")
+                self._register_task.cancel()
+                try:
+                    await self._register_task
+                except asyncio.CancelledError:
+                    # Expected when task is cancelled: await on a
+                    # cancelled task raises CancelledError
+                    pass
+                except Exception as e:
+                    logger.debug(
+                        "[NacosRegistry] Error cancelling task: %s",
+                        e,
+                        exc_info=True,
+                    )
+
+            # Wait for thread if it's still running (matching
+            # LocalDeployManager pattern)
+            if (
+                self._register_thread is not None
+                and self._register_thread.is_alive()
+            ):
+                logger.info(
+                    "[NacosRegistry] Waiting for registration thread "
+                    "to finish",
+                )
+                thread_timeout = (
+                    timeout
+                    if (
+                        wait_for_completion
+                        and current_status == RegistrationStatus.IN_PROGRESS
+                    )
+                    else 1.0
+                )
+                self._register_thread.join(timeout=thread_timeout)
+                if self._register_thread.is_alive():
+                    logger.warning(
+                        "[NacosRegistry] Registration thread did not "
+                        "terminate, potential resource leak",
+                    )
+
+        # Always clean up service connection, regardless of status
+        if self._nacos_ai_service is not None:
+            try:
+                await self._nacos_ai_service.shutdown()
+            except Exception:
+                logger.debug(
+                    "[NacosRegistry] Error closing service during cleanup",
+                    exc_info=True,
+                )
+            finally:
+                self._nacos_ai_service = None
+
+        logger.info("[NacosRegistry] Cleanup completed")

agentscope_runtime/engine/deployers/kubernetes_deployer.py

@@ -144,6 +144,7 @@ class KubernetesDeployManager(DeployManager):
         image_tag: str = "latest",
         push_to_registry: bool = False,
         use_cache: bool = True,
+        pypi_mirror: Optional[str] = None,
         **kwargs,
     ) -> Dict[str, Any]:
         """
@@ -170,6 +171,7 @@ class KubernetesDeployManager(DeployManager):
             mount_dir: Mount directory
             runtime_config: K8s runtime configuration
             use_cache: Enable build cache (default: True)
+            pypi_mirror: PyPI mirror URL for pip package installation
             # Backward compatibility
             image_name: Image name
             image_tag: Image tag
@@ -209,6 +211,7 @@ class KubernetesDeployManager(DeployManager):
                     protocol_adapters=protocol_adapters,
                     custom_endpoints=custom_endpoints,
                     use_cache=use_cache,
+                    pypi_mirror=pypi_mirror,
                     **kwargs,
                 )
                 if not built_image_name:

agentscope_runtime/engine/deployers/utils/docker_image_utils/dockerfile_generator.py

@@ -22,6 +22,7 @@ class DockerfileConfig(BaseModel):
     health_check_endpoint: str = "/health"
     custom_template: Optional[str] = None
     platform: Optional[str] = None
+    pypi_mirror: Optional[str] = None
 
 
 class DockerfileGenerator:
@@ -73,8 +74,7 @@ COPY . {working_dir}/
 # Install Python dependencies
 RUN pip install --no-cache-dir --upgrade pip
 RUN if [ -f requirements.txt ]; then \\
-        pip install --no-cache-dir -r requirements.txt \\
-        -i https://pypi.tuna.tsinghua.edu.cn/simple; fi
+        pip install --no-cache-dir -r requirements.txt{pypi_mirror_flag}; fi
 
 # Create non-root user for security
 RUN adduser --disabled-password --gecos '' {user} && \\
@@ -136,6 +136,11 @@ HEALTHCHECK --interval=30s --timeout=10s --start-period=5s --retries=3 \\
                 f'"--port", "{config.port}"]'
             )
 
+        # Prepare PyPI mirror flag
+        pypi_mirror_flag = ""
+        if config.pypi_mirror:
+            pypi_mirror_flag = f" -i {config.pypi_mirror}"
+
         # Format template with configuration values
         content = template.format(
             base_image=config.base_image,
@@ -147,6 +152,7 @@ HEALTHCHECK --interval=30s --timeout=10s --start-period=5s --retries=3 \\
             env_vars_section=env_vars_section,
             startup_command_section=startup_command_section,
             platform=config.platform,
+            pypi_mirror_flag=pypi_mirror_flag,
         )
 
         return content