digitalkin 0.2.25rc0__py3-none-any.whl → 0.3.2.dev14__py3-none-any.whl

digitalkin/core/task_manager/task_session.py

@@ -0,0 +1,368 @@
+"""Task session easing task lifecycle management."""
+
+import asyncio
+import datetime
+from collections.abc import AsyncGenerator
+
+from digitalkin.core.task_manager.surrealdb_repository import SurrealDBConnection
+from digitalkin.logger import logger
+from digitalkin.models.core.task_monitor import (
+    CancellationReason,
+    HeartbeatMessage,
+    SignalMessage,
+    SignalType,
+    TaskStatus,
+)
+from digitalkin.modules._base_module import BaseModule
+
+
+class TaskSession:
+    """Task Session with lifecycle management.
+
+    The Session defined the whole lifecycle of a task as an epheneral context.
+    """
+
+    db: SurrealDBConnection
+    module: BaseModule
+
+    status: TaskStatus
+    signal_queue: AsyncGenerator | None
+
+    task_id: str
+    mission_id: str
+    signal_record_id: str | None
+    heartbeat_record_id: str | None
+
+    started_at: datetime.datetime | None
+    completed_at: datetime.datetime | None
+
+    is_cancelled: asyncio.Event
+    cancellation_reason: CancellationReason
+    _paused: asyncio.Event
+    _heartbeat_interval: datetime.timedelta
+    _last_heartbeat: datetime.datetime
+
+    def __init__(
+        self,
+        task_id: str,
+        mission_id: str,
+        db: SurrealDBConnection,
+        module: BaseModule,
+        heartbeat_interval: datetime.timedelta = datetime.timedelta(seconds=2),
+        queue_maxsize: int = 1000,
+    ) -> None:
+        """Initialize Task Session.
+
+        Args:
+            task_id: Unique task identifier
+            mission_id: Mission identifier
+            db: SurrealDB connection
+            module: Module instance
+            heartbeat_interval: Interval between heartbeats
+            queue_maxsize: Maximum size for the queue (0 = unlimited)
+        """
+        self.db = db
+        self.module = module
+
+        self.status = TaskStatus.PENDING
+        # Bounded queue to prevent unbounded memory growth (max 1000 items)
+        self.queue: asyncio.Queue = asyncio.Queue(maxsize=queue_maxsize)
+
+        self.task_id = task_id
+        self.mission_id = mission_id
+
+        self.heartbeat = None
+        self.started_at = None
+        self.completed_at = None
+
+        self.signal_record_id = None
+        self.heartbeat_record_id = None
+
+        self.is_cancelled = asyncio.Event()
+        self.cancellation_reason = CancellationReason.UNKNOWN
+        self._paused = asyncio.Event()
+        self._heartbeat_interval = heartbeat_interval
+
+        logger.info(
+            "TaskSession initialized",
+            extra={
+                "task_id": task_id,
+                "mission_id": mission_id,
+                "heartbeat_interval": str(heartbeat_interval),
+            },
+        )
+
+    @property
+    def cancelled(self) -> bool:
+        """Task cancellation status."""
+        return self.is_cancelled.is_set()
+
+    @property
+    def paused(self) -> bool:
+        """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.
+
+        Returns:
+            bool: True if heartbeat was successful, False otherwise
+        """
+        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),
+        )
+
+        if self.heartbeat_record_id is None:
+            try:
+                success = await self.db.create("heartbeats", heartbeat.model_dump())
+                if "code" not in success:
+                    self.heartbeat_record_id = success.get("id")  # type: ignore
+                    self._last_heartbeat = heartbeat.timestamp
+                    return True
+            except Exception as e:
+                logger.error(
+                    "Heartbeat exception",
+                    extra={**self.session_ids, "error": str(e)},
+                    exc_info=True,
+                )
+            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",
+                extra={**self.session_ids, "delta": str(heartbeat.timestamp - self._last_heartbeat)},
+            )
+            return True
+
+        try:
+            success = await self.db.merge("heartbeats", self.heartbeat_record_id, heartbeat.model_dump())
+            if "code" not in success:
+                self._last_heartbeat = heartbeat.timestamp
+                return True
+        except Exception as e:
+            logger.error(
+                "Heartbeat exception",
+                extra={**self.session_ids, "error": str(e)},
+                exc_info=True,
+            )
+        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", extra=self.session_ids)
+        while not self.cancelled:
+            logger.debug(
+                "Heartbeat tick",
+                extra={**self.session_ids, "cancelled": self.cancelled},
+            )
+            success = await self.send_heartbeat()
+            if not success:
+                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())
+
+    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", extra=self.session_ids)
+            await self._paused.wait()
+
+    async def listen_signals(self) -> None:  # noqa: C901
+        """Enhanced signal listener with comprehensive handling.
+
+        Raises:
+            CancelledError: Asyncio when task cancelling
+        """
+        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", extra={**self.session_ids, "signal": signal})
+                if self.cancelled:
+                    break
+
+                if signal is None or signal["id"] == self.signal_record_id or "payload" not in signal:
+                    continue
+
+                if signal["action"] == "cancel":
+                    await self._handle_cancel(CancellationReason.SIGNAL)
+                elif signal["action"] == "pause":
+                    await self._handle_pause()
+                elif signal["action"] == "resume":
+                    await self._handle_resume()
+                elif signal["action"] == "status":
+                    await self._handle_status_request()
+
+        except asyncio.CancelledError:
+            logger.debug("Signal listener cancelled", extra=self.session_ids)
+            raise
+        except Exception as e:
+            logger.error(
+                "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", extra=self.session_ids)
+
+    async def _handle_cancel(self, reason: CancellationReason = CancellationReason.UNKNOWN) -> None:
+        """Idempotent cancellation with acknowledgment and reason tracking.
+
+        Args:
+            reason: The reason for cancellation (signal, heartbeat failure, cleanup, etc.)
+        """
+        if self.is_cancelled.is_set():
+            logger.debug(
+                "Cancel ignored - already cancelled",
+                extra={
+                    **self.session_ids,
+                    "existing_reason": self.cancellation_reason.value,
+                    "new_reason": reason.value,
+                },
+            )
+            return
+
+        self.cancellation_reason = reason
+        self.status = TaskStatus.CANCELLED
+        self.is_cancelled.set()
+
+        # Log with appropriate level based on reason
+        if reason in {CancellationReason.SUCCESS_CLEANUP, CancellationReason.FAILURE_CLEANUP}:
+            logger.debug(
+                "Task cancelled (cleanup)",
+                extra={**self.session_ids, "cancellation_reason": reason.value},
+            )
+        else:
+            logger.info(
+                "Task cancelled",
+                extra={**self.session_ids, "cancellation_reason": reason.value},
+            )
+
+        # Resume if paused so cancellation can proceed
+        if self._paused.is_set():
+            self._paused.set()
+
+        await self.db.update(
+            "tasks",
+            self.signal_record_id,  # type: ignore
+            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(),
+        )
+
+    async def _handle_pause(self) -> None:
+        """Pause task execution."""
+        if not self._paused.is_set():
+            logger.info("Task paused", extra=self.session_ids)
+            self._paused.set()
+
+        await self.db.update(
+            "tasks",
+            self.signal_record_id,  # type: ignore
+            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(),
+        )
+
+    async def _handle_resume(self) -> None:
+        """Resume paused task."""
+        if self._paused.is_set():
+            logger.info("Task resumed", extra=self.session_ids)
+            self._paused.clear()
+
+        await self.db.update(
+            "tasks",
+            self.signal_record_id,  # type: ignore
+            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(),
+        )
+
+    async def _handle_status_request(self) -> None:
+        """Send current task status."""
+        await self.db.update(
+            "tasks",
+            self.signal_record_id,  # type: ignore
+            SignalMessage(
+                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", extra=self.session_ids)
+
+    async def cleanup(self) -> None:
+        """Clean up task session resources.
+
+        This includes:
+        - Clearing queue to free memory
+        - Stopping module
+        - Closing database connection
+        - Clearing module reference
+        """
+        # Clear queue to free memory
+        try:
+            while not self.queue.empty():
+                self.queue.get_nowait()
+        except asyncio.QueueEmpty:
+            pass
+
+        # Stop module
+        try:
+            await self.module.stop()
+        except Exception:
+            logger.exception(
+                "Error stopping module during cleanup",
+                extra={"mission_id": self.mission_id, "task_id": self.task_id},
+            )
+
+        # Close DB connection (kills all live queries)
+        await self.db.close()
+
+        # Clear module reference to allow garbage collection
+        self.module = None  # type: ignore

digitalkin/grpc_servers/__init__.py

@@ -1,19 +1 @@
-"""This package contains the gRPC server and client implementations.
-
-```shell
-digitalkin/grpc/
-├── __init__.py
-├── base_server.py        # Base server implementation with common functionality
-├── module_server.py      # Module-specific server implementation
-├── registry_server.py    # Registry-specific server implementation
-├── module_servicer.py    # gRPC servicer for Module service
-├── registry_servicer.py  # gRPC servicer for Registry service
-├── client/               # Client libraries for connecting to servers
-│   ├── __init__.py
-│   ├── module_client.py
-│   └── registry_client.py
-└── utils/                # Utility functions
-    ├── __init__.py
-    └── server_utils.py   # Common server utilities
-```
-"""
+"""This package contains the gRPC server and client implementations."""

digitalkin/grpc_servers/_base_server.py

@@ -17,9 +17,9 @@ from digitalkin.grpc_servers.utils.exceptions import (
     ServerStateError,
     ServicerError,
 )
-from digitalkin.grpc_servers.utils.models import SecurityMode, ServerConfig, ServerMode
-from digitalkin.grpc_servers.utils.types import GrpcServer, ServiceDescriptor, T
 from digitalkin.logger import logger
+from digitalkin.models.grpc_servers.models import SecurityMode, ServerConfig, ServerMode
+from digitalkin.models.grpc_servers.types import GrpcServer, ServiceDescriptor, T
 
 
 class BaseServer(abc.ABC):
@@ -299,7 +299,7 @@ class BaseServer(abc.ABC):
         self._add_reflection()
 
         # Start the server
-        logger.debug("Starting gRPC server on %s", self.config.address)
+        logger.debug("Starting gRPC server on %s", self.config.address, extra={"config": self.config})
         try:
             if self.config.mode == ServerMode.ASYNC:
                 # For async server, use the event loop