digitalkin 0.3.1.dev2__py3-none-any.whl → 0.3.2.dev14__py3-none-any.whl

digitalkin/core/task_manager/task_session.py

@@ -84,9 +84,12 @@ class TaskSession:
         self._heartbeat_interval = heartbeat_interval
 
         logger.info(
-            "TaskContext initialized for task: '%s'",
-            task_id,
-            extra={"task_id": task_id, "mission_id": mission_id, "heartbeat_interval": heartbeat_interval},
+            "TaskSession initialized",
+            extra={
+                "task_id": task_id,
+                "mission_id": mission_id,
+                "heartbeat_interval": str(heartbeat_interval),
+            },
         )
 
     @property
@@ -99,6 +102,21 @@ class TaskSession:
         """Task paused status."""
         return self._paused.is_set()
 
+    @property
+    def setup_id(self) -> str:
+        """Get setup_id from module context."""
+        return self.module.context.session.setup_id
+
+    @property
+    def setup_version_id(self) -> str:
+        """Get setup_version_id from module context."""
+        return self.module.context.session.setup_version_id
+
+    @property
+    def session_ids(self) -> dict[str, str]:
+        """Get all session IDs from module context for structured logging."""
+        return self.module.context.session.current_ids()
+
     async def send_heartbeat(self) -> bool:
         """Rate-limited heartbeat with connection resilience.
 
@@ -108,6 +126,8 @@ class TaskSession:
         heartbeat = HeartbeatMessage(
             task_id=self.task_id,
             mission_id=self.mission_id,
+            setup_id=self.setup_id,
+            setup_version_id=self.setup_version_id,
             timestamp=datetime.datetime.now(datetime.timezone.utc),
         )
 
@@ -120,23 +140,17 @@ class TaskSession:
                     return True
             except Exception as e:
                 logger.error(
-                    "Heartbeat exception for task: '%s'",
-                    self.task_id,
-                    extra={"task_id": self.task_id, "error": str(e)},
+                    "Heartbeat exception",
+                    extra={**self.session_ids, "error": str(e)},
                     exc_info=True,
                 )
-            logger.error(
-                "Initial heartbeat failed for task: '%s'",
-                self.task_id,
-                extra={"task_id": self.task_id},
-            )
+            logger.error("Initial heartbeat failed", extra=self.session_ids)
             return False
 
         if (heartbeat.timestamp - self._last_heartbeat) < self._heartbeat_interval:
             logger.debug(
-                "Heartbeat skipped due to rate limiting for task: '%s' | delta=%s",
-                self.task_id,
-                heartbeat.timestamp - self._last_heartbeat,
+                "Heartbeat skipped due to rate limiting",
+                extra={**self.session_ids, "delta": str(heartbeat.timestamp - self._last_heartbeat)},
             )
             return True
 
@@ -147,39 +161,24 @@ class TaskSession:
                 return True
         except Exception as e:
             logger.error(
-                "Heartbeat exception for task: '%s'",
-                self.task_id,
-                extra={"task_id": self.task_id, "error": str(e)},
+                "Heartbeat exception",
+                extra={**self.session_ids, "error": str(e)},
                 exc_info=True,
             )
-        logger.warning(
-            "Heartbeat failed for task: '%s'",
-            self.task_id,
-            extra={"task_id": self.task_id},
-        )
+        logger.warning("Heartbeat failed", extra=self.session_ids)
         return False
 
     async def generate_heartbeats(self) -> None:
         """Periodic heartbeat generator with cancellation support."""
-        logger.debug(
-            "Heartbeat generator started for task: '%s'",
-            self.task_id,
-            extra={"task_id": self.task_id, "mission_id": self.mission_id},
-        )
+        logger.debug("Heartbeat generator started", extra=self.session_ids)
         while not self.cancelled:
             logger.debug(
-                "Heartbeat tick for task: '%s', cancelled=%s",
-                self.task_id,
-                self.cancelled,
-                extra={"task_id": self.task_id, "mission_id": self.mission_id},
+                "Heartbeat tick",
+                extra={**self.session_ids, "cancelled": self.cancelled},
             )
             success = await self.send_heartbeat()
             if not success:
-                logger.error(
-                    "Heartbeat failed, cancelling task: '%s'",
-                    self.task_id,
-                    extra={"task_id": self.task_id, "mission_id": self.mission_id},
-                )
+                logger.error("Heartbeat failed, cancelling task", extra=self.session_ids)
                 await self._handle_cancel(CancellationReason.HEARTBEAT_FAILURE)
                 break
             await asyncio.sleep(self._heartbeat_interval.total_seconds())
@@ -187,11 +186,7 @@ class TaskSession:
     async def wait_if_paused(self) -> None:
         """Block execution if task is paused."""
         if self._paused.is_set():
-            logger.info(
-                "Task paused, waiting for resume: '%s'",
-                self.task_id,
-                extra={"task_id": self.task_id},
-            )
+            logger.info("Task paused, waiting for resume", extra=self.session_ids)
             await self._paused.wait()
 
     async def listen_signals(self) -> None:  # noqa: C901
@@ -200,18 +195,14 @@ class TaskSession:
         Raises:
             CancelledError: Asyncio when task cancelling
         """
-        logger.info(
-            "Signal listener started for task: '%s'",
-            self.task_id,
-            extra={"task_id": self.task_id},
-        )
+        logger.info("Signal listener started", extra=self.session_ids)
         if self.signal_record_id is None:
             self.signal_record_id = (await self.db.select_by_task_id("tasks", self.task_id)).get("id")
 
         live_id, live_signals = await self.db.start_live("tasks")
         try:
             async for signal in live_signals:
-                logger.debug("Signal received for task '%s': %s", self.task_id, signal)
+                logger.debug("Signal received", extra={**self.session_ids, "signal": signal})
                 if self.cancelled:
                     break
 
@@ -228,26 +219,17 @@ class TaskSession:
                     await self._handle_status_request()
 
         except asyncio.CancelledError:
-            logger.debug(
-                "Signal listener cancelled for task: '%s'",
-                self.task_id,
-                extra={"task_id": self.task_id},
-            )
+            logger.debug("Signal listener cancelled", extra=self.session_ids)
             raise
         except Exception as e:
             logger.error(
-                "Signal listener fatal error for task: '%s'",
-                self.task_id,
-                extra={"task_id": self.task_id, "error": str(e)},
+                "Signal listener fatal error",
+                extra={**self.session_ids, "error": str(e)},
                 exc_info=True,
             )
         finally:
             await self.db.stop_live(live_id)
-            logger.info(
-                "Signal listener stopped for task: '%s'",
-                self.task_id,
-                extra={"task_id": self.task_id},
-            )
+            logger.info("Signal listener stopped", extra=self.session_ids)
 
     async def _handle_cancel(self, reason: CancellationReason = CancellationReason.UNKNOWN) -> None:
         """Idempotent cancellation with acknowledgment and reason tracking.
@@ -257,13 +239,9 @@ class TaskSession:
         """
         if self.is_cancelled.is_set():
             logger.debug(
-                "Cancel ignored - task already cancelled: '%s' (existing reason: %s, new reason: %s)",
-                self.task_id,
-                self.cancellation_reason.value,
-                reason.value,
+                "Cancel ignored - already cancelled",
                 extra={
-                    "task_id": self.task_id,
-                    "mission_id": self.mission_id,
+                    **self.session_ids,
                     "existing_reason": self.cancellation_reason.value,
                     "new_reason": reason.value,
                 },
@@ -277,25 +255,13 @@ class TaskSession:
         # Log with appropriate level based on reason
         if reason in {CancellationReason.SUCCESS_CLEANUP, CancellationReason.FAILURE_CLEANUP}:
             logger.debug(
-                "Task cancelled (cleanup): '%s', reason: %s",
-                self.task_id,
-                reason.value,
-                extra={
-                    "task_id": self.task_id,
-                    "mission_id": self.mission_id,
-                    "cancellation_reason": reason.value,
-                },
+                "Task cancelled (cleanup)",
+                extra={**self.session_ids, "cancellation_reason": reason.value},
             )
         else:
             logger.info(
-                "Task cancelled: '%s', reason: %s",
-                self.task_id,
-                reason.value,
-                extra={
-                    "task_id": self.task_id,
-                    "mission_id": self.mission_id,
-                    "cancellation_reason": reason.value,
-                },
+                "Task cancelled",
+                extra={**self.session_ids, "cancellation_reason": reason.value},
             )
 
         # Resume if paused so cancellation can proceed
@@ -308,6 +274,8 @@ class TaskSession:
             SignalMessage(
                 task_id=self.task_id,
                 mission_id=self.mission_id,
+                setup_id=self.setup_id,
+                setup_version_id=self.setup_version_id,
                 action=SignalType.ACK_CANCEL,
                 status=self.status,
             ).model_dump(),
@@ -316,11 +284,7 @@ class TaskSession:
     async def _handle_pause(self) -> None:
         """Pause task execution."""
         if not self._paused.is_set():
-            logger.info(
-                "Pausing task: '%s'",
-                self.task_id,
-                extra={"task_id": self.task_id},
-            )
+            logger.info("Task paused", extra=self.session_ids)
             self._paused.set()
 
         await self.db.update(
@@ -329,6 +293,8 @@ class TaskSession:
             SignalMessage(
                 task_id=self.task_id,
                 mission_id=self.mission_id,
+                setup_id=self.setup_id,
+                setup_version_id=self.setup_version_id,
                 action=SignalType.ACK_PAUSE,
                 status=self.status,
             ).model_dump(),
@@ -337,11 +303,7 @@ class TaskSession:
     async def _handle_resume(self) -> None:
         """Resume paused task."""
         if self._paused.is_set():
-            logger.info(
-                "Resuming task: '%s'",
-                self.task_id,
-                extra={"task_id": self.task_id},
-            )
+            logger.info("Task resumed", extra=self.session_ids)
             self._paused.clear()
 
         await self.db.update(
@@ -350,6 +312,8 @@ class TaskSession:
             SignalMessage(
                 task_id=self.task_id,
                 mission_id=self.mission_id,
+                setup_id=self.setup_id,
+                setup_version_id=self.setup_version_id,
                 action=SignalType.ACK_RESUME,
                 status=self.status,
             ).model_dump(),
@@ -361,18 +325,16 @@ class TaskSession:
             "tasks",
             self.signal_record_id,  # type: ignore
             SignalMessage(
-                mission_id=self.mission_id,
                 task_id=self.task_id,
+                mission_id=self.mission_id,
+                setup_id=self.setup_id,
+                setup_version_id=self.setup_version_id,
                 status=self.status,
                 action=SignalType.ACK_STATUS,
             ).model_dump(),
         )
 
-        logger.debug(
-            "Status report sent for task: '%s'",
-            self.task_id,
-            extra={"task_id": self.task_id},
-        )
+        logger.debug("Status report sent", extra=self.session_ids)
 
     async def cleanup(self) -> None:
         """Clean up task session resources.

digitalkin/grpc_servers/module_server.py

@@ -1,22 +1,14 @@
 """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,
@@ -24,13 +16,17 @@ from digitalkin.models.grpc_servers.models import (
     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.
@@ -57,6 +53,12 @@ class ModuleServer(BaseServer):
         self.server_config = server_config
         self.client_config = client_config
         self.module_servicer: ModuleServicer | None = None
+        self.registry: RegistryStrategy | None = None
+
+        self._registry_client_config: ClientConfig | None = None
+        if self.server_config.registry_address:
+            self._registry_client_config = self._build_registry_client_config()
+        self._prepare_registry_config()
 
     def _register_servicers(self) -> None:
         """Register the module servicer with the gRPC server.
@@ -77,17 +79,53 @@ 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._registry_client_config:
+            return
+
+        self.module_class.services_config_params["registry"] = {"client_config": self._registry_client_config}
+
+    def _build_registry_client_config(self) -> ClientConfig:
+        """Build ClientConfig for registry from server_config.registry_address.
+
+        Returns:
+            ClientConfig configured for registry connection.
+        """
+        host, port = self.server_config.registry_address.rsplit(":", 1)
+        return ClientConfig(
+            host=host,
+            port=int(port),
+            mode=self.server_config.mode,
+            security=self.client_config.security if self.client_config else SecurityMode.INSECURE,
+            credentials=self.client_config.credentials if self.client_config else None,
+        )
+
+    def _init_registry(self) -> None:
+        """Initialize server-level registry client for registration.
+
+        Note: services_config_params["registry"] is already set in _prepare_registry_config()
+        which runs in __init__(). This method only creates the server-level client instance.
+        """
+        if not self._registry_client_config:
+            return
+
+        self.registry = GrpcRegistry("", "", "", self._registry_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,12 +135,12 @@ 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})
@@ -110,158 +148,61 @@ class ModuleServer(BaseServer):
             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:
-            try:
-                self._deregister_from_registry()
-            except ServerError:
-                logger.exception("Failed to deregister from registry")
+        """Stop the module server.
 
+        Modules become inactive when they stop sending heartbeats
+        """
         super().stop(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"],
-            )
-
-            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,
+        """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__},
             )
-
-            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
+
+        logger.info(
+            "Attempting to register module with registry",
+            extra={
+                "module_id": module_id,
+                "address": self.server_config.host,
+                "port": self.server_config.port,
+                "version": version,
+                "registry_address": self.server_config.registry_address,
+            },
         )
 
-        # 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=self.server_config.host,
+            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": self.server_config.host,
+                    "port": self.server_config.port,
+                    "registry_address": self.server_config.registry_address,
+                },
             )
-        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,
+                    "registry_address": self.server_config.registry_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