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

base_server/server_async_insecure.py

@@ -9,8 +9,9 @@ from pathlib import Path
 # Add parent directory to path to enable imports
 sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent))
 
-from digitalkin.grpc_servers._base_server import BaseServer
 from digitalkin.grpc_servers.utils.models import SecurityMode, ServerConfig, ServerMode
+
+from digitalkin.grpc_servers._base_server import BaseServer
 from examples.base_server.mock.mock_pb2 import DESCRIPTOR, HelloReply  # type: ignore
 from examples.base_server.mock.mock_pb2_grpc import (
     Greeter,
@@ -30,7 +31,7 @@ class AsyncGreeterImpl(Greeter):
 
     async def SayHello(self, request, context):  # noqa: N802
         """Asynchronous implementation of SayHello method."""
-        logger.info(f"Received request object: {request}")
+        logger.info("Received request object: %s", request)
         logger.info(f"Request attributes: {vars(request)}")
         logger.info(f"Received request with name: {request.name}")
 
@@ -40,7 +41,7 @@ class AsyncGreeterImpl(Greeter):
             name = "unknown"
             # Check context metadata
             for key, value in context.invocation_metadata():
-                logger.info(f"Metadata: {key}={value}")
+                logger.info("Metadata: %s=%s", key, value)
                 if key.lower() == "name":
                     name = value
 
@@ -97,7 +98,7 @@ async def main_async() -> int:
         # as the KeyboardInterrupt usually breaks out of asyncio.run()
         logger.info("Server stopping due to keyboard interrupt...")
     except Exception as e:
-        logger.exception(f"Error running server: {e}")
+        logger.exception("Error running server: %s", e)
         return 1
     finally:
         # Clean up resources if server was started
@@ -116,7 +117,7 @@ def main():
         logger.info("Server stopped by keyboard interrupt")
         return 0  # Clean exit
     except Exception as e:
-        logger.exception(f"Fatal error: {e}")
+        logger.exception("Fatal error: %s", e)
         return 1
 
 

base_server/server_async_secure.py

@@ -9,13 +9,14 @@ from pathlib import Path
 # Add parent directory to path to enable imports
 sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent))
 
-from digitalkin.grpc_servers._base_server import BaseServer
 from digitalkin.grpc_servers.utils.models import (
     SecurityMode,
     ServerConfig,
     ServerCredentials,
     ServerMode,
 )
+
+from digitalkin.grpc_servers._base_server import BaseServer
 from examples.base_server.mock.mock_pb2 import DESCRIPTOR, HelloReply  # type: ignore
 from examples.base_server.mock.mock_pb2_grpc import (
     Greeter,
@@ -35,7 +36,7 @@ class AsyncGreeterImpl(Greeter):
 
     async def SayHello(self, request, context):  # noqa: N802
         """Asynchronous implementation of SayHello method."""
-        logger.info(f"Received request object: {request}")
+        logger.info("Received request object: %s", request)
         logger.info(f"Request attributes: {vars(request)}")
         logger.info(f"Received request with name: {request.name}")
 
@@ -45,7 +46,7 @@ class AsyncGreeterImpl(Greeter):
             name = "unknown"
             # Check context metadata
             for key, value in context.invocation_metadata():
-                logger.info(f"Metadata: {key}={value}")
+                logger.info("Metadata: %s=%s", key, value)
                 if key.lower() == "name":
                     name = value
 
@@ -115,7 +116,7 @@ async def main_async() -> int:
         # as the KeyboardInterrupt usually breaks out of asyncio.run()
         logger.info("Server stopping due to keyboard interrupt...")
     except Exception as e:
-        logger.exception(f"Error running server: {e}")
+        logger.exception("Error running server: %s", e)
         return 1
     finally:
         # Clean up resources if server was started
@@ -134,7 +135,7 @@ def main():
         logger.info("Server stopped by keyboard interrupt")
         return 0  # Clean exit
     except Exception as e:
-        logger.exception(f"Fatal error: {e}")
+        logger.exception("Fatal error: %s", e)
         return 1
 
 

base_server/server_sync_insecure.py

@@ -8,8 +8,9 @@ from pathlib import Path
 # Add parent directory to path to enable imports
 sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent))
 
-from digitalkin.grpc_servers._base_server import BaseServer
 from digitalkin.grpc_servers.utils.models import SecurityMode, ServerConfig, ServerMode
+
+from digitalkin.grpc_servers._base_server import BaseServer
 from examples.base_server.mock.mock_pb2 import DESCRIPTOR, HelloReply  # type: ignore
 from examples.base_server.mock.mock_pb2_grpc import (
     Greeter,
@@ -29,7 +30,7 @@ class SyncGreeterServicer(Greeter):
 
     def SayHello(self, request, context):  # noqa: N802
         """Implementation of SayHello method."""
-        logger.info(f"Received request object: {request}")
+        logger.info("Received request object: %s", request)
         logger.info(f"Request attributes: {vars(request)}")
         logger.info(f"Received request with name: {request.name}")
 
@@ -39,7 +40,7 @@ class SyncGreeterServicer(Greeter):
             name = "unknown"
             # Check context metadata
             for key, value in context.invocation_metadata():
-                logger.info(f"Metadata: {key}={value}")
+                logger.info("Metadata: %s=%s", key, value)
                 if key.lower() == "name":
                     name = value
 
@@ -92,7 +93,7 @@ def main() -> int:
             server.stop()
 
     except Exception as e:
-        logger.exception(f"Error running server: {e}")
+        logger.exception("Error running server: %s", e)
         return 1
 
     return 0

base_server/server_sync_secure.py

@@ -8,13 +8,14 @@ from pathlib import Path
 # Add parent directory to path to enable imports
 sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent))
 
-from digitalkin.grpc_servers._base_server import BaseServer
 from digitalkin.grpc_servers.utils.models import (
     SecurityMode,
     ServerConfig,
     ServerCredentials,
     ServerMode,
 )
+
+from digitalkin.grpc_servers._base_server import BaseServer
 from examples.base_server.mock.mock_pb2 import DESCRIPTOR, HelloReply  # type: ignore
 from examples.base_server.mock.mock_pb2_grpc import (
     Greeter,
@@ -34,7 +35,7 @@ class SyncGreeterServicer(Greeter):
 
     def SayHello(self, request, context):  # noqa: N802
         """Implementation of SayHello method."""
-        logger.info(f"Received request object: {request}")
+        logger.info("Received request object: %s", request)
         logger.info(f"Request attributes: {vars(request)}")
         logger.info(f"Received request with name: {request.name}")
 
@@ -44,7 +45,7 @@ class SyncGreeterServicer(Greeter):
             name = "unknown"
             # Check context metadata
             for key, value in context.invocation_metadata():
-                logger.info(f"Metadata: {key}={value}")
+                logger.info("Metadata: %s=%s", key, value)
                 if key.lower() == "name":
                     name = value
 
@@ -111,7 +112,7 @@ def main() -> int:
             server.stop()
 
     except Exception as e:
-        logger.exception(f"Error running server: {e}")
+        logger.exception("Error running server: %s", e)
         return 1
 
     return 0

digitalkin/__version__.py

@@ -5,4 +5,4 @@ from importlib.metadata import PackageNotFoundError, version
 try:
     __version__ = version("digitalkin")
 except PackageNotFoundError:
-    __version__ = "0.3.1.dev1"
+    __version__ = "0.3.2.a2"

digitalkin/core/job_manager/base_job_manager.py

@@ -8,8 +8,8 @@ from typing import Any, Generic
 from digitalkin.core.task_manager.base_task_manager import BaseTaskManager
 from digitalkin.core.task_manager.task_session import TaskSession
 from digitalkin.models.core.task_monitor import TaskStatus
-from digitalkin.models.module import InputModelT, OutputModelT, SetupModelT
 from digitalkin.models.module.module import ModuleCodeModel
+from digitalkin.models.module.module_types import InputModelT, OutputModelT, SetupModelT
 from digitalkin.modules._base_module import BaseModule
 from digitalkin.services.services_config import ServicesConfig
 from digitalkin.services.services_models import ServicesMode

digitalkin/core/job_manager/single_job_manager.py

@@ -5,7 +5,7 @@ import datetime
 import uuid
 from collections.abc import AsyncGenerator, AsyncIterator
 from contextlib import asynccontextmanager
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import grpc
 
@@ -13,9 +13,12 @@ from digitalkin.core.common import ConnectionFactory, ModuleFactory
 from digitalkin.core.job_manager.base_job_manager import BaseJobManager
 from digitalkin.core.task_manager.local_task_manager import LocalTaskManager
 from digitalkin.core.task_manager.task_session import TaskSession
+
+if TYPE_CHECKING:
+    from digitalkin.core.task_manager.surrealdb_repository import SurrealDBConnection
 from digitalkin.logger import logger
 from digitalkin.models.core.task_monitor import TaskStatus
-from digitalkin.models.module import InputModelT, OutputModelT, SetupModelT
+from digitalkin.models.module.base_types import InputModelT, OutputModelT, SetupModelT
 from digitalkin.models.module.module import ModuleCodeModel
 from digitalkin.modules._base_module import BaseModule
 from digitalkin.services.services_models import ServicesMode
@@ -29,10 +32,6 @@ class SingleJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
     to handle their output data.
     """
 
-    async def start(self) -> None:
-        """Start manager."""
-        self.channel = await ConnectionFactory.create_surreal_connection("task_manager", datetime.timedelta(seconds=5))
-
     def __init__(
         self,
         module_class: type[BaseModule],
@@ -55,6 +54,11 @@ class SingleJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
         super().__init__(module_class, services_mode, task_manager)
 
         self._lock = asyncio.Lock()
+        self.channel: SurrealDBConnection | None = None
+
+    async def start(self) -> None:
+        """Start manager."""
+        self.channel = await ConnectionFactory.create_surreal_connection("task_manager", datetime.timedelta(seconds=5))
 
     async def generate_config_setup_module_response(self, job_id: str) -> SetupModelT | ModuleCodeModel:
         """Generate a stream consumer for a module's output data.
@@ -86,7 +90,10 @@ class SingleJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
                 message=f"Module {job_id} did not respond within 30 seconds",
             )
         finally:
-            logger.info(f"{job_id=}: {session.queue.empty()}")
+            logger.debug(
+                "Config setup response retrieved",
+                extra={"job_id": job_id, "queue_empty": session.queue.empty()},
+            )
 
     async def create_config_setup_instance_job(
         self,
@@ -110,11 +117,14 @@ class SingleJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
             str: The unique identifier (job ID) of the created job.
 
         Raises:
+            RuntimeError: If start() was not called before creating jobs.
             Exception: If the module fails to start.
         """
         job_id = str(uuid.uuid4())
-        # TODO: Ensure the job_id is unique.
         module = ModuleFactory.create_module_instance(self.module_class, job_id, mission_id, setup_id, setup_version_id)
+        if self.channel is None:
+            msg = "JobManager.start() must be called before creating jobs"
+            raise RuntimeError(msg)
         self.tasks_sessions[job_id] = TaskSession(job_id, mission_id, self.channel, module)
 
         try:
@@ -126,7 +136,7 @@ class SingleJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
         except Exception:
             # Remove the module from the manager in case of an error.
             del self.tasks_sessions[job_id]
-            logger.exception("Failed to start module %s: %s", job_id)
+            logger.exception("Failed to start module", extra={"job_id": job_id})
             raise
         else:
             return job_id
@@ -134,13 +144,33 @@ class SingleJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
     async def add_to_queue(self, job_id: str, output_data: OutputModelT | ModuleCodeModel) -> None:
         """Add output data to the queue for a specific job.
 
-        This method is used as a callback to handle output data generated by a module job.
+        Uses timeout-based backpressure: if the queue is full after 5s,
+        drops the oldest message to make room for the new one.
+        Rejects writes after stream is closed to prevent message loss.
 
         Args:
             job_id: The unique identifier of the job.
             output_data: The output data produced by the job.
         """
-        await self.tasks_sessions[job_id].queue.put(output_data.model_dump())
+        session = self.tasks_sessions.get(job_id)
+        if session is None:
+            logger.warning("Queue write rejected - session not found", extra={"job_id": job_id})
+            return
+
+        if session.stream_closed:
+            logger.debug("Queue write rejected - stream closed", extra={"job_id": job_id})
+            return
+
+        try:
+            await asyncio.wait_for(session.queue.put(output_data.model_dump()), timeout=5.0)
+        except asyncio.TimeoutError:
+            logger.warning("Queue full, dropping oldest message", extra={"job_id": job_id})
+            try:
+                session.queue.get_nowait()
+                session.queue.task_done()
+            except asyncio.QueueEmpty:
+                pass
+            session.queue.put_nowait(output_data.model_dump())
 
     @asynccontextmanager  # type: ignore
     async def generate_stream_consumer(self, job_id: str) -> AsyncIterator[AsyncGenerator[dict[str, Any], None]]:  # type: ignore
@@ -177,42 +207,39 @@ class SingleJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
         logger.debug("Session: %s with Module %s", job_id, session.module)
 
         async def _stream() -> AsyncGenerator[dict[str, Any], Any]:
-            """Stream output data from the module with simple blocking pattern.
+            """Stream output data from the module with bounded blocking.
 
-            This implementation uses a simple one-item-at-a-time pattern optimized
-            for local execution where we have direct access to session status:
-            1. Block waiting for each item
-            2. Check termination conditions after each item
-            3. Clean shutdown when task completes
-
-            This pattern provides:
-            - Immediate termination when task completes
-            - Direct session status monitoring
-            - Simple, predictable behavior for local tasks
+            Uses a 1-second timeout on queue.get() to periodically re-check
+            termination flags, preventing indefinite hangs when the task crashes
+            without producing output.
 
             Yields:
                 dict: Output data generated by the module.
             """
             while True:
-                # Block for next item - if queue is empty but producer not finished yet
-                msg = await session.queue.get()
+                if session.stream_closed or session.is_cancelled.is_set():
+                    logger.debug("Stream ending for job %s (pre-check)", job_id)
+                    break
+
+                try:
+                    msg = await asyncio.wait_for(session.queue.get(), timeout=1.0)
+                except asyncio.TimeoutError:
+                    continue
+
                 try:
                     yield msg
                 finally:
-                    # Always mark task as done, even if consumer raises exception
                     session.queue.task_done()
 
-                # Check termination conditions after each message
-                # This allows immediate shutdown when the task completes
                 if (
-                    session.is_cancelled.is_set()
+                    session.stream_closed
+                    or session.is_cancelled.is_set()
                     or (session.status is TaskStatus.COMPLETED and session.queue.empty())
                     or session.status is TaskStatus.FAILED
                 ):
                     logger.debug(
-                        "Stream ending for job %s: cancelled=%s, status=%s, queue_empty=%s",
+                        "Stream ending for job %s: status=%s, queue_empty=%s",
                         job_id,
-                        session.is_cancelled.is_set(),
                         session.status,
                         session.queue.empty(),
                     )
@@ -259,6 +286,18 @@ class SingleJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
         logger.info("Managed task started: '%s'", job_id, extra={"task_id": job_id})
         return job_id
 
+    async def clean_session(self, task_id: str, mission_id: str) -> bool:
+        """Clean a task's session.
+
+        Args:
+            task_id: Unique identifier for the task.
+            mission_id: Mission identifier.
+
+        Returns:
+            bool: True if the task was successfully cleaned, False otherwise.
+        """
+        return await self._task_manager.clean_session(task_id, mission_id)
+
     async def stop_module(self, job_id: str) -> bool:
         """Stop a running module job.
 
@@ -271,20 +310,23 @@ class SingleJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
         Raises:
             Exception: If an error occurs while stopping the module.
         """
-        logger.info(f"STOP required for {job_id=}")
+        logger.info("Stop module requested", extra={"job_id": job_id})
 
         async with self._lock:
             session = self.tasks_sessions.get(job_id)
 
             if not session:
-                logger.warning(f"session with id: {job_id} not found")
+                logger.warning("Session not found", extra={"job_id": job_id})
                 return False
             try:
                 await session.module.stop()
                 await self.cancel_task(job_id, session.mission_id)
-                logger.debug(f"session {job_id} ({session.module.name}) stopped successfully")
-            except Exception as e:
-                logger.error(f"Error while stopping module {job_id}: {e}")
+                logger.debug(
+                    "Module stopped successfully",
+                    extra={"job_id": job_id, "mission_id": session.mission_id},
+                )
+            except Exception:
+                logger.exception("Error stopping module", extra={"job_id": job_id})
                 raise
             else:
                 return True
@@ -331,7 +373,7 @@ class SingleJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
             await asyncio.gather(*stop_tasks, return_exceptions=True)
 
         # Close SurrealDB connection after stopping all modules
-        if hasattr(self, "channel"):
+        if self.channel is not None:
             try:
                 await self.channel.close()
                 logger.info("SingleJobManager: SurrealDB connection closed")

digitalkin/core/job_manager/taskiq_broker.py

@@ -21,8 +21,9 @@ from digitalkin.core.job_manager.base_job_manager import BaseJobManager
 from digitalkin.core.task_manager.task_executor import TaskExecutor
 from digitalkin.core.task_manager.task_session import TaskSession
 from digitalkin.logger import logger
-from digitalkin.models.core.job_manager_models import StreamCodeModel
-from digitalkin.models.module.module_types import OutputModelT
+from digitalkin.models.module.module import ModuleCodeModel
+from digitalkin.models.module.module_types import DataModel, OutputModelT
+from digitalkin.models.module.utility import EndOfStreamOutput
 from digitalkin.modules._base_module import BaseModule
 from digitalkin.services.services_config import ServicesConfig
 from digitalkin.services.services_models import ServicesMode
@@ -141,7 +142,7 @@ async def cleanup_global_resources() -> None:
         logger.warning("Failed to shutdown Taskiq broker: %s", e)
 
 
-async def send_message_to_stream(job_id: str, output_data: OutputModelT) -> None:  # type: ignore
+async def send_message_to_stream(job_id: str, output_data: OutputModelT | ModuleCodeModel) -> None:  # type: ignore[type-var]
     """Callback define to add a message frame to the Rstream.
 
     Args:
@@ -186,7 +187,7 @@ async def run_start_module(
     module_class.discover()
 
     job_id = context.message.task_id
-    callback = await BaseJobManager.job_specific_callback(send_message_to_stream, job_id)
+    callback = await BaseJobManager.job_specific_callback(send_message_to_stream, job_id)  # type: ignore[type-var]
     module = ModuleFactory.create_module_instance(module_class, job_id, mission_id, setup_id, setup_version_id)
 
     channel = None
@@ -201,14 +202,14 @@ async def run_start_module(
         # Create a proper done callback that handles errors
         async def send_end_of_stream(_: Any) -> None:  # noqa: ANN401
             try:
-                await callback(StreamCodeModel(code="__END_OF_STREAM__"))
+                await callback(DataModel(root=EndOfStreamOutput()))
             except Exception as e:
                 logger.error("Error sending end of stream: %s", e, exc_info=True)
 
         # Reconstruct Pydantic models from dicts for type safety
         try:
             input_model = module_class.create_input_model(input_data)
-            setup_model = module_class.create_setup_model(setup_data)
+            setup_model = await module_class.create_setup_model(setup_data)
         except Exception as e:
             logger.error("Failed to reconstruct models for job %s: %s", job_id, e, exc_info=True)
             raise
@@ -272,7 +273,7 @@ async def run_config_module(
     logger.debug("Services config: %s | Module config: %s", services_config, module_class.services_config)
 
     job_id = context.message.task_id
-    callback = await BaseJobManager.job_specific_callback(send_message_to_stream, job_id)
+    callback = await BaseJobManager.job_specific_callback(send_message_to_stream, job_id)  # type: ignore[type-var]
     module = ModuleFactory.create_module_instance(module_class, job_id, mission_id, setup_id, setup_version_id)
 
     # Override environment variables temporarily to use manager's SurrealDB

digitalkin/core/job_manager/taskiq_job_manager.py

@@ -22,9 +22,10 @@ from digitalkin.core.common import ConnectionFactory, QueueFactory
 from digitalkin.core.job_manager.base_job_manager import BaseJobManager
 from digitalkin.core.job_manager.taskiq_broker import STREAM, STREAM_RETENTION, TASKIQ_BROKER, cleanup_global_resources
 from digitalkin.core.task_manager.remote_task_manager import RemoteTaskManager
+from digitalkin.core.task_manager.surrealdb_repository import SurrealDBConnection
 from digitalkin.logger import logger
 from digitalkin.models.core.task_monitor import TaskStatus
-from digitalkin.models.module import InputModelT, OutputModelT, SetupModelT
+from digitalkin.models.module.module_types import InputModelT, OutputModelT, SetupModelT
 from digitalkin.modules._base_module import BaseModule
 from digitalkin.services.services_models import ServicesMode
 
@@ -36,6 +37,7 @@ class TaskiqJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
     """Taskiq job manager for running modules in Taskiq tasks."""
 
     services_mode: ServicesMode
+    channel: SurrealDBConnection | None
 
     @staticmethod
     def _define_consumer() -> Consumer:
@@ -113,7 +115,7 @@ class TaskiqJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
     async def _stop(self) -> None:
         """Stop the TaskiqJobManager and clean up all resources."""
         # Close SurrealDB connection
-        if hasattr(self, "channel"):
+        if self.channel is not None:
             try:
                 await self.channel.close()
                 logger.info("TaskiqJobManager: SurrealDB connection closed")
@@ -128,8 +130,9 @@ class TaskiqJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
             await self.stream_consumer_task
 
         # Clean up job queues
+        queue_count = len(self.job_queues)
         self.job_queues.clear()
-        logger.info("TaskiqJobManager: Cleared %d job queues", len(self.job_queues))
+        logger.info("TaskiqJobManager: Cleared %d job queues", queue_count)
 
         # Call global cleanup for producer and broker
         await cleanup_global_resources()
@@ -161,6 +164,7 @@ class TaskiqJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
         self.job_queues: dict[str, asyncio.Queue] = {}
         self.max_queue_size = 1000
         self.stream_timeout = stream_timeout
+        self.channel = None
 
     async def generate_config_setup_module_response(self, job_id: str) -> SetupModelT:
         """Generate a stream consumer for a module's output data.
@@ -429,7 +433,7 @@ class TaskiqJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
             return TaskStatus.FAILED
 
         # Safety check: if channel not initialized (start() wasn't called), return FAILED
-        if not hasattr(self, "channel") or self.channel is None:
+        if self.channel is None:
             logger.warning("Job %s status check failed - channel not initialized", job_id)
             return TaskStatus.FAILED
 
@@ -521,7 +525,7 @@ class TaskiqJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
         for job_id in self.tasks_sessions:
             try:
                 status = await self.get_module_status(job_id)
-                task_record = await self.channel.select_by_task_id("tasks", job_id)
+                task_record = await self.channel.select_by_task_id("tasks", job_id)  # type: ignore
 
                 modules_info[job_id] = {
                     "name": self.module_class.__name__,

digitalkin/core/task_manager/base_task_manager.py

@@ -8,6 +8,8 @@ from abc import ABC, abstractmethod
 from collections.abc import Coroutine
 from typing import Any
 
+from typing_extensions import Self
+
 from digitalkin.core.task_manager.surrealdb_repository import SurrealDBConnection
 from digitalkin.core.task_manager.task_session import TaskSession
 from digitalkin.logger import logger
@@ -507,7 +509,7 @@ class BaseTaskManager(ABC):
             },
         )
 
-    async def __aenter__(self) -> "BaseTaskManager":
+    async def __aenter__(self) -> Self:
         """Enter async context manager.
 
         Returns:

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)
@@ -146,7 +149,7 @@ class SurrealDBConnection(Generic[TSurreal]):
         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
+        return cast("list[dict[str, Any]] | dict[str, Any]", result)
 
     async def merge(
         self,
@@ -170,7 +173,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 +197,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 +212,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 +263,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