digitalkin 0.3.0rc2__py3-none-any.whl → 0.3.1__py3-none-any.whl

digitalkin/core/task_manager/remote_task_manager.py

@@ -0,0 +1,87 @@
+"""Remote task manager for distributed execution."""
+
+import datetime
+from collections.abc import Coroutine
+from typing import Any
+
+from digitalkin.core.task_manager.base_task_manager import BaseTaskManager
+from digitalkin.logger import logger
+from digitalkin.modules._base_module import BaseModule
+
+
+class RemoteTaskManager(BaseTaskManager):
+    """Task manager for distributed/remote execution.
+
+    Only manages task metadata and signals - actual execution happens in remote workers.
+    Suitable for horizontally scaled deployments with Taskiq/Celery workers.
+    """
+
+    async def create_task(
+        self,
+        task_id: str,
+        mission_id: str,
+        module: BaseModule,
+        coro: Coroutine[Any, Any, None],
+        heartbeat_interval: datetime.timedelta = datetime.timedelta(seconds=2),
+        connection_timeout: datetime.timedelta = datetime.timedelta(seconds=5),
+    ) -> None:
+        """Register task for remote execution (metadata only).
+
+        Creates TaskSession for signal handling and monitoring, but doesn't execute the coroutine.
+        The coroutine will be recreated and executed by a remote worker.
+
+        Args:
+            task_id: Unique identifier for the task
+            mission_id: Mission identifier
+            module: Module instance for metadata (not executed here)
+            coro: Coroutine (will be closed - execution happens in worker)
+            heartbeat_interval: Interval between heartbeats
+            connection_timeout: Connection timeout for SurrealDB
+
+        Raises:
+            ValueError: If task_id duplicated
+            RuntimeError: If task overload
+        """
+        # Validation
+        await self._validate_task_creation(task_id, mission_id, coro)
+
+        logger.info(
+            "Registering remote task: '%s'",
+            task_id,
+            extra={
+                "mission_id": mission_id,
+                "task_id": task_id,
+                "heartbeat_interval": heartbeat_interval,
+                "connection_timeout": connection_timeout,
+            },
+        )
+
+        try:
+            # Create session for metadata and signal handling
+            _channel, _session = await self._create_session(
+                task_id, mission_id, module, heartbeat_interval, connection_timeout
+            )
+
+            # Close coroutine - worker will recreate and execute it
+            coro.close()
+
+            logger.info(
+                "Remote task registered: '%s'",
+                task_id,
+                extra={
+                    "mission_id": mission_id,
+                    "task_id": task_id,
+                    "total_sessions": len(self.tasks_sessions),
+                },
+            )
+
+        except Exception as e:
+            logger.error(
+                "Failed to register remote task: '%s'",
+                task_id,
+                extra={"mission_id": mission_id, "task_id": task_id, "error": str(e)},
+                exc_info=True,
+            )
+            # Cleanup on failure
+            await self._cleanup_task(task_id, mission_id=mission_id)
+            raise

digitalkin/core/task_manager/surrealdb_repository.py

@@ -1,5 +1,6 @@
 """SurrealDB connection management."""
 
+import asyncio
 import datetime
 import os
 from collections.abc import AsyncGenerator
@@ -38,6 +39,7 @@ class SurrealDBConnection(Generic[TSurreal]):
 
     db: TSurreal
     timeout: datetime.timedelta
+    _live_queries: set[UUID]  # Track active live queries for cleanup
 
     @staticmethod
     def _valid_id(raw_id: str, table_name: str) -> RecordID:
@@ -82,6 +84,7 @@ class SurrealDBConnection(Generic[TSurreal]):
         self.password = os.getenv("SURREALDB_PASSWORD", "root")
         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
 
     async def init_surreal_instance(self) -> None:
         """Init a SurrealDB connection instance."""
@@ -92,7 +95,37 @@ class SurrealDBConnection(Generic[TSurreal]):
         logger.debug("Successfully connected to SurrealDB")
 
     async def close(self) -> None:
-        """Close the SurrealDB connection if it exists."""
+        """Close the SurrealDB connection if it exists.
+
+        This will also kill all active live queries to prevent memory leaks.
+        """
+        # 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))
+            live_query_ids = list(self._live_queries)
+
+            # Kill all queries concurrently, capturing any exceptions
+            results = await asyncio.gather(
+                *[self.db.kill(live_id) for live_id in live_query_ids], return_exceptions=True
+            )
+
+            # Process results and track failures
+            failed_queries = []
+            for live_id, result in zip(live_query_ids, results):
+                if isinstance(result, (ConnectionError, TimeoutError, Exception)):
+                    failed_queries.append((live_id, str(result)))
+                else:
+                    self._live_queries.discard(live_id)
+
+            # Log aggregated failures once instead of per-query
+            if failed_queries:
+                logger.warning(
+                    "Failed to kill %d live queries: %s",
+                    len(failed_queries),
+                    failed_queries[:5],  # Only log first 5 to avoid log spam
+                    extra={"total_failed": len(failed_queries)},
+                )
+
         logger.debug("Closing SurrealDB connection")
         await self.db.close()
 
@@ -208,20 +241,26 @@ class SurrealDBConnection(Generic[TSurreal]):
     ) -> tuple[UUID, AsyncGenerator[dict[str, Any], None]]:
         """Create and subscribe to a live SurrealQL query.
 
+        The live query ID is tracked to ensure proper cleanup on connection close.
+
         Args:
             table_name: Name of the table to insert into
 
         Returns:
-            List[Dict[str, Any]]: Query results
+            tuple[UUID, AsyncGenerator]: Live query ID and subscription generator
         """
         live_id = await self.db.live(table_name, diff=False)
+        self._live_queries.add(live_id)  # Track for cleanup
+        logger.debug("Started live query %s for table %s (total: %d)", live_id, table_name, len(self._live_queries))
         return live_id, await self.db.subscribe_live(live_id)
 
     async def stop_live(self, live_id: UUID) -> None:
         """Kill a live SurrealQL query.
 
         Args:
-            live_id: record ID to watch for
+            live_id: Live query ID to kill
         """
-        logger.debug("KILL Subscribe live for: %s", live_id)
+        logger.debug("Killing live query: %s", live_id)
         await self.db.kill(live_id)
+        self._live_queries.discard(live_id)  # Remove from tracker
+        logger.debug("Stopped live query %s (remaining: %d)", live_id, len(self._live_queries))

digitalkin/core/task_manager/task_executor.py

@@ -0,0 +1,249 @@
+"""Task executor for running tasks with full lifecycle management."""
+
+import asyncio
+import datetime
+from collections.abc import Coroutine
+from typing import Any
+
+from digitalkin.core.task_manager.surrealdb_repository import SurrealDBConnection
+from digitalkin.core.task_manager.task_session import TaskSession
+from digitalkin.logger import logger
+from digitalkin.models.core.task_monitor import (
+    CancellationReason,
+    SignalMessage,
+    SignalType,
+    TaskStatus,
+)
+
+
+class TaskExecutor:
+    """Executes tasks with the supervisor pattern (main + heartbeat + signal listener).
+
+    Pure execution logic - no task registry or orchestration.
+    Used by workers to run distributed tasks or by TaskManager for local execution.
+    """
+
+    @staticmethod
+    async def execute_task(  # noqa: C901, PLR0915
+        task_id: str,
+        mission_id: str,
+        coro: Coroutine[Any, Any, None],
+        session: TaskSession,
+        channel: SurrealDBConnection,
+    ) -> asyncio.Task[None]:
+        """Execute a task using the supervisor pattern.
+
+        Runs three concurrent sub-tasks:
+        - Main coroutine (the actual work)
+        - Heartbeat generator (sends heartbeats to SurrealDB)
+        - Signal listener (watches for stop/pause/resume signals)
+
+        The first task to complete determines the outcome.
+
+        Args:
+            task_id: Unique identifier for the task
+            mission_id: Mission identifier for the task
+            coro: The coroutine to execute (module.start(...))
+            session: TaskSession for state management
+            channel: SurrealDB connection for signals
+
+        Returns:
+            asyncio.Task: The supervisor task managing the lifecycle
+        """
+
+        async def signal_wrapper() -> None:
+            """Create initial signal record and listen for signals."""
+            try:
+                await channel.create(
+                    "tasks",
+                    SignalMessage(
+                        task_id=task_id,
+                        mission_id=mission_id,
+                        status=session.status,
+                        action=SignalType.START,
+                    ).model_dump(),
+                )
+                await session.listen_signals()
+            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(),
+                )
+                logger.info("Signal listener ended", extra={"mission_id": mission_id, "task_id": task_id})
+
+        async def heartbeat_wrapper() -> None:
+            """Generate heartbeats for task health monitoring."""
+            try:
+                await session.generate_heartbeats()
+            except asyncio.CancelledError:
+                logger.debug("Heartbeat cancelled", extra={"mission_id": mission_id, "task_id": task_id})
+            finally:
+                logger.info("Heartbeat task ended", extra={"mission_id": mission_id, "task_id": task_id})
+
+        async def supervisor() -> None:  # noqa: C901, PLR0912, PLR0915
+            """Supervise the three concurrent tasks and handle outcomes.
+
+            Raises:
+                RuntimeError: If the heartbeat task stops unexpectedly.
+                asyncio.CancelledError: If the supervisor task is cancelled.
+            """
+            session.started_at = datetime.datetime.now(datetime.timezone.utc)
+            session.status = TaskStatus.RUNNING
+
+            # Create tasks with proper exception handling
+            main_task = None
+            hb_task = None
+            sig_task = None
+            cleanup_reason = CancellationReason.UNKNOWN
+
+            try:
+                main_task = asyncio.create_task(coro, name=f"{task_id}_main")
+                hb_task = asyncio.create_task(heartbeat_wrapper(), name=f"{task_id}_heartbeat")
+                sig_task = asyncio.create_task(signal_wrapper(), name=f"{task_id}_listener")
+                done, pending = await asyncio.wait(
+                    [main_task, sig_task, hb_task],
+                    return_when=asyncio.FIRST_COMPLETED,
+                )
+
+                # Determine cleanup reason based on which task completed first
+                completed = next(iter(done))
+
+                if completed is main_task:
+                    # Main task finished - cleanup is due to success
+                    cleanup_reason = CancellationReason.SUCCESS_CLEANUP
+                elif completed is sig_task or (completed is hb_task and sig_task.done()):
+                    # Signal task finished - external cancellation
+                    cleanup_reason = CancellationReason.SIGNAL
+                elif completed is hb_task:
+                    # Heartbeat stopped - failure cleanup
+                    cleanup_reason = CancellationReason.FAILURE_CLEANUP
+
+                # Cancel pending tasks with proper reason logging
+                if pending:
+                    pending_names = [t.get_name() for t in pending]
+                    logger.debug(
+                        "Cancelling pending tasks: %s, reason: %s",
+                        pending_names,
+                        cleanup_reason.value,
+                        extra={
+                            "mission_id": mission_id,
+                            "task_id": task_id,
+                            "pending_tasks": pending_names,
+                            "cancellation_reason": cleanup_reason.value,
+                        },
+                    )
+                    for t in pending:
+                        t.cancel()
+
+                # Propagate exception/result from the finished task
+                await completed
+
+                # Determine final status based on which task completed
+                if completed is main_task:
+                    session.status = TaskStatus.COMPLETED
+                    logger.info(
+                        "Main task completed successfully",
+                        extra={"mission_id": mission_id, "task_id": task_id},
+                    )
+                elif completed is sig_task or (completed is hb_task and sig_task.done()):
+                    session.status = TaskStatus.CANCELLED
+                    session.cancellation_reason = CancellationReason.SIGNAL
+                    logger.info(
+                        "Task cancelled via external signal",
+                        extra={
+                            "mission_id": mission_id,
+                            "task_id": task_id,
+                            "cancellation_reason": CancellationReason.SIGNAL.value,
+                        },
+                    )
+                elif completed is hb_task:
+                    session.status = TaskStatus.FAILED
+                    session.cancellation_reason = CancellationReason.HEARTBEAT_FAILURE
+                    logger.error(
+                        "Heartbeat stopped unexpectedly for task: '%s'",
+                        task_id,
+                        extra={
+                            "mission_id": mission_id,
+                            "task_id": task_id,
+                            "cancellation_reason": CancellationReason.HEARTBEAT_FAILURE.value,
+                        },
+                    )
+                    msg = f"Heartbeat stopped for {task_id}"
+                    raise RuntimeError(msg)  # noqa: TRY301
+
+            except asyncio.CancelledError:
+                session.status = TaskStatus.CANCELLED
+                # Only set reason if not already set (preserve original reason)
+                logger.info(
+                    "Task cancelled externally: '%s', reason: %s",
+                    task_id,
+                    session.cancellation_reason.value,
+                    extra={
+                        "mission_id": mission_id,
+                        "task_id": task_id,
+                        "cancellation_reason": session.cancellation_reason.value,
+                    },
+                )
+                cleanup_reason = CancellationReason.FAILURE_CLEANUP
+                raise
+            except Exception:
+                session.status = TaskStatus.FAILED
+                cleanup_reason = CancellationReason.FAILURE_CLEANUP
+                logger.exception(
+                    "Task failed with exception: '%s'",
+                    task_id,
+                    extra={"mission_id": mission_id, "task_id": task_id},
+                )
+                raise
+            finally:
+                session.completed_at = datetime.datetime.now(datetime.timezone.utc)
+                # Ensure all tasks are cleaned up with proper reason
+                tasks_to_cleanup = [t for t in [main_task, hb_task, sig_task] if t is not None and not t.done()]
+                if tasks_to_cleanup:
+                    cleanup_names = [t.get_name() for t in tasks_to_cleanup]
+                    logger.debug(
+                        "Final cleanup of %d remaining tasks: %s, reason: %s",
+                        len(tasks_to_cleanup),
+                        cleanup_names,
+                        cleanup_reason.value,
+                        extra={
+                            "mission_id": mission_id,
+                            "task_id": task_id,
+                            "cleanup_count": len(tasks_to_cleanup),
+                            "cleanup_tasks": cleanup_names,
+                            "cancellation_reason": cleanup_reason.value,
+                        },
+                    )
+                    for t in tasks_to_cleanup:
+                        t.cancel()
+                    await asyncio.gather(*tasks_to_cleanup, return_exceptions=True)
+
+                duration = (
+                    (session.completed_at - session.started_at).total_seconds()
+                    if session.started_at and session.completed_at
+                    else None
+                )
+                logger.info(
+                    "Task execution completed: '%s', status: %s, reason: %s, duration: %.2fs",
+                    task_id,
+                    session.status.value,
+                    session.cancellation_reason.value if session.status == TaskStatus.CANCELLED else "n/a",
+                    duration or 0,
+                    extra={
+                        "mission_id": mission_id,
+                        "task_id": task_id,
+                        "status": session.status.value,
+                        "cancellation_reason": session.cancellation_reason.value,
+                        "duration": duration,
+                    },
+                )
+
+        # Return the supervisor task to be awaited by caller
+        return asyncio.create_task(supervisor(), name=f"{task_id}_supervisor")

digitalkin/core/task_manager/task_session.py

@@ -6,7 +6,13 @@ 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 HeartbeatMessage, SignalMessage, SignalType, TaskStatus
+from digitalkin.models.core.task_monitor import (
+    CancellationReason,
+    HeartbeatMessage,
+    SignalMessage,
+    SignalType,
+    TaskStatus,
+)
 from digitalkin.modules._base_module import BaseModule
 
 
@@ -31,6 +37,7 @@ class TaskSession:
     completed_at: datetime.datetime | None
 
     is_cancelled: asyncio.Event
+    cancellation_reason: CancellationReason
     _paused: asyncio.Event
     _heartbeat_interval: datetime.timedelta
     _last_heartbeat: datetime.datetime
@@ -58,6 +65,7 @@ class TaskSession:
         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
@@ -71,6 +79,7 @@ class TaskSession:
         self.heartbeat_record_id = None
 
         self.is_cancelled = asyncio.Event()
+        self.cancellation_reason = CancellationReason.UNKNOWN
         self._paused = asyncio.Event()
         self._heartbeat_interval = heartbeat_interval
 
@@ -152,17 +161,26 @@ class TaskSession:
 
     async def generate_heartbeats(self) -> None:
         """Periodic heartbeat generator with cancellation support."""
-        logger.debug("Heartbeat started")
+        logger.debug(
+            "Heartbeat generator started for task: '%s'",
+            self.task_id,
+            extra={"task_id": self.task_id, "mission_id": self.mission_id},
+        )
         while not self.cancelled:
-            logger.debug(f"Heartbeat tick for task: '{self.task_id}' | {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},
+            )
             success = await self.send_heartbeat()
             if not success:
                 logger.error(
                     "Heartbeat failed, cancelling task: '%s'",
                     self.task_id,
-                    extra={"task_id": self.task_id},
+                    extra={"task_id": self.task_id, "mission_id": self.mission_id},
                 )
-                await self._handle_cancel()
+                await self._handle_cancel(CancellationReason.HEARTBEAT_FAILURE)
                 break
             await asyncio.sleep(self._heartbeat_interval.total_seconds())
 
@@ -201,7 +219,7 @@ class TaskSession:
                     continue
 
                 if signal["action"] == "cancel":
-                    await self._handle_cancel()
+                    await self._handle_cancel(CancellationReason.SIGNAL)
                 elif signal["action"] == "pause":
                     await self._handle_pause()
                 elif signal["action"] == "resume":
@@ -231,26 +249,55 @@ class TaskSession:
                 extra={"task_id": self.task_id},
             )
 
-    async def _handle_cancel(self) -> None:
-        """Idempotent cancellation with acknowledgment."""
-        logger.debug("Handle cancel called")
+    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 signal ignored - task already cancelled: '%s'",
+                "Cancel ignored - task already cancelled: '%s' (existing reason: %s, new reason: %s)",
                 self.task_id,
-                extra={"task_id": self.task_id},
+                self.cancellation_reason.value,
+                reason.value,
+                extra={
+                    "task_id": self.task_id,
+                    "mission_id": self.mission_id,
+                    "existing_reason": self.cancellation_reason.value,
+                    "new_reason": reason.value,
+                },
             )
             return
 
-        logger.info(
-            "Cancelling task: '%s'",
-            self.task_id,
-            extra={"task_id": self.task_id},
-        )
-
+        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): '%s', reason: %s",
+                self.task_id,
+                reason.value,
+                extra={
+                    "task_id": self.task_id,
+                    "mission_id": self.mission_id,
+                    "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,
+                },
+            )
+
         # Resume if paused so cancellation can proceed
         if self._paused.is_set():
             self._paused.set()
@@ -326,3 +373,34 @@ class TaskSession:
             self.task_id,
             extra={"task_id": self.task_id},
         )
+
+    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/module_server.py

@@ -4,11 +4,11 @@ import uuid
 from pathlib import Path
 
 import grpc
-from digitalkin_proto.digitalkin.module.v2 import (
+from digitalkin_proto.agentic_mesh_protocol.module.v1 import (
     module_service_pb2,
     module_service_pb2_grpc,
 )
-from digitalkin_proto.digitalkin.module_registry.v2 import (
+from digitalkin_proto.agentic_mesh_protocol.module_registry.v1 import (
     metadata_pb2,
     module_registry_service_pb2_grpc,
     registration_pb2,