digitalkin 0.3.0rc0__tar.gz → 0.3.0rc1__tar.gz

{digitalkin-0.3.0rc0 → digitalkin-0.3.0rc1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: digitalkin
-Version: 0.3.0rc0
+Version: 0.3.0rc1
 Summary: SDK to build kin used in DigitalKin
 Author-email: "DigitalKin.ai" <contact@digitalkin.ai>
 License: Attribution-NonCommercial-ShareAlike 4.0 International

{digitalkin-0.3.0rc0 → digitalkin-0.3.0rc1}/pyproject.toml

@@ -12,7 +12,7 @@
 
     keywords = [ "digitalkin", "kin", "agent", "gprc", "sdk" ]
 
-    version = "0.3.0-rc0"
+    version = "0.3.0-rc1"
     classifiers = [
         "Development Status :: 3 - Alpha",
         "Intended Audience :: Developers",
@@ -157,9 +157,9 @@
             "N",     # pep8-naming
             "ERA",   # eradicate
             "RUF",   # Ruff-specific rules
-            # "DOC",   # pydoclint
-            "FURB", # Refurb
-            "PL",   # Pylint
+            "DOC",   # pydoclint
+            "FURB",  # Refurb
+            "PL",    # Pylint
         ]
 
         ignore = [

{digitalkin-0.3.0rc0 → digitalkin-0.3.0rc1}/src/digitalkin/__version__.py

@@ -5,4 +5,4 @@ from importlib.metadata import PackageNotFoundError, version
 try:
     __version__ = version("digitalkin")
 except PackageNotFoundError:
-    __version__ = "0.3.0-rc0"
+    __version__ = "0.3.0-rc1"

digitalkin-0.3.0rc1/src/digitalkin/core/__init__.py

@@ -0,0 +1 @@
+"""Core of Digitlakin defining the task management and sub-modules."""

digitalkin-0.3.0rc1/src/digitalkin/core/job_manager/__init__.py

@@ -0,0 +1 @@
+"""Job Manager logic."""

{digitalkin-0.3.0rc0/src/digitalkin/modules → digitalkin-0.3.0rc1/src/digitalkin/core}/job_manager/base_job_manager.py

@@ -5,15 +5,16 @@ from collections.abc import AsyncGenerator, AsyncIterator, Callable, Coroutine
 from contextlib import asynccontextmanager
 from typing import Any, Generic
 
+from digitalkin.core.task_manager.task_manager import TaskManager
+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.task_monitor import TaskStatus
 from digitalkin.modules._base_module import BaseModule
 from digitalkin.services.services_config import ServicesConfig
 from digitalkin.services.services_models import ServicesMode
 
 
-class BaseJobManager(abc.ABC, Generic[InputModelT, SetupModelT, OutputModelT]):
+class BaseJobManager(abc.ABC, TaskManager, Generic[InputModelT, SetupModelT, OutputModelT]):
     """Abstract base class for managing background module jobs."""
 
     async def start(self) -> None:
@@ -25,7 +26,8 @@ class BaseJobManager(abc.ABC, Generic[InputModelT, SetupModelT, OutputModelT]):
 
     @staticmethod
     async def job_specific_callback(
-        callback: Callable[[str, OutputModelT | ModuleCodeModel], Coroutine[Any, Any, None]], job_id: str
+        callback: Callable[[str, OutputModelT | ModuleCodeModel], Coroutine[Any, Any, None]],
+        job_id: str,
     ) -> Callable[[OutputModelT | ModuleCodeModel], Coroutine[Any, Any, None]]:
         """Generate a job-specific callback function.
 

{digitalkin-0.3.0rc0/src/digitalkin/modules → digitalkin-0.3.0rc1/src/digitalkin/core}/job_manager/single_job_manager.py

@@ -9,19 +9,18 @@ from typing import Any, Generic
 
 import grpc
 
+from digitalkin.core.job_manager.base_job_manager import BaseJobManager
+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 TaskStatus
 from digitalkin.models.module import InputModelT, OutputModelT, SetupModelT
 from digitalkin.models.module.module import ModuleCodeModel
-from digitalkin.models.module.task_monitor import TaskStatus
 from digitalkin.modules._base_module import BaseModule
-from digitalkin.modules.job_manager.base_job_manager import BaseJobManager
-from digitalkin.modules.job_manager.surrealdb_repository import SurrealDBConnection
-from digitalkin.modules.job_manager.task_manager import TaskManager
-from digitalkin.modules.job_manager.task_session import TaskSession
 from digitalkin.services.services_models import ServicesMode
 
 
-class SingleJobManager(BaseJobManager, TaskManager, Generic[InputModelT, OutputModelT, SetupModelT]):
+class SingleJobManager(BaseJobManager, Generic[InputModelT, OutputModelT, SetupModelT]):
     """Manages a single instance of a module job.
 
     This class ensures that only one instance of a module job is active at a time.
@@ -31,7 +30,7 @@ class SingleJobManager(BaseJobManager, TaskManager, Generic[InputModelT, OutputM
 
     async def start(self) -> None:
         """Start manager."""
-        self.channel = SurrealDBConnection("task_manager", datetime.timedelta(seconds=5))
+        self.channel: SurrealDBConnection = SurrealDBConnection("task_manager", datetime.timedelta(seconds=5))
         await self.channel.init_surreal_instance()
 
     def __init__(
@@ -87,7 +86,6 @@ class SingleJobManager(BaseJobManager, TaskManager, Generic[InputModelT, OutputM
 
         Args:
             config_setup_data: The input data required to start the job.
-            setup_data: The setup configuration for the module.
             mission_id: The mission ID associated with the job.
             setup_id: The setup ID associated with the module.
             setup_version_id: The setup ID.
@@ -101,7 +99,7 @@ class SingleJobManager(BaseJobManager, TaskManager, Generic[InputModelT, OutputM
         job_id = str(uuid.uuid4())
         # TODO: Ensure the job_id is unique.
         module = self.module_class(job_id, mission_id=mission_id, setup_id=setup_id, setup_version_id=setup_version_id)
-        self.tasks_sessions[job_id] = TaskSession(job_id, self.channel, module)
+        self.tasks_sessions[job_id] = TaskSession(job_id, mission_id, self.channel, module)
 
         try:
             await module.start_config_setup(
@@ -224,6 +222,7 @@ class SingleJobManager(BaseJobManager, TaskManager, Generic[InputModelT, OutputM
 
         await self.create_task(
             job_id,
+            mission_id,
             module,
             module.start(input_data, setup_data, callback, done_callback=None),
         )
@@ -254,7 +253,7 @@ class SingleJobManager(BaseJobManager, TaskManager, Generic[InputModelT, OutputM
                 await session.module.stop()
 
                 if job_id in self.tasks:
-                    await self.cancel_task(job_id)
+                    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}")

{digitalkin-0.3.0rc0/src/digitalkin/modules → digitalkin-0.3.0rc1/src/digitalkin/core}/job_manager/taskiq_broker.py

@@ -14,11 +14,11 @@ from taskiq.compat import model_validate
 from taskiq.message import BrokerMessage
 from taskiq_aio_pika import AioPikaBroker
 
+from digitalkin.core.job_manager.base_job_manager import BaseJobManager
 from digitalkin.logger import logger
+from digitalkin.models.core.job_manager_models import StreamCodeModel
 from digitalkin.models.module.module_types import OutputModelT
 from digitalkin.modules._base_module import BaseModule
-from digitalkin.modules.job_manager.base_job_manager import BaseJobManager
-from digitalkin.modules.job_manager.job_manager_models import StreamCodeModel
 from digitalkin.services.services_config import ServicesConfig
 from digitalkin.services.services_models import ServicesMode
 
@@ -194,7 +194,6 @@ async def run_config_module(
         module_class: type[BaseModule],
         services_mode: ServicesMode,
         config_setup_data: dict,
-        setup_data: dict,
         context: Allow TaskIQ context access
     """
     logger.warning("%s", services_mode)

{digitalkin-0.3.0rc0/src/digitalkin/modules → digitalkin-0.3.0rc1/src/digitalkin/core}/job_manager/taskiq_job_manager.py

@@ -17,12 +17,12 @@ from typing import TYPE_CHECKING, Any, Generic
 
 from rstream import Consumer, ConsumerOffsetSpecification, MessageContext, OffsetType
 
+from digitalkin.core.job_manager.base_job_manager import BaseJobManager
+from digitalkin.core.job_manager.taskiq_broker import STREAM, STREAM_RETENTION, TASKIQ_BROKER
 from digitalkin.logger import logger
+from digitalkin.models.core.task_monitor import TaskStatus
 from digitalkin.models.module import InputModelT, SetupModelT
-from digitalkin.models.module.task_monitor import TaskStatus
 from digitalkin.modules._base_module import BaseModule
-from digitalkin.modules.job_manager.base_job_manager import BaseJobManager
-from digitalkin.modules.job_manager.taskiq_broker import STREAM, STREAM_RETENTION, TASKIQ_BROKER
 from digitalkin.services.services_models import ServicesMode
 
 if TYPE_CHECKING:
@@ -146,7 +146,6 @@ class TaskiqJobManager(BaseJobManager, Generic[InputModelT, SetupModelT]):
 
         Args:
             config_setup_data: The input data required to start the job.
-            setup_data: The setup configuration for the module.
             mission_id: The mission ID associated with the job.
             setup_id: The setup ID associated with the module.
             setup_version_id: The setup ID.
@@ -158,7 +157,7 @@ class TaskiqJobManager(BaseJobManager, Generic[InputModelT, SetupModelT]):
             TypeError: If the function is called with bad data type.
             ValueError: If the module fails to start.
         """
-        task = TASKIQ_BROKER.find_task("digitalkin.modules.job_manager.taskiq_broker:run_config_module")
+        task = TASKIQ_BROKER.find_task("digitalkin.core.taskiq_broker:run_config_module")
 
         if task is None:
             msg = "Task not found"
@@ -242,7 +241,7 @@ class TaskiqJobManager(BaseJobManager, Generic[InputModelT, SetupModelT]):
         Raises:
             ValueError: If the task is not found.
         """
-        task = TASKIQ_BROKER.find_task("digitalkin.modules.job_manager.taskiq_broker:run_start_module")
+        task = TASKIQ_BROKER.find_task("digitalkin.core.taskiq_broker:run_start_module")
 
         if task is None:
             msg = "Task not found"

digitalkin-0.3.0rc1/src/digitalkin/core/task_manager/__init__.py

@@ -0,0 +1 @@
+"""Base task manager logic."""

{digitalkin-0.3.0rc0/src/digitalkin/modules/job_manager → digitalkin-0.3.0rc1/src/digitalkin/core/task_manager}/surrealdb_repository.py

@@ -71,7 +71,6 @@ class SurrealDBConnection(Generic[TSurreal]):
 
         Args:
             database: AsyncSurrealDB connection to a specific database
-            table_name: Name of the table to interact with
             timeout: Timeout for database operations
         """
         self.timeout = timeout

{digitalkin-0.3.0rc0/src/digitalkin/modules/job_manager → digitalkin-0.3.0rc1/src/digitalkin/core/task_manager}/task_manager.py

@@ -6,14 +6,18 @@ 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.module.task_monitor import SignalMessage, SignalType, TaskStatus
+from digitalkin.models.core.task_monitor import SignalMessage, SignalType, TaskStatus
 from digitalkin.modules._base_module import BaseModule
-from digitalkin.modules.job_manager.task_session import SurrealDBConnection, TaskSession
 
 
 class TaskManager:
-    """Task manager with comprehensive lifecycle management."""
+    """Task manager with comprehensive lifecycle management.
+
+    Handle the tasks creation, execution, monitoring, signaling, and cancellation.
+    """
 
     tasks: dict[str, asyncio.Task]
     tasks_sessions: dict[str, TaskSession]
@@ -23,7 +27,7 @@ class TaskManager:
     _shutdown_event: asyncio.Event
 
     def __init__(self, default_timeout: float = 10.0, max_concurrent_tasks: int = 100) -> None:
-        """."""
+        """Defining task manager properties."""
         self.tasks = {}
         self.tasks_sessions = {}
         self.default_timeout = default_timeout
@@ -34,22 +38,32 @@ class TaskManager:
             "TaskManager initialized with max_concurrent_tasks: %d, default_timeout: %.1f",
             max_concurrent_tasks,
             default_timeout,
-            extra={"max_concurrent_tasks": max_concurrent_tasks, "default_timeout": default_timeout},
+            extra={
+                "max_concurrent_tasks": max_concurrent_tasks,
+                "default_timeout": default_timeout,
+            },
         )
 
     @property
     def task_count(self) -> int:
-        """."""
+        """Number of managed tasks."""
         return len(self.tasks_sessions)
 
     @property
     def running_tasks(self) -> set[str]:
-        """."""
+        """Get IDs of currently running tasks."""
         return {task_id for task_id, task in self.tasks.items() if not task.done()}
 
-    async def _cleanup_task(self, task_id: str) -> None:
-        """Clean up task resources."""
-        logger.debug("Cleaning up resources for task: '%s'", task_id, extra={"task_id": task_id})
+    async def _cleanup_task(self, task_id: str, mission_id: str) -> None:
+        """Clean up task resources.
+
+        Args:
+            task_id (str): The ID of the task to clean up.
+            mission_id (str): The ID of the mission associated with the task.
+        """
+        logger.debug(
+            "Cleaning up resources for task: '%s'", task_id, extra={"mission_id": mission_id, "task_id": task_id}
+        )
         if task_id in self.tasks_sessions:
             await self.tasks_sessions[task_id].db.close()
         # Remove from collections
@@ -57,6 +71,7 @@ class TaskManager:
     async def _task_wrapper(  # noqa: C901, PLR0915
         self,
         task_id: str,
+        mission_id: str,
         coro: Coroutine[Any, Any, None],
         session: TaskSession,
     ) -> asyncio.Task[None]:
@@ -75,31 +90,33 @@ class TaskManager:
                     "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={"task_id": task_id})
+                logger.debug("Signal listener cancelled", extra={"mission_id": mission_id, "task_id": task_id})
             finally:
                 await self.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={"task_id": task_id})
+                logger.info("Signal listener ended", extra={"mission_id": mission_id, "task_id": task_id})
 
         async def heartbeat_wrapper() -> None:
             try:
                 await session.generate_heartbeats()
             except asyncio.CancelledError:
-                logger.debug("Signal listener cancelled", extra={"task_id": task_id})
+                logger.debug("Signal listener cancelled", extra={"mission_id": mission_id, "task_id": task_id})
             finally:
-                logger.info("Heartbeat task ended", extra={"task_id": task_id})
+                logger.info("Heartbeat task ended", extra={"mission_id": mission_id, "task_id": task_id})
 
         async def supervisor() -> None:
             session.started_at = datetime.datetime.now(datetime.timezone.utc)
@@ -153,6 +170,7 @@ class TaskManager:
     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),
@@ -167,7 +185,11 @@ class TaskManager:
         if task_id in self.tasks:
             # close Coroutine during runtime
             coro.close()
-            logger.warning("Task creation failed - task already exists: '%s'", task_id, extra={"task_id": task_id})
+            logger.warning(
+                "Task creation failed - task already exists: '%s'",
+                task_id,
+                extra={"mission_id": mission_id, "task_id": task_id},
+            )
             msg = f"Task {task_id} already exists"
             raise ValueError(msg)
 
@@ -177,6 +199,7 @@ class TaskManager:
                 "Task creation failed - max concurrent tasks reached: %d",
                 self.max_concurrent_tasks,
                 extra={
+                    "mission_id": mission_id,
                     "task_id": task_id,
                     "current_count": len(self.tasks),
                     "max_concurrent": self.max_concurrent_tasks,
@@ -189,6 +212,7 @@ class TaskManager:
             "Creating new task: '%s'",
             task_id,
             extra={
+                "mission_id": mission_id,
                 "task_id": task_id,
                 "heartbeat_interval": heartbeat_interval,
                 "connection_timeout": connection_timeout,
@@ -199,17 +223,26 @@ class TaskManager:
             # Initialize components
             channel: SurrealDBConnection = SurrealDBConnection("task_manager", connection_timeout)
             await channel.init_surreal_instance()
-            session = TaskSession(task_id, channel, module, heartbeat_interval)
+            session = TaskSession(task_id, mission_id, channel, module, heartbeat_interval)
 
             self.tasks_sessions[task_id] = session
 
             # Create wrapper task
-            self.tasks[task_id] = asyncio.create_task(self._task_wrapper(task_id, coro, session), name=task_id)
+            self.tasks[task_id] = asyncio.create_task(
+                self._task_wrapper(
+                    task_id,
+                    mission_id,
+                    coro,
+                    session,
+                ),
+                name=task_id,
+            )
 
             logger.info(
                 "Task created successfully: '%s'",
                 task_id,
                 extra={
+                    "mission_id": mission_id,
                     "task_id": task_id,
                     "total_tasks": len(self.tasks),
                 },
@@ -217,13 +250,16 @@ class TaskManager:
 
         except Exception as e:
             logger.error(
-                "Failed to create task: '%s'", task_id, extra={"task_id": task_id, "error": str(e)}, exc_info=True
+                "Failed to create 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)
+            await self._cleanup_task(task_id, mission_id=mission_id)
             raise
 
-    async def send_signal(self, task_id: str, signal_type: str, payload: dict) -> bool:
+    async def send_signal(self, task_id: str, mission_id: str, signal_type: str, payload: dict) -> bool:
         """Send signal to a specific task.
 
         Returns:
@@ -233,7 +269,7 @@ class TaskManager:
             logger.warning(
                 "Cannot send signal - task not found: '%s'",
                 task_id,
-                extra={"task_id": task_id, "signal_type": signal_type},
+                extra={"mission_id": mission_id, "task_id": task_id, "signal_type": signal_type},
             )
             return False
 
@@ -241,20 +277,22 @@ class TaskManager:
             "Sending signal '%s' to task: '%s'",
             signal_type,
             task_id,
-            extra={"task_id": task_id, "signal_type": signal_type, "payload": payload},
+            extra={"mission_id": mission_id, "task_id": task_id, "signal_type": signal_type, "payload": payload},
         )
 
         await self.channel.update("tasks", signal_type, payload)
         return True
 
-    async def cancel_task(self, task_id: str, timeout: float | None = None) -> bool:
+    async def cancel_task(self, task_id: str, mission_id: str, timeout: float | None = None) -> bool:
         """Cancel a task with graceful shutdown and fallback.
 
         Returns:
             bool: True if the task was cancelled successfully, False otherwise.
         """
         if task_id not in self.tasks:
-            logger.warning("Cannot cancel - task not found: '%s'", task_id, extra={"task_id": task_id})
+            logger.warning(
+                "Cannot cancel - task not found: '%s'", task_id, extra={"mission_id": mission_id, "task_id": task_id}
+            )
             return True
 
         timeout = timeout or self.default_timeout
@@ -264,23 +302,25 @@ class TaskManager:
             "Initiating task cancellation: '%s', timeout: %.1fs",
             task_id,
             timeout,
-            extra={"task_id": task_id, "timeout": timeout},
+            extra={"mission_id": mission_id, "task_id": task_id, "timeout": timeout},
         )
 
         try:
             # Phase 1: Cooperative cancellation
-            # await self.send_signal(task_id, "cancel")  # noqa: ERA001
+            # await self.send_signal(task_id, mission_id, "cancel")  # noqa: ERA001
 
             # Wait for graceful shutdown
             await asyncio.wait_for(task, timeout=timeout)
 
-            logger.info("Task cancelled gracefully: '%s'", task_id, extra={"task_id": task_id})
+            logger.info(
+                "Task cancelled gracefully: '%s'", task_id, extra={"mission_id": mission_id, "task_id": task_id}
+            )
 
         except asyncio.TimeoutError:
             logger.warning(
                 "Graceful cancellation timed out for task: '%s', forcing cancellation",
                 task_id,
-                extra={"task_id": task_id, "timeout": timeout},
+                extra={"mission_id": mission_id, "task_id": task_id, "timeout": timeout},
             )
 
             # Phase 2: Force cancellation
@@ -288,61 +328,65 @@ class TaskManager:
             with contextlib.suppress(asyncio.CancelledError):
                 await task
 
-            logger.warning("Task force-cancelled: '%s'", task_id, extra={"task_id": task_id})
+            logger.warning("Task force-cancelled: '%s'", task_id, extra={"mission_id": mission_id, "task_id": task_id})
             return True
 
         except Exception as e:
             logger.error(
                 "Error during task cancellation: '%s'",
                 task_id,
-                extra={"task_id": task_id, "error": str(e)},
+                extra={"mission_id": mission_id, "task_id": task_id, "error": str(e)},
                 exc_info=True,
             )
             return False
         return True
 
-    async def clean_session(self, task_id: str) -> bool:
+    async def clean_session(self, task_id: str, mission_id: str) -> bool:
         """Clean up task session without cancelling the task.
 
         Returns:
             bool: True if the task was cleaned successfully, False otherwise.
         """
         if task_id not in self.tasks_sessions:
-            logger.warning("Cannot clean session - task not found: '%s'", task_id, extra={"task_id": task_id})
+            logger.warning(
+                "Cannot clean session - task not found: '%s'",
+                task_id,
+                extra={"mission_id": mission_id, "task_id": task_id},
+            )
             return False
 
         await self.tasks_sessions[task_id].module.stop()
-        await self.cancel_task(task_id)
+        await self.cancel_task(mission_id, task_id)
 
-        logger.info("Cleaning up session for task: '%s'", task_id, extra={"task_id": task_id})
+        logger.info("Cleaning up session for task: '%s'", task_id, extra={"mission_id": mission_id, "task_id": task_id})
         self.tasks_sessions.pop(task_id, None)
         return True
 
-    async def pause_task(self, task_id: str) -> bool:
+    async def pause_task(self, task_id: str, mission_id: str) -> bool:
         """Pause a running task.
 
         Returns:
             bool: True if the task was paused successfully, False otherwise.
         """
-        return await self.send_signal(task_id, "pause", {})
+        return await self.send_signal(task_id, mission_id, "pause", {})
 
-    async def resume_task(self, task_id: str) -> bool:
+    async def resume_task(self, task_id: str, mission_id: str) -> bool:
         """Resume a paused task.
 
         Returns:
             bool: True if the task was paused successfully, False otherwise.
         """
-        return await self.send_signal(task_id, "resume", {})
+        return await self.send_signal(task_id, mission_id, "resume", {})
 
-    async def get_task_status(self, task_id: str) -> bool:
+    async def get_task_status(self, task_id: str, mission_id: str) -> bool:
         """Request status from a task.
 
         Returns:
             bool: True if the task was paused successfully, False otherwise.
         """
-        return await self.send_signal(task_id, "status", {})
+        return await self.send_signal(task_id, mission_id, "status", {})
 
-    async def cancel_all_tasks(self, timeout: float | None = None) -> dict[str, bool]:
+    async def cancel_all_tasks(self, mission_id: str, timeout: float | None = None) -> dict[str, bool]:
         """Cancel all running tasks.
 
         Returns:
@@ -352,25 +396,27 @@ class TaskManager:
         task_ids = list(self.running_tasks)
 
         logger.info(
-            "Cancelling all tasks: %d tasks", len(task_ids), extra={"task_count": len(task_ids), "timeout": timeout}
+            "Cancelling all tasks: %d tasks",
+            len(task_ids),
+            extra={"mission_id": mission_id, "task_count": len(task_ids), "timeout": timeout},
         )
 
         results = {}
         for task_id in task_ids:
-            results[task_id] = await self.cancel_task(task_id, timeout)
+            results[task_id] = await self.cancel_task(task_id, mission_id, timeout)
 
         return results
 
-    async def shutdown(self, timeout: float = 30.0) -> None:
+    async def shutdown(self, mission_id: str, timeout: float = 30.0) -> None:
         """Graceful shutdown of all tasks."""
         logger.info(
             "TaskManager shutdown initiated, timeout: %.1fs",
             timeout,
-            extra={"timeout": timeout, "active_tasks": len(self.running_tasks)},
+            extra={"mission_id": mission_id, "timeout": timeout, "active_tasks": len(self.running_tasks)},
         )
 
         self._shutdown_event.set()
-        results = await self.cancel_all_tasks(timeout)
+        results = await self.cancel_all_tasks(mission_id, timeout)
 
         failed_tasks = [task_id for task_id, success in results.items() if not success]
         if failed_tasks:
@@ -378,12 +424,16 @@ class TaskManager:
                 "Failed to cancel %d tasks during shutdown: %s",
                 len(failed_tasks),
                 failed_tasks,
-                extra={"failed_tasks": failed_tasks, "failed_count": len(failed_tasks)},
+                extra={"mission_id": mission_id, "failed_tasks": failed_tasks, "failed_count": len(failed_tasks)},
             )
 
         logger.info(
             "TaskManager shutdown completed, cancelled: %d, failed: %d",
             len(results) - len(failed_tasks),
             len(failed_tasks),
-            extra={"cancelled_count": len(results) - len(failed_tasks), "failed_count": len(failed_tasks)},
+            extra={
+                "mission_id": mission_id,
+                "cancelled_count": len(results) - len(failed_tasks),
+                "failed_count": len(failed_tasks),
+            },
         )