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

digitalkin/core/task_manager/surrealdb_repository.py

@@ -4,7 +4,7 @@ import asyncio
 import datetime
 import os
 from collections.abc import AsyncGenerator
-from typing import Any, Generic, TypeVar
+from typing import Any, Generic, TypeVar, cast
 from uuid import UUID
 
 from surrealdb import AsyncHttpSurrealConnection, AsyncSurreal, AsyncWsSurrealConnection, RecordID
@@ -40,6 +40,7 @@ class SurrealDBConnection(Generic[TSurreal]):
     db: TSurreal
     timeout: datetime.timedelta
     _live_queries: set[UUID]  # Track active live queries for cleanup
+    _closed: bool  # Flag to prevent operations on closed connection
 
     @staticmethod
     def _valid_id(raw_id: str, table_name: str) -> RecordID:
@@ -85,13 +86,14 @@ class SurrealDBConnection(Generic[TSurreal]):
         self.namespace = os.getenv("SURREALDB_NAMESPACE", "test")
         self.database = database or os.getenv("SURREALDB_DATABASE", "task_manager")
         self._live_queries = set()  # Initialize live queries tracker
+        self._closed = False
 
     async def init_surreal_instance(self) -> None:
         """Init a SurrealDB connection instance."""
         logger.debug("Connecting to SurrealDB at %s", self.url)
         self.db = AsyncSurreal(self.url)  # type: ignore
         await self.db.signin({"username": self.username, "password": self.password})
-        await self.db.use(self.namespace, self.database)
+        await self.db.use(self.namespace, self.database)  # type: ignore[arg-type]
         logger.debug("Successfully connected to SurrealDB")
 
     async def close(self) -> None:
@@ -99,6 +101,7 @@ class SurrealDBConnection(Generic[TSurreal]):
 
         This will also kill all active live queries to prevent memory leaks.
         """
+        self._closed = True
         # Kill all tracked live queries before closing connection
         if self._live_queries:
             logger.debug("Killing %d active live queries before closing", len(self._live_queries))
@@ -112,7 +115,7 @@ class SurrealDBConnection(Generic[TSurreal]):
             # Process results and track failures
             failed_queries = []
             for live_id, result in zip(live_query_ids, results):
-                if isinstance(result, (ConnectionError, TimeoutError, Exception)):
+                if isinstance(result, ConnectionError | TimeoutError | Exception):
                     failed_queries.append((live_id, str(result)))
                 else:
                     self._live_queries.discard(live_id)
@@ -142,11 +145,27 @@ class SurrealDBConnection(Generic[TSurreal]):
 
         Returns:
             Dict[str, Any]: The created record as returned by the database
+
+        Raises:
+            RuntimeError: If the database returns an error response
         """
         logger.debug("Creating record in %s with data: %s", table_name, data)
         result = await self.db.create(table_name, data)
         logger.debug("create result: %s", result)
-        return result
+
+        # Check for error response from SurrealDB
+        if isinstance(result, dict) and "code" in result:
+            error_msg = result.get("message", result.get("information", "Unknown error"))
+            logger.error(
+                "SurrealDB create failed: %s (code: %s)",
+                error_msg,
+                result.get("code"),
+                extra={"table": table_name, "error": result},
+            )
+            msg = f"SurrealDB create failed in '{table_name}': {error_msg}"
+            raise RuntimeError(msg)
+
+        return cast("list[dict[str, Any]] | dict[str, Any]", result)
 
     async def merge(
         self,
@@ -170,7 +189,7 @@ class SurrealDBConnection(Generic[TSurreal]):
         logger.debug("Updating record in %s with data: %s", record_id, data)
         result = await self.db.merge(record_id, data)
         logger.debug("update result: %s", result)
-        return result
+        return cast("list[dict[str, Any]] | dict[str, Any]", result)
 
     async def update(
         self,
@@ -194,7 +213,7 @@ class SurrealDBConnection(Generic[TSurreal]):
         logger.debug("Updating record in %s with data: %s", record_id, data)
         result = await self.db.update(record_id, data)
         logger.debug("update result: %s", result)
-        return result
+        return cast("list[dict[str, Any]] | dict[str, Any]", result)
 
     async def execute_query(self, query: str, params: dict[str, Any] | None = None) -> list[dict[str, Any]]:
         """Execute a custom SurrealQL query.
@@ -209,7 +228,7 @@ class SurrealDBConnection(Generic[TSurreal]):
         logger.debug("execute_query: %s with params: %s", query, params)
         result = await self.db.query(query, params or {})
         logger.debug("execute_query result: %s", result)
-        return [result] if isinstance(result, dict) else result
+        return cast("list[dict[str, Any]]", [result] if isinstance(result, dict) else result)
 
     async def select_by_task_id(self, table: str, value: str) -> dict[str, Any]:
         """Fetch a record from a table by a unique field.
@@ -260,6 +279,9 @@ class SurrealDBConnection(Generic[TSurreal]):
         Args:
             live_id: Live query ID to kill
         """
+        if self._closed:
+            self._live_queries.discard(live_id)
+            return
         logger.debug("Killing live query: %s", live_id)
         await self.db.kill(live_id)
         self._live_queries.discard(live_id)  # Remove from tracker

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
@@ -54,28 +55,53 @@ class TaskExecutor:
         async def signal_wrapper() -> None:
             """Create initial signal record and listen for signals."""
             try:
-                await channel.create(
+                # Create task record and capture the record ID directly
+                # This avoids a race condition where SELECT might run before CREATE completes
+                result = 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.START,
                     ).model_dump(),
                 )
-                await session.listen_signals()
+                # Store the record ID in session - required before starting live query
+                if isinstance(result, dict) and "id" in result:
+                    session.signal_record_id = result["id"]
+                    logger.debug(
+                        "Task signal record created",
+                        extra={"mission_id": mission_id, "task_id": task_id, "record_id": result["id"]},
+                    )
+                    # Only start listening if we have a valid record ID
+                    await session.listen_signals()
+                else:
+                    # Create failed - wait for cancellation instead of listening
+                    logger.error(
+                        "Failed to get record ID from task creation, waiting for cancellation",
+                        extra={"mission_id": mission_id, "task_id": task_id, "result": result},
+                    )
+                    await session.is_cancelled.wait()
             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 +151,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 +180,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 +226,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,
@@ -81,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
@@ -99,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.
 
@@ -108,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),
         )
 
@@ -120,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
 
@@ -147,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())
@@ -187,32 +222,32 @@ 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
         """Enhanced signal listener with comprehensive handling.
 
         Raises:
-            CancelledError: Asyncio when task cancelling
+            CancelledError: If task is cancelled during signal listening.
         """
-        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)
+
+        # signal_record_id must be set by TaskExecutor before calling this method.
+        # If not set, we cannot filter signals correctly - abort early.
         if self.signal_record_id is None:
-            self.signal_record_id = (await self.db.select_by_task_id("tasks", self.task_id)).get("id")
+            logger.error(
+                "signal_record_id not set - cannot start signal listener without valid record ID",
+                extra=self.session_ids,
+            )
+            return
 
         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)
-                if self.cancelled:
+                logger.debug("Signal received", extra={**self.session_ids, "signal": signal})
+                # Check both cancelled and stream_closed to ensure clean shutdown
+                if self.cancelled or self.stream_closed:
                     break
 
                 if signal is None or signal["id"] == self.signal_record_id or "payload" not in signal:
@@ -228,26 +263,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.
@@ -257,13 +284,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 +300,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,19 +319,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(
@@ -329,6 +339,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 +349,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 +358,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,28 +371,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():
@@ -390,6 +410,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()