digitalkin 0.3.1.dev1__py3-none-any.whl → 0.3.2a2__py3-none-any.whl

digitalkin/grpc_servers/module_server.py

@@ -1,36 +1,31 @@
 """Module gRPC server implementation for DigitalKin."""
 
-import uuid
-from pathlib import Path
+from typing import TYPE_CHECKING
 
-import grpc
-from digitalkin_proto.agentic_mesh_protocol.module.v1 import (
+from agentic_mesh_protocol.module.v1 import (
     module_service_pb2,
     module_service_pb2_grpc,
 )
-from digitalkin_proto.agentic_mesh_protocol.module_registry.v1 import (
-    metadata_pb2,
-    module_registry_service_pb2_grpc,
-    registration_pb2,
-)
 
 from digitalkin.grpc_servers._base_server import BaseServer
 from digitalkin.grpc_servers.module_servicer import ModuleServicer
-from digitalkin.grpc_servers.utils.exceptions import ServerError
 from digitalkin.logger import logger
 from digitalkin.models.grpc_servers.models import (
     ClientConfig,
     ModuleServerConfig,
-    SecurityMode,
 )
 from digitalkin.modules._base_module import BaseModule
+from digitalkin.services.registry import GrpcRegistry
+
+if TYPE_CHECKING:
+    from digitalkin.services.registry import RegistryStrategy
 
 
 class ModuleServer(BaseServer):
     """gRPC server for a DigitalKin module.
 
     This server exposes the module's functionality through the ModuleService gRPC interface.
-    It can optionally register itself with a ModuleRegistry server.
+    It can optionally register itself with a Registry server.
 
     Attributes:
         module: The module instance being served.
@@ -49,14 +44,17 @@ class ModuleServer(BaseServer):
 
         Args:
             module_class: The module instance to be served.
-            server_config: Server configuration including registry address if auto-registration is desired.
-            client_config: Client configuration used by services.
+            server_config: Server configuration.
+            client_config: Client configuration used by services and registry connection.
         """
         super().__init__(server_config)
         self.module_class = module_class
         self.server_config = server_config
         self.client_config = client_config
         self.module_servicer: ModuleServicer | None = None
+        self.registry: RegistryStrategy | None = None
+
+        self._prepare_registry_config()
 
     def _register_servicers(self) -> None:
         """Register the module servicer with the gRPC server.
@@ -77,17 +75,34 @@ class ModuleServer(BaseServer):
         )
         logger.debug("Registered Module servicer")
 
+    def _prepare_registry_config(self) -> None:
+        """Prepare registry client config on module_class before server starts.
+
+        This ensures ServicesConfig created by JobManager will have registry config,
+        allowing spawned module instances to inherit the registry configuration.
+        """
+        if not self.client_config:
+            return
+
+        self.module_class.services_config_params["registry"] = {"client_config": self.client_config}
+
+    def _init_registry(self) -> None:
+        """Initialize server-level registry client for registration."""
+        if not self.client_config:
+            return
+
+        self.registry = GrpcRegistry("", "", "", self.client_config)
+
     def start(self) -> None:
         """Start the module server and register with the registry if configured."""
         logger.info("Starting module server", extra={"server_config": self.server_config})
         super().start()
 
-        # If a registry address is provided, register the module
-        if self.server_config.registry_address:
-            try:
-                self._register_with_registry()
-            except Exception:
-                logger.exception("Failed to register with registry")
+        try:
+            self._init_registry()
+            self._register_with_registry()
+        except Exception:
+            logger.exception("Failed to register with registry")
 
         if self.module_servicer is not None:
             logger.debug("Setup post init started", extra={"client_config": self.client_config})
@@ -97,171 +112,80 @@ class ModuleServer(BaseServer):
         """Start the module server and register with the registry if configured."""
         logger.info("Starting module server", extra={"server_config": self.server_config})
         await super().start_async()
-        # If a registry address is provided, register the module
-        if self.server_config.registry_address:
-            try:
-                self._register_with_registry()
-            except Exception:
-                logger.exception("Failed to register with registry")
+
+        try:
+            self._init_registry()
+            self._register_with_registry()
+        except Exception:
+            logger.exception("Failed to register with registry")
 
         if self.module_servicer is not None:
             logger.info("Setup post init started", extra={"client_config": self.client_config})
             await self.module_servicer.job_manager.start()
             self.module_servicer.setup.__post_init__(self.client_config)
 
-    def stop(self, grace: float | None = None) -> None:
-        """Stop the module server and deregister from the registry if needed."""
-        # If registered with a registry, deregister
-        if self.server_config.registry_address:
+    async def stop_async(self, grace: float | None = None) -> None:
+        """Stop the module server with async cleanup.
+
+        Deregisters from registry and stops the server. Modules also become
+        inactive when they stop sending heartbeats as a fallback.
+        """
+        if self.registry is not None:
             try:
-                self._deregister_from_registry()
-            except ServerError:
+                module_id = self.module_class.get_module_id()
+                if module_id and module_id != "unknown":
+                    await self.registry.deregister(module_id)
+            except Exception:
                 logger.exception("Failed to deregister from registry")
 
-        super().stop(grace)
+        await super().stop_async(grace)
 
     def _register_with_registry(self) -> None:
-        """Register this module with the registry server.
-
-        Raises:
-            ServerError: If communication with the registry server fails.
-        """
-        logger.debug(
-            "Registering module with registry at %s",
-            self.server_config.registry_address,
-            extra={"server_config": self.server_config},
-        )
-
-        # Create appropriate channel based on security mode
-        channel = self._create_registry_channel()
-
-        with channel:
-            # Create a stub (client)
-            stub = module_registry_service_pb2_grpc.ModuleRegistryServiceStub(channel)
-
-            # Determine module type
-            module_type = self._determine_module_type()
-
-            metadata = metadata_pb2.Metadata(
-                name=self.module_class.metadata["name"],
-                tags=[metadata_pb2.Tag(tag=tag) for tag in self.module_class.metadata["tags"]],
-                description=self.module_class.metadata["description"],
+        """Register this module with the registry server."""
+        if not self.registry:
+            logger.debug("No registry configured, skipping registration")
+            return
+
+        module_id = self.module_class.get_module_id()
+        version = self.module_class.metadata.get("version", "0.0.0")
+
+        if not module_id or module_id == "unknown":
+            logger.warning(
+                "Module has no valid module_id, skipping registration",
+                extra={"module_class": self.module_class.__name__},
             )
-
-            self.module_class.metadata["module_id"] = f"{self.module_class.metadata['name']}:{uuid.uuid4()}"
-            # Create registration request
-            request = registration_pb2.RegisterRequest(
-                module_id=self.module_class.metadata["module_id"],
-                version=self.module_class.metadata["version"],
-                module_type=module_type,
-                address=self.server_config.address,
-                metadata=metadata,
-            )
-
-            try:
-                # Call the register method
-                logger.debug(
-                    "Request sent to registry for module: %s:%s",
-                    self.module_class.metadata["name"],
-                    self.module_class.metadata["module_id"],
-                    extra={"module_info": self.module_class.metadata},
-                )
-                response = stub.RegisterModule(request)
-
-                if response.success:
-                    logger.debug("Module registered successfully")
-                else:
-                    logger.error("Module registration failed")
-            except grpc.RpcError:
-                logger.exception("RPC error during registration:")
-                raise ServerError
-
-    def _deregister_from_registry(self) -> None:
-        """Deregister this module from the registry server.
-
-        Raises:
-            ServerError: If communication with the registry server fails.
-        """
-        logger.debug(
-            "Deregistering module from registry at %s",
-            self.server_config.registry_address,
+            return
+
+        advertise_address = self.server_config.advertise_host or self.server_config.host
+
+        logger.info(
+            "Attempting to register module with registry",
+            extra={
+                "module_id": module_id,
+                "address": advertise_address,
+                "port": self.server_config.port,
+                "version": version,
+            },
         )
 
-        # Create appropriate channel based on security mode
-        channel = self._create_registry_channel()
-
-        with channel:
-            # Create a stub (client)
-            stub = module_registry_service_pb2_grpc.ModuleRegistryServiceStub(channel)
+        result = self.registry.register(
+            module_id=module_id,
+            address=advertise_address,
+            port=self.server_config.port,
+            version=version,
+        )
 
-            # Create deregistration request
-            request = registration_pb2.DeregisterRequest(
-                module_id=self.module_class.metadata["module_id"],
+        if result:
+            logger.info(
+                "Module registered successfully",
+                extra={
+                    "module_id": result.module_id,
+                    "address": advertise_address,
+                    "port": self.server_config.port,
+                },
             )
-        try:
-            # Call the deregister method
-            response = stub.DeregisterModule(request)
-
-            if response.success:
-                logger.debug("Module deregistered successfull")
-            else:
-                logger.error("Module deregistration failed")
-        except grpc.RpcError:
-            logger.exception("RPC error during deregistration")
-            raise ServerError
-
-    def _create_registry_channel(self) -> grpc.Channel:
-        """Create an appropriate channel to the registry server.
-
-        Returns:
-            A gRPC channel for communication with the registry.
-
-        Raises:
-            ValueError: If credentials are required but not provided.
-        """
-        if (
-            self.client_config is not None
-            and self.client_config.security == SecurityMode.SECURE
-            and self.client_config.credentials
-        ):
-            # Secure channel
-            # Secure channel
-            root_certificates = Path(self.client_config.credentials.root_cert_path).read_bytes()
-
-            # mTLS channel
-            private_key = None
-            certificate_chain = None
-            if (
-                self.client_config.credentials.client_cert_path is not None
-                and self.client_config.credentials.client_key_path is not None
-            ):
-                private_key = Path(self.client_config.credentials.client_key_path).read_bytes()
-                certificate_chain = Path(self.client_config.credentials.client_cert_path).read_bytes()
-
-            # Create channel credentials
-            channel_credentials = grpc.ssl_channel_credentials(
-                root_certificates=root_certificates,
-                certificate_chain=certificate_chain,
-                private_key=private_key,
+        else:
+            logger.warning(
+                "Module registration returned None (module may not exist in registry)",
+                extra={"module_id": module_id, "address": advertise_address},
             )
-            return grpc.secure_channel(self.server_config.registry_address, channel_credentials)
-        # Insecure channel
-        return grpc.insecure_channel(self.server_config.registry_address)
-
-    def _determine_module_type(self) -> str:
-        """Determine the module type based on its class.
-
-        Returns:
-            A string representing the module type.
-        """
-        module_type = "UNKNOWN"
-        class_name = self.module_class.__name__
-
-        if class_name == "ToolModule":
-            module_type = "TOOL"
-        elif class_name == "TriggerModule":
-            module_type = "TRIGGER"
-        elif class_name == "ArchetypeModule":
-            module_type = "KIN"
-
-        return module_type

digitalkin/grpc_servers/module_servicer.py

@@ -1,11 +1,12 @@
 """Module servicer implementation for DigitalKin."""
 
+import asyncio
 from argparse import ArgumentParser, Namespace
 from collections.abc import AsyncGenerator
 from typing import Any
 
 import grpc
-from digitalkin_proto.agentic_mesh_protocol.module.v1 import (
+from agentic_mesh_protocol.module.v1 import (
     information_pb2,
     lifecycle_pb2,
     module_service_pb2_grpc,
@@ -17,8 +18,9 @@ from digitalkin.core.job_manager.base_job_manager import BaseJobManager
 from digitalkin.grpc_servers.utils.exceptions import ServicerError
 from digitalkin.logger import logger
 from digitalkin.models.core.job_manager_models import JobManagerMode
-from digitalkin.models.module.module import ModuleStatus
+from digitalkin.models.module.module import ModuleCodeModel, ModuleStatus
 from digitalkin.modules._base_module import BaseModule
+from digitalkin.services.registry import GrpcRegistry, RegistryStrategy
 from digitalkin.services.services_models import ServicesMode
 from digitalkin.services.setup.default_setup import DefaultSetup
 from digitalkin.services.setup.grpc_setup import GrpcSetup
@@ -40,6 +42,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
     args: Namespace
     setup: SetupStrategy
     job_manager: BaseJobManager
+    _registry_cache: RegistryStrategy | None = None
 
     def _add_parser_args(self, parser: ArgumentParser) -> None:
         super()._add_parser_args(parser)
@@ -82,6 +85,26 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         )
         self.setup = GrpcSetup() if self.args.services_mode == ServicesMode.REMOTE else DefaultSetup()
 
+    def _get_registry(self) -> RegistryStrategy | None:
+        """Get a cached registry instance if configured.
+
+        Returns:
+            Cached GrpcRegistry instance if registry config exists, None otherwise.
+        """
+        if self._registry_cache is not None:
+            return self._registry_cache
+
+        registry_config = self.module_class.services_config_params.get("registry")
+        if not registry_config:
+            return None
+
+        client_config = registry_config.get("client_config")
+        if not client_config:
+            return None
+
+        self._registry_cache = GrpcRegistry("", "", "", client_config)
+        return self._registry_cache
+
     async def ConfigSetupModule(  # noqa: N802
         self,
         request: lifecycle_pb2.ConfigSetupModuleRequest,
@@ -108,11 +131,9 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
                 "mission_id": request.mission_id,
             },
         )
-        # Process the module input
-        # TODO: Secret should be used here as well
         setup_version = request.setup_version
         config_setup_data = self.module_class.create_config_setup_model(json_format.MessageToDict(request.content))
-        setup_version_data = self.module_class.create_setup_model(
+        setup_version_data = await self.module_class.create_setup_model(
             json_format.MessageToDict(request.setup_version.content),
             config_fields=True,
         )
@@ -139,8 +160,33 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
             return lifecycle_pb2.ConfigSetupModuleResponse(success=False)
 
         updated_setup_data = await self.job_manager.generate_config_setup_module_response(job_id)
-        logger.info("Setup updated")
-        logger.debug(f"Updated setup data: {updated_setup_data=}")
+        logger.info("Setup response received", extra={"job_id": job_id})
+
+        # Check if response is an error
+        if isinstance(updated_setup_data, ModuleCodeModel):
+            logger.error(
+                "Config setup failed",
+                extra={"job_id": job_id, "code": updated_setup_data.code, "error_message": updated_setup_data.message},
+            )
+            context.set_code(grpc.StatusCode.INTERNAL)
+            context.set_details(updated_setup_data.message or "Config setup failed")
+            return lifecycle_pb2.ConfigSetupModuleResponse(success=False)
+
+        if isinstance(updated_setup_data, dict) and "code" in updated_setup_data:
+            # ModuleCodeModel was serialized to dict
+            logger.error(
+                "Config setup failed",
+                extra={
+                    "job_id": job_id,
+                    "code": updated_setup_data["code"],
+                    "error_message": updated_setup_data.get("message"),
+                },
+            )
+            context.set_code(grpc.StatusCode.INTERNAL)
+            context.set_details(updated_setup_data.get("message") or "Config setup failed")
+            return lifecycle_pb2.ConfigSetupModuleResponse(success=False)
+
+        logger.debug("Updated setup data", extra={"job_id": job_id, "setup_data": updated_setup_data})
         setup_version.content = json_format.ParseDict(
             updated_setup_data,
             struct_pb2.Struct(),
@@ -185,7 +231,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
             msg = "No setup data returned."
             raise ServicerError(msg)
 
-        setup_data = self.module_class.create_setup_model(setup_data_class.current_setup_version.content)
+        setup_data = await self.module_class.create_setup_model(setup_data_class.current_setup_version.content)
 
         # create a task to run the module in background
         job_id = await self.job_manager.create_module_instance_job(
@@ -205,6 +251,11 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         try:
             async with self.job_manager.generate_stream_consumer(job_id) as stream:  # type: ignore
                 async for message in stream:
+                    # Early detection of client disconnection
+                    if context.cancelled():
+                        logger.info("Client disconnected", extra={"job_id": job_id})
+                        break
+
                     if message.get("error", None) is not None:
                         logger.error("Error in output_data", extra={"message": message})
                         context.set_code(message["error"]["code"])
@@ -219,19 +270,34 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
                         yield lifecycle_pb2.StartModuleResponse(success=False, job_id=job_id)
                         break
 
-                    if message.get("code", None) is not None and message.get("code") == "__END_OF_STREAM__":
+                    logger.info("Yielding message from job %s: %s", job_id, message)
+                    proto = json_format.ParseDict(message, struct_pb2.Struct(), ignore_unknown_fields=True)
+                    yield lifecycle_pb2.StartModuleResponse(success=True, output=proto, job_id=job_id)
+
+                    if message.get("root", {}).get("protocol") == "end_of_stream":
                         logger.info(
-                            "End of stream via __END_OF_STREAM__",
+                            "End of stream signal received",
                             extra={"job_id": job_id, "mission_id": request.mission_id},
                         )
                         break
-
-                    logger.info("Yielding message from job %s: %s", job_id, message)
-                    proto = json_format.ParseDict(message, struct_pb2.Struct(), ignore_unknown_fields=True)
-                    yield lifecycle_pb2.StartModuleResponse(success=True, output=proto, job_id=job_id)
         finally:
-            await self.job_manager.wait_for_completion(job_id)
-            await self.job_manager.clean_session(job_id, mission_id=request.mission_id)
+            try:
+                await asyncio.wait_for(
+                    self.job_manager.wait_for_completion(job_id),
+                    timeout=30.0,
+                )
+            except Exception:
+                logger.exception(
+                    "Error waiting for job completion",
+                    extra={"job_id": job_id, "mission_id": request.mission_id},
+                )
+            try:
+                await self.job_manager.clean_session(job_id, mission_id=request.mission_id)
+            except Exception:
+                logger.exception(
+                    "Error cleaning session",
+                    extra={"job_id": job_id, "mission_id": request.mission_id},
+                )
 
         logger.info("Job %s finished", job_id)
 
@@ -249,17 +315,19 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         Returns:
             A response indicating success or failure.
         """
-        logger.debug("StopModule called for module: '%s'", self.module_class.__name__)
+        logger.debug(
+            "StopModule called",
+            extra={"module_class": self.module_class.__name__, "job_id": request.job_id},
+        )
 
         response: bool = await self.job_manager.stop_module(request.job_id)
         if not response:
-            message = f"Job {request.job_id} not found"
-            logger.warning(message)
+            logger.warning("Job not found for stop request", extra={"job_id": request.job_id})
             context.set_code(grpc.StatusCode.NOT_FOUND)
-            context.set_details(message)
+            context.set_details(f"Job {request.job_id} not found")
             return lifecycle_pb2.StopModuleResponse(success=False)
 
-        logger.debug("Job %s stopped successfully", request.job_id, extra={"job_id": request.job_id})
+        logger.debug("Job stopped successfully", extra={"job_id": request.job_id})
         return lifecycle_pb2.StopModuleResponse(success=True)
 
     async def GetModuleStatus(  # noqa: N802
@@ -350,7 +418,9 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         # Get input schema if available
         try:
             # Convert schema to proto format
-            input_schema_proto = self.module_class.get_input_format(llm_format=request.llm_format)
+            input_schema_proto = await self.module_class.get_input_format(
+                llm_format=request.llm_format,
+            )
             input_format_struct = json_format.Parse(
                 text=input_schema_proto,
                 message=struct_pb2.Struct(),  # pylint: disable=no-member
@@ -386,7 +456,9 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         # Get output schema if available
         try:
             # Convert schema to proto format
-            output_schema_proto = self.module_class.get_output_format(llm_format=request.llm_format)
+            output_schema_proto = await self.module_class.get_output_format(
+                llm_format=request.llm_format,
+            )
             output_format_struct = json_format.Parse(
                 text=output_schema_proto,
                 message=struct_pb2.Struct(),  # pylint: disable=no-member
@@ -422,7 +494,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         # Get setup schema if available
         try:
             # Convert schema to proto format
-            setup_schema_proto = self.module_class.get_setup_format(llm_format=request.llm_format)
+            setup_schema_proto = await self.module_class.get_setup_format(llm_format=request.llm_format)
             setup_format_struct = json_format.Parse(
                 text=setup_schema_proto,
                 message=struct_pb2.Struct(),  # pylint: disable=no-member
@@ -439,7 +511,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
             setup_schema=setup_format_struct,
         )
 
-    def GetModuleSecret(  # noqa: N802
+    async def GetModuleSecret(  # noqa: N802
         self,
         request: information_pb2.GetModuleSecretRequest,
         context: grpc.ServicerContext,
@@ -458,7 +530,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         # Get secret schema if available
         try:
             # Convert schema to proto format
-            secret_schema_proto = self.module_class.get_secret_format(llm_format=request.llm_format)
+            secret_schema_proto = await self.module_class.get_secret_format(llm_format=request.llm_format)
             secret_format_struct = json_format.Parse(
                 text=secret_schema_proto,
                 message=struct_pb2.Struct(),  # pylint: disable=no-member
@@ -494,7 +566,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         # Get setup schema if available
         try:
             # Convert schema to proto format
-            config_setup_schema_proto = self.module_class.get_config_setup_format(llm_format=request.llm_format)
+            config_setup_schema_proto = await self.module_class.get_config_setup_format(llm_format=request.llm_format)
             config_setup_format_struct = json_format.Parse(
                 text=config_setup_schema_proto,
                 message=struct_pb2.Struct(),  # pylint: disable=no-member
@@ -510,3 +582,37 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
             success=True,
             config_setup_schema=config_setup_format_struct,
         )
+
+    async def GetModuleCost(  # noqa: N802
+        self,
+        request: information_pb2.GetModuleCostRequest,
+        context: grpc.ServicerContext,
+    ) -> information_pb2.GetModuleCostResponse:
+        """Get information about the module's cost configuration.
+
+        Args:
+            request: The get module cost request.
+            context: The gRPC context.
+
+        Returns:
+            A response with the module's cost schema.
+        """
+        logger.debug("GetModuleCost called for module: '%s'", self.module_class.__name__)
+
+        try:
+            cost_schema_proto = await self.module_class.get_cost_format(llm_format=request.llm_format)
+            cost_format_struct = json_format.Parse(
+                text=cost_schema_proto,
+                message=struct_pb2.Struct(),
+                ignore_unknown_fields=True,
+            )
+        except NotImplementedError as e:
+            logger.warning(e)
+            context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+            context.set_details(str(e))
+            return information_pb2.GetModuleCostResponse()
+
+        return information_pb2.GetModuleCostResponse(
+            success=True,
+            cost_schema=cost_format_struct,
+        )