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

digitalkin/core/task_manager/task_executor.py

@@ -1,6 +1,7 @@
 """Task executor for running tasks with full lifecycle management."""
 
 import asyncio
+import contextlib
 import datetime
 from collections.abc import Coroutine
 from typing import Any
@@ -59,6 +60,8 @@ class TaskExecutor:
                     SignalMessage(
                         task_id=task_id,
                         mission_id=mission_id,
+                        setup_id=session.setup_id,
+                        setup_version_id=session.setup_version_id,
                         status=session.status,
                         action=SignalType.START,
                     ).model_dump(),
@@ -67,15 +70,21 @@ class TaskExecutor:
             except asyncio.CancelledError:
                 logger.debug("Signal listener cancelled", extra={"mission_id": mission_id, "task_id": task_id})
             finally:
-                await channel.create(
-                    "tasks",
-                    SignalMessage(
-                        task_id=task_id,
-                        mission_id=mission_id,
-                        status=session.status,
-                        action=SignalType.STOP,
-                    ).model_dump(),
-                )
+                with contextlib.suppress(Exception):  # Connection may already be closed
+                    await channel.create(
+                        "tasks",
+                        SignalMessage(
+                            task_id=task_id,
+                            mission_id=mission_id,
+                            setup_id=session.setup_id,
+                            setup_version_id=session.setup_version_id,
+                            status=session.status,
+                            action=SignalType.STOP,
+                            cancellation_reason=session.cancellation_reason,
+                            error_message=session._last_exception,  # noqa: SLF001
+                            exception_traceback=session._last_traceback,  # noqa: SLF001
+                        ).model_dump(),
+                    )
                 logger.info("Signal listener ended", extra={"mission_id": mission_id, "task_id": task_id})
 
         async def heartbeat_wrapper() -> None:
@@ -125,8 +134,14 @@ class TaskExecutor:
                     # Heartbeat stopped - failure cleanup
                     cleanup_reason = CancellationReason.FAILURE_CLEANUP
 
+                # Signal stream to close FIRST before any cleanup
+                session.close_stream()
+
                 # Cancel pending tasks with proper reason logging
                 if pending:
+                    # Give stream time to see the signal and exit gracefully
+                    await asyncio.sleep(0.01)  # Allow one event loop cycle
+
                     pending_names = [t.get_name() for t in pending]
                     logger.debug(
                         "Cancelling pending tasks: %s, reason: %s",
@@ -148,6 +163,7 @@ class TaskExecutor:
                 # Determine final status based on which task completed
                 if completed is main_task:
                     session.status = TaskStatus.COMPLETED
+                    session.cancellation_reason = CancellationReason.COMPLETED
                     logger.info(
                         "Main task completed successfully",
                         extra={"mission_id": mission_id, "task_id": task_id},
@@ -193,9 +209,10 @@ class TaskExecutor:
                 )
                 cleanup_reason = CancellationReason.FAILURE_CLEANUP
                 raise
-            except Exception:
+            except Exception as e:
                 session.status = TaskStatus.FAILED
                 cleanup_reason = CancellationReason.FAILURE_CLEANUP
+                session.record_exception(e)
                 logger.exception(
                     "Task failed with exception: '%s'",
                     task_id,

digitalkin/core/task_manager/task_session.py

@@ -1,7 +1,9 @@
 """Task session easing task lifecycle management."""
 
 import asyncio
+import contextlib
 import datetime
+import traceback
 from collections.abc import AsyncGenerator
 
 from digitalkin.core.task_manager.surrealdb_repository import SurrealDBConnection
@@ -39,9 +41,17 @@ class TaskSession:
     is_cancelled: asyncio.Event
     cancellation_reason: CancellationReason
     _paused: asyncio.Event
+    _stream_closed: asyncio.Event
     _heartbeat_interval: datetime.timedelta
     _last_heartbeat: datetime.datetime
 
+    # Exception tracking for enhanced DB logging
+    _last_exception: str | None
+    _last_traceback: str | None
+
+    # Cleanup guard for idempotent cleanup
+    _cleanup_done: bool
+
     def __init__(
         self,
         task_id: str,
@@ -49,14 +59,24 @@ class TaskSession:
         db: SurrealDBConnection,
         module: BaseModule,
         heartbeat_interval: datetime.timedelta = datetime.timedelta(seconds=2),
+        queue_maxsize: int = 1000,
     ) -> None:
-        """Initialize Task Session."""
+        """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=1000)
+        self.queue: asyncio.Queue = asyncio.Queue(maxsize=queue_maxsize)
 
         self.task_id = task_id
         self.mission_id = mission_id
@@ -71,12 +91,23 @@ class TaskSession:
         self.is_cancelled = asyncio.Event()
         self.cancellation_reason = CancellationReason.UNKNOWN
         self._paused = asyncio.Event()
+        self._stream_closed = asyncio.Event()
         self._heartbeat_interval = heartbeat_interval
 
+        # Exception tracking
+        self._last_exception = None
+        self._last_traceback = None
+
+        # Cleanup guard
+        self._cleanup_done = False
+
         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
@@ -89,6 +120,39 @@ class TaskSession:
         """Task paused status."""
         return self._paused.is_set()
 
+    @property
+    def stream_closed(self) -> bool:
+        """Check if stream termination was signaled."""
+        return self._stream_closed.is_set()
+
+    def close_stream(self) -> None:
+        """Signal that the stream should terminate."""
+        self._stream_closed.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()
+
+    def record_exception(self, exc: Exception) -> None:
+        """Record exception details for DB logging.
+
+        Args:
+            exc: The exception that caused the task to fail.
+        """
+        self._last_exception = str(exc)
+        self._last_traceback = traceback.format_exc()
+
     async def send_heartbeat(self) -> bool:
         """Rate-limited heartbeat with connection resilience.
 
@@ -98,6 +162,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),
         )
 
@@ -110,23 +176,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
 
@@ -137,39 +197,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())
@@ -177,11 +222,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
@@ -190,18 +231,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
 
@@ -218,26 +255,18 @@ 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},
-            )
+            with contextlib.suppress(Exception):  # Connection may already be closed
+                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.
@@ -247,13 +276,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,
                 },
@@ -267,25 +292,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
@@ -298,19 +311,18 @@ 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,
+                cancellation_reason=reason,
             ).model_dump(),
         )
 
     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(
@@ -319,6 +331,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(),
@@ -327,11 +341,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(
@@ -340,6 +350,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(),
@@ -351,28 +363,38 @@ 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.
 
+        This method is idempotent - safe to call multiple times.
+        Second and subsequent calls are no-ops.
+
         This includes:
         - Clearing queue to free memory
+        - Cleaning up module context services
         - Stopping module
         - Closing database connection
         - Clearing module reference
         """
+        if self._cleanup_done:
+            logger.debug(
+                "Cleanup already done, skipping",
+                extra={"task_id": self.task_id, "mission_id": self.mission_id},
+            )
+            return
+        self._cleanup_done = True
+
         # Clear queue to free memory
         try:
             while not self.queue.empty():
@@ -380,6 +402,16 @@ class TaskSession:
         except asyncio.QueueEmpty:
             pass
 
+        # Clean up module context services (e.g., gRPC channel pool)
+        if self.module is not None and self.module.context is not None:
+            try:
+                await self.module.context.cleanup()
+            except Exception:
+                logger.exception(
+                    "Error cleaning up module context",
+                    extra={"mission_id": self.mission_id, "task_id": self.task_id},
+                )
+
         # Stop module
         try:
             await self.module.stop()