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

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,173 @@
+"""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 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:
+            """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
+
+            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,
+                )
+
+                # One task completed -> cancel the others
+                for t in pending:
+                    t.cancel()
+
+                # Propagate exception/result from the finished task
+                completed = next(iter(done))
+                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()):
+                    logger.debug(
+                        f"Task cancelled due to signal {sig_task=}",
+                        extra={"mission_id": mission_id, "task_id": task_id},
+                    )
+                    session.status = TaskStatus.CANCELLED
+                elif completed is hb_task:
+                    session.status = TaskStatus.FAILED
+                    logger.error(
+                        f"Heartbeat stopped for {task_id}",
+                        extra={"mission_id": mission_id, "task_id": task_id},
+                    )
+                    msg = f"Heartbeat stopped for {task_id}"
+                    raise RuntimeError(msg)  # noqa: TRY301
+
+            except asyncio.CancelledError:
+                session.status = TaskStatus.CANCELLED
+                logger.info("Task cancelled", extra={"mission_id": mission_id, "task_id": task_id})
+                raise
+            except Exception:
+                session.status = TaskStatus.FAILED
+                logger.exception("Task failed", 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
+                tasks_to_cleanup = [t for t in [main_task, hb_task, sig_task] if t is not None]
+                for t in tasks_to_cleanup:
+                    if not t.done():
+                        t.cancel()
+                if tasks_to_cleanup:
+                    await asyncio.gather(*tasks_to_cleanup, return_exceptions=True)
+
+                logger.info(
+                    "Task execution completed with status: %s",
+                    session.status,
+                    extra={
+                        "mission_id": mission_id,
+                        "task_id": task_id,
+                        "status": session.status,
+                        "duration": (
+                            (session.completed_at - session.started_at).total_seconds()
+                            if session.started_at and session.completed_at
+                            else None
+                        ),
+                    },
+                )
+
+        # 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

@@ -42,23 +42,14 @@ class TaskSession:
         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)
-        """
+        """Initialize Task Session."""
         self.db = db
         self.module = module
 
         self.status = TaskStatus.PENDING
-        self.queue: asyncio.Queue = asyncio.Queue(maxsize=queue_maxsize)
+        # Bounded queue to prevent unbounded memory growth (max 1000 items)
+        self.queue: asyncio.Queue = asyncio.Queue(maxsize=1000)
 
         self.task_id = task_id
         self.mission_id = mission_id
@@ -326,3 +317,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,

digitalkin/grpc_servers/module_servicer.py

@@ -5,7 +5,7 @@ from collections.abc import AsyncGenerator
 from typing import Any
 
 import grpc
-from digitalkin_proto.digitalkin.module.v2 import (
+from digitalkin_proto.agentic_mesh_protocol.module.v1 import (
     information_pb2,
     lifecycle_pb2,
     module_service_pb2_grpc,
@@ -172,7 +172,8 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         )
         # Process the module input
         # TODO: Check failure of input data format
-        input_data = self.module_class.create_input_model(dict(request.input.items()))
+        input_data = self.module_class.create_input_model(json_format.MessageToDict(request.input))
+
         setup_data_class = self.setup.get_setup(
             setup_dict={
                 "setup_id": request.setup_id,
@@ -225,7 +226,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
                     proto = json_format.ParseDict(message, struct_pb2.Struct(), ignore_unknown_fields=True)
                     yield lifecycle_pb2.StartModuleResponse(success=True, output=proto, job_id=job_id)
         finally:
-            await self.job_manager.tasks[job_id]
+            await self.job_manager.wait_for_completion(job_id)
             await self.job_manager.clean_session(job_id, mission_id=request.mission_id)
 
         logger.info("Job %s finished", job_id)

digitalkin/grpc_servers/registry_server.py

@@ -1,6 +1,6 @@
 """Registry gRPC server implementation for DigitalKin."""
 
-from digitalkin_proto.digitalkin.module_registry.v2 import (
+from digitalkin_proto.agentic_mesh_protocol.module_registry.v1 import (
     module_registry_service_pb2,
     module_registry_service_pb2_grpc,
 )

digitalkin/grpc_servers/registry_servicer.py

@@ -9,7 +9,7 @@ from collections.abc import Iterator
 from enum import Enum
 
 import grpc
-from digitalkin_proto.digitalkin.module_registry.v2 import (
+from digitalkin_proto.agentic_mesh_protocol.module_registry.v1 import (
     discover_pb2,
     metadata_pb2,
     module_registry_service_pb2_grpc,
@@ -344,7 +344,7 @@ class RegistryServicer(module_registry_service_pb2_grpc.ModuleRegistryServiceSer
             return status_pb2.ModuleStatusResponse()
 
         module = self.registered_modules[request.module_id]
-        return status_pb2.ModuleStatusResponse(module_id=module.module_id, status=module.status.name)
+        return status_pb2.ModuleStatusResponse(module_id=module.module_id, status=module.status.value)
 
     def ListModuleStatus(  # noqa: N802
         self,
@@ -379,7 +379,7 @@ class RegistryServicer(module_registry_service_pb2_grpc.ModuleRegistryServiceSer
             list_size = len(self.registered_modules)
 
         modules_statuses = [
-            status_pb2.ModuleStatusResponse(module_id=module.module_id, status=module.status.name)
+            status_pb2.ModuleStatusResponse(module_id=module.module_id, status=module.status.value)
             for module in list(self.registered_modules.values())[request.offset : request.offset + list_size]
         ]
 
@@ -409,7 +409,7 @@ class RegistryServicer(module_registry_service_pb2_grpc.ModuleRegistryServiceSer
         for module in self.registered_modules.values():
             yield status_pb2.ModuleStatusResponse(
                 module_id=module.module_id,
-                status=module.status.name,
+                status=module.status.value,
             )
 
     def UpdateModuleStatus(  # noqa: N802

digitalkin/grpc_servers/utils/grpc_error_handler.py

@@ -0,0 +1,53 @@
+"""Shared error handling utilities for gRPC services."""
+
+from collections.abc import Generator
+from contextlib import contextmanager
+from typing import Any
+
+from digitalkin.grpc_servers.utils.exceptions import ServerError
+from digitalkin.logger import logger
+
+
+class GrpcErrorHandlerMixin:
+    """Mixin class providing common gRPC error handling functionality."""
+
+    @contextmanager
+    def handle_grpc_errors(  # noqa: PLR6301
+        self,
+        operation: str,
+        service_error_class: type[Exception] | None = None,
+    ) -> Generator[Any, Any, Any]:
+        """Handle gRPC errors for the given operation.
+
+        Args:
+            operation: Name of the operation being performed.
+            service_error_class: Optional specific service exception class to raise.
+                                If not provided, uses the generic ServerError.
+
+        Yields:
+            Context for the operation.
+
+        Raises:
+            ServerError: For gRPC-related errors.
+            service_error_class: For service-specific errors if provided.
+        """
+        if service_error_class is None:
+            service_error_class = ServerError
+
+        try:
+            yield
+        except service_error_class as e:
+            # Re-raise service-specific errors as-is
+            msg = f"{service_error_class.__name__} in {operation}: {e}"
+            logger.exception(msg)
+            raise service_error_class(msg) from e
+        except ServerError as e:
+            # Handle gRPC server errors
+            msg = f"gRPC {operation} failed: {e}"
+            logger.exception(msg)
+            raise ServerError(msg) from e
+        except Exception as e:
+            # Handle unexpected errors
+            msg = f"Unexpected error in {operation}: {e}"
+            logger.exception(msg)
+            raise service_error_class(msg) from e

digitalkin/models/grpc_servers/models.py

@@ -175,8 +175,8 @@ class ClientConfig(ChannelConfig):
     credentials: ClientCredentials | None = Field(None, description="Client credentials for secure mode")
     channel_options: list[tuple[str, Any]] = Field(
         default_factory=lambda: [
-            ("grpc.max_receive_message_length", 100 * 1024 * 1024),  # 100MB
-            ("grpc.max_send_message_length", 100 * 1024 * 1024),  # 100MB
+            ("grpc.max_receive_message_length", 50 * 1024 * 1024),  # 50MB
+            ("grpc.max_send_message_length", 50 * 1024 * 1024),  # 50MB
         ],
         description="Additional channel options",
     )
@@ -223,8 +223,8 @@ class ServerConfig(ChannelConfig):
     credentials: ServerCredentials | None = Field(None, description="Server credentials for secure mode")
     server_options: list[tuple[str, Any]] = Field(
         default_factory=lambda: [
-            ("grpc.max_receive_message_length", 100 * 1024 * 1024),  # 100MB
-            ("grpc.max_send_message_length", 100 * 1024 * 1024),  # 100MB
+            ("grpc.max_receive_message_length", 50 * 1024 * 1024),  # 50MB
+            ("grpc.max_send_message_length", 50 * 1024 * 1024),  # 50MB
         ],
         description="Additional server options",
     )

digitalkin/services/cost/grpc_cost.py

@@ -1,14 +1,12 @@
 """This module implements the gRPC Cost strategy."""
 
-from collections.abc import Generator
-from contextlib import contextmanager
-from typing import Any, Literal
+from typing import Literal
 
-from digitalkin_proto.digitalkin.cost.v1 import cost_pb2, cost_service_pb2_grpc
+from digitalkin_proto.agentic_mesh_protocol.cost.v1 import cost_pb2, cost_service_pb2_grpc
 from google.protobuf import json_format
 
-from digitalkin.grpc_servers.utils.exceptions import ServerError
 from digitalkin.grpc_servers.utils.grpc_client_wrapper import GrpcClientWrapper
+from digitalkin.grpc_servers.utils.grpc_error_handler import GrpcErrorHandlerMixin
 from digitalkin.logger import logger
 from digitalkin.models.grpc_servers.models import ClientConfig
 from digitalkin.services.cost.cost_strategy import (
@@ -20,40 +18,9 @@ from digitalkin.services.cost.cost_strategy import (
 )
 
 
-class GrpcCost(CostStrategy, GrpcClientWrapper):
+class GrpcCost(CostStrategy, GrpcClientWrapper, GrpcErrorHandlerMixin):
     """This class implements the default Cost strategy."""
 
-    @staticmethod
-    @contextmanager
-    def _handle_grpc_errors(operation: str) -> Generator[Any, Any, Any]:
-        """Context manager for consistent gRPC error handling.
-
-        Yields:
-            Allow error handling in context.
-
-        Args:
-            operation: Description of the operation being performed.
-
-        Raises:
-            ValueError: Error with the model validation.
-            ServerError: from gRPC Client.
-            CostServiceError: Unexpected error.
-        """
-        try:
-            yield
-        except CostServiceError as e:
-            msg = f"CostServiceError in {operation}: {e}"
-            logger.exception(msg)
-            raise CostServiceError(msg) from e
-        except ServerError as e:
-            msg = f"gRPC {operation} failed: {e}"
-            logger.exception(msg)
-            raise ServerError(msg) from e
-        except Exception as e:
-            msg = f"Unexpected error in {operation}"
-            logger.exception(msg)
-            raise CostServiceError(msg) from e
-
     def __init__(
         self,
         mission_id: str,
@@ -66,7 +33,7 @@ class GrpcCost(CostStrategy, GrpcClientWrapper):
         super().__init__(mission_id=mission_id, setup_id=setup_id, setup_version_id=setup_version_id, config=config)
         channel = self._init_channel(client_config)
         self.stub = cost_service_pb2_grpc.CostServiceStub(channel)
-        logger.debug("Channel client 'Cost' initialized succesfully")
+        logger.debug("Channel client 'Cost' initialized successfully")
 
     def add(
         self,
@@ -84,7 +51,7 @@ class GrpcCost(CostStrategy, GrpcClientWrapper):
         Raises:
             CostServiceError: If the cost config is invalid
         """
-        with self._handle_grpc_errors("AddCost"):
+        with self.handle_grpc_errors("AddCost", CostServiceError):
             cost_config = self.config.get(cost_config_name)
             if cost_config is None:
                 msg = f"Cost config {cost_config_name} not found in the configuration."
@@ -122,7 +89,7 @@ class GrpcCost(CostStrategy, GrpcClientWrapper):
         Returns:
             CostData: The cost data
         """
-        with self._handle_grpc_errors("GetCost"):
+        with self.handle_grpc_errors("GetCost", CostServiceError):
             request = cost_pb2.GetCostRequest(name=name, mission_id=self.mission_id)
             response: cost_pb2.GetCostResponse = self.exec_grpc_query("GetCost", request)
             cost_data_list = [
@@ -150,7 +117,7 @@ class GrpcCost(CostStrategy, GrpcClientWrapper):
         Returns:
             list[CostData]: The cost data
         """
-        with self._handle_grpc_errors("GetCosts"):
+        with self.handle_grpc_errors("GetCosts", CostServiceError):
             request = cost_pb2.GetCostsRequest(
                 mission_id=self.mission_id,
                 filter=cost_pb2.CostFilter(