digitalkin 0.3.3.dev6__tar.gz → 0.3.3.dev8__tar.gz

{digitalkin-0.3.3.dev6 → digitalkin-0.3.3.dev8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: digitalkin
-Version: 0.3.3.dev6
+Version: 0.3.3.dev8
 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.3.dev6 → digitalkin-0.3.3.dev8}/pyproject.toml

@@ -34,7 +34,7 @@
         "grpcio-status==1.78.0",
         "pydantic==2.12.5",
     ]
-    version = "0.3.3.dev6"
+    version = "0.3.3.dev8"
 
     [project.optional-dependencies]
         profiling = [
@@ -66,7 +66,7 @@
         "mypy==1.19.1",
         "pre-commit==4.5.1",
         "pyright==1.1.408",
-        "ruff==0.15.5",
+        "ruff==0.15.7",
         "twine==6.2.0",
         "types-grpcio-health-checking==1.0.0.20250506",
         "types-grpcio-reflection==1.0.0.20250506",
@@ -86,13 +86,13 @@
         "mkdocs-git-revision-date-localized-plugin==1.5.1",
         "mkdocs-glightbox==0.5.2",
         "mkdocs-include-markdown-plugin==7.2.1",
-        "mkdocs-literate-nav==0.6.2",
+        "mkdocs-literate-nav==0.6.3",
         "mkdocs-llmstxt==0.5.0",
-        "mkdocs-material[imaging]==9.7.5",
+        "mkdocs-material[imaging]==9.7.6",
         "mkdocs-minify-plugin==0.8.0",
         "mkdocs-open-in-new-tab==1.0.8",
         "mkdocs-redirects==1.2.2",
-        "mkdocs-section-index==0.3.10",
+        "mkdocs-section-index==0.3.11",
         "mkdocs==1.6.1",
         "mkdocstrings-python==2.0.3",
         "mkdocstrings==1.0.3",
@@ -104,7 +104,7 @@
         "hdrhistogram==0.10.3",
         "psutil==7.2.2",
         "pytest-asyncio==1.3.0",
-        "pytest-cov==7.0.0",
+        "pytest-cov==7.1.0",
         "pytest-html==4.2.0",
         "pytest-json-report==1.5.0",
         "pytest-timeout==2.4.0",

{digitalkin-0.3.3.dev6 → digitalkin-0.3.3.dev8}/src/digitalkin/__version__.py

@@ -5,4 +5,4 @@ from importlib.metadata import PackageNotFoundError, version
 try:
     __version__ = version("digitalkin")
 except PackageNotFoundError:
-    __version__ = "0.3.3.dev6"
+    __version__ = "0.3.3.dev8"

{digitalkin-0.3.3.dev6 → digitalkin-0.3.3.dev8}/src/digitalkin/core/job_manager/base_job_manager.py

@@ -130,6 +130,12 @@ class BaseJobManager(abc.ABC, Generic[InputModelT, OutputModelT, SetupModelT]):
         required for the job manager to function.
         """
 
+    async def stop(self) -> None:
+        """Stop the job manager and clean up resources.
+
+        Default no-op. Subclasses with external connections override.
+        """
+
     @staticmethod
     async def job_specific_callback(
         callback: Callable[[str, DataModel | ModuleCodeModel], Coroutine[Any, Any, None]],

{digitalkin-0.3.3.dev6 → digitalkin-0.3.3.dev8}/src/digitalkin/core/job_manager/single_job_manager.py

@@ -91,10 +91,11 @@ class SingleJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
                 message=f"Module {job_id} did not respond within {self._config_setup_timeout} seconds",
             )
         finally:
-            logger.debug(
-                "Config setup response retrieved",
-                extra={"job_id": job_id, "queue_empty": session.queue.empty()},
-            )
+            self.tasks_sessions.pop(job_id, None)
+            try:
+                await session.cleanup()
+            except Exception:
+                logger.exception("Config setup session cleanup failed", extra={"job_id": job_id})
 
     async def create_config_setup_instance_job(
         self,
@@ -132,8 +133,12 @@ class SingleJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
             )
             logger.debug("Module %s (%s) started successfully", job_id, module.name)
         except Exception:
-            # Remove the module from the manager in case of an error.
-            del self.tasks_sessions[job_id]
+            session = self.tasks_sessions.pop(job_id, None)
+            if session is not None:
+                try:
+                    await session.cleanup()
+                except Exception:
+                    logger.debug("Session cleanup failed during error handling", exc_info=True)
             logger.exception("Failed to start module", extra={"job_id": job_id})
             raise
         else:

{digitalkin-0.3.3.dev6 → digitalkin-0.3.3.dev8}/src/digitalkin/core/job_manager/taskiq_broker.py

@@ -11,8 +11,10 @@ from rstream import Producer
 from rstream.exceptions import PreconditionFailed
 from taskiq import Context, TaskiqDepends, TaskiqMessage
 from taskiq.abc.formatter import TaskiqFormatter
+from taskiq.abc.middleware import TaskiqMiddleware
 from taskiq.compat import model_validate
 from taskiq.message import BrokerMessage
+from taskiq.result import TaskiqResult
 from taskiq_aio_pika import AioPikaBroker
 
 from digitalkin.core.common import ModuleFactory
@@ -175,6 +177,11 @@ class TaskiqBrokerConfig:
             startup=[TaskiqBrokerConfig.init_rstream],
         )
         broker.formatter = PickleFormatter()
+        redis_url = os.environ.get("DIGITALKIN_TASKIQ_RESULT_BACKEND_URL")
+        if redis_url:
+            from taskiq_redis import RedisAsyncResultBackend
+
+            broker.with_result_backend(RedisAsyncResultBackend(redis_url))
         return broker
 
     @staticmethod
@@ -224,9 +231,54 @@ class TaskiqBrokerConfig:
         await RSTREAM_PRODUCER.send(stream=TaskiqBrokerConfig.STREAM, message=body)
 
 
+class TaskiqLifecycleMiddleware(TaskiqMiddleware):
+    """Lifecycle middleware for structured logging and safety-net EndOfStreamOutput."""
+
+    async def pre_execute(self, message: TaskiqMessage) -> TaskiqMessage:  # noqa: PLR6301
+        """Log task start.
+
+        Returns:
+            The unmodified message.
+        """
+        logger.info("Taskiq task starting: %s (task_name=%s)", message.task_id, message.task_name)
+        return message
+
+    async def post_execute(self, message: TaskiqMessage, result: TaskiqResult) -> None:  # noqa: PLR6301
+        """Log task completion."""
+        log_fn = logger.info if not result.is_err else logger.error
+        log_fn(
+            "Taskiq task finished: %s (task_name=%s, is_err=%s, exec_time=%.3fs)",
+            message.task_id,
+            message.task_name,
+            result.is_err,
+            result.execution_time,
+        )
+
+    async def on_error(  # noqa: PLR6301
+        self,
+        message: TaskiqMessage,
+        result: TaskiqResult,  # noqa: ARG002
+        exception: BaseException,
+    ) -> None:
+        """Safety net: send EndOfStreamOutput if worker task failed to."""
+        logger.error("Taskiq task error: %s (task_name=%s, error=%s)", message.task_id, message.task_name, exception)
+        try:
+            await TaskiqBrokerConfig.send_message_to_stream(
+                message.task_id,
+                ModuleCodeModel(code="WorkerCrash", short_description="Middleware safety net", message=str(exception)),
+            )
+            await TaskiqBrokerConfig.send_message_to_stream(
+                message.task_id,
+                DataModel(root=EndOfStreamOutput()),
+            )
+        except Exception:
+            logger.exception("Middleware safety net failed for %s", message.task_id)
+
+
 # Module-level globals required by Taskiq framework (decorator needs broker at import time)
 RSTREAM_PRODUCER = TaskiqBrokerConfig.define_producer()
 TASKIQ_BROKER = TaskiqBrokerConfig.define_broker()
+TASKIQ_BROKER.add_middlewares(TaskiqLifecycleMiddleware())
 
 
 @TASKIQ_BROKER.task(task_name="__discarded__")

{digitalkin-0.3.3.dev6 → digitalkin-0.3.3.dev8}/src/digitalkin/core/job_manager/taskiq_job_manager.py

@@ -9,6 +9,7 @@ except ImportError:
 
 import asyncio
 import contextlib
+import datetime
 import json
 import os
 from collections.abc import AsyncGenerator, AsyncIterator
@@ -37,6 +38,9 @@ class TaskiqJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
     """Taskiq job manager for running modules in Taskiq tasks."""
 
     services_mode: ServicesMode
+    stream_consumer: Consumer
+    stream_consumer_task: asyncio.Task[None]
+    _reaper_task: asyncio.Task[None]
 
     @staticmethod
     async def _on_consumer_closed(reason: Any) -> None:
@@ -86,8 +90,95 @@ class TaskiqJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
         job_id = data.get("job_id")
         if not job_id:
             return
+        output_data = data.get("output_data")
         if queue := self.job_queues.get(job_id):
-            await queue.put(data.get("output_data"))
+            await queue.put(output_data)
+
+        # Bridge session status from RStream terminal markers
+        session = self.tasks_sessions.get(job_id)
+        if session is None or not isinstance(output_data, dict):
+            return
+        if "code" in output_data:
+            if session.status not in {"cancelled", "failed"}:
+                session.status = "failed"
+                logger.info("Job %s marked failed from RStream error (code=%s)", job_id, output_data.get("code"))
+        elif isinstance(output_data.get("root"), dict) and output_data["root"].get("protocol") == "end_of_stream":
+            if session.status not in {"cancelled", "failed"}:
+                session.status = "completed"
+                logger.info("Job %s marked completed from RStream end_of_stream", job_id)
+            session.close_stream()
+
+    async def _run_consumer_with_restart(self) -> None:
+        """Run the RStream consumer with automatic restart on failure.
+
+        Raises:
+            CancelledError: If the task is cancelled.
+        """
+        max_retries = int(os.environ.get("DIGITALKIN_RSTREAM_MAX_RETRIES", "10"))
+        base_delay = 1.0
+        max_delay = 60.0
+        attempt = 0
+
+        while True:
+            try:
+                await self.stream_consumer.run()
+                break  # Normal exit (consumer closed gracefully)
+            except asyncio.CancelledError:
+                raise
+            except Exception:
+                attempt += 1
+                if attempt > max_retries:
+                    logger.exception("Stream consumer failed after %d retries, giving up", max_retries)
+                    for session in list(self.tasks_sessions.values()):
+                        if session.status == "pending":
+                            session.status = "failed"
+                            session.close_stream()
+                    break
+                delay = min(base_delay * (2 ** (attempt - 1)), max_delay)
+                logger.exception(
+                    "Stream consumer failed (attempt %d/%d), restarting in %.1fs", attempt, max_retries, delay
+                )
+                await asyncio.sleep(delay)
+                # Reconnect
+                self.stream_consumer = self._define_consumer()
+                await self.stream_consumer.create_stream(
+                    TaskiqBrokerConfig.STREAM,
+                    exists_ok=True,
+                    arguments={"max-length-bytes": TaskiqBrokerConfig.STREAM_RETENTION},
+                )
+                await self.stream_consumer.start()
+                await self.stream_consumer.subscribe(
+                    stream=TaskiqBrokerConfig.STREAM,
+                    subscriber_name=f"""subscriber_{os.environ.get("SERVER_NAME", "module_servicer")}""",
+                    callback=self._on_message,  # type: ignore[arg-type]
+                    offset_specification=ConsumerOffsetSpecification(OffsetType.LAST),
+                    initial_credit=int(os.environ.get("DIGITALKIN_RSTREAM_INITIAL_CREDIT", "50")),
+                )
+                logger.info("Stream consumer reconnected (attempt %d)", attempt)
+
+    async def _reap_orphan_sessions(self) -> None:
+        """Mark sessions stuck in pending beyond timeout as failed.
+
+        Handles hard worker crashes where no EndOfStreamOutput arrives.
+        """
+        orphan_timeout = float(os.environ.get("DIGITALKIN_ORPHAN_SESSION_TIMEOUT", "600.0"))
+        check_interval = float(os.environ.get("DIGITALKIN_ORPHAN_CHECK_INTERVAL", "60.0"))
+
+        while True:
+            try:
+                await asyncio.sleep(check_interval)
+            except asyncio.CancelledError:
+                return
+            now = datetime.datetime.now(datetime.timezone.utc)
+            for task_id, session in list(self.tasks_sessions.items()):
+                if session.status != "pending":
+                    continue
+                elapsed = (now - session.created_at).total_seconds()
+                if elapsed > orphan_timeout:
+                    logger.warning("Orphan session: %s (pending %.0fs)", task_id, elapsed)
+                    session.status = "failed"
+                    session.close_stream()
+                    await self._task_manager._cleanup_task(task_id, session.mission_id)  # noqa: SLF001
 
     async def start(self) -> None:
         """Start the TaskiqJobManager (no-op for external connections)."""
@@ -116,37 +207,41 @@ class TaskiqJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
             initial_credit=int(os.environ.get("DIGITALKIN_RSTREAM_INITIAL_CREDIT", "50")),
         )
 
-        # Wrap the consumer task with error handling
-        async def run_consumer_with_error_handling() -> None:
-            try:
-                await self.stream_consumer.run()
-            except asyncio.CancelledError:
-                logger.debug("Stream consumer task cancelled")
-                raise
-            except Exception:
-                logger.exception("Stream consumer task failed")
-                raise
-
         self.stream_consumer_task = asyncio.create_task(
-            run_consumer_with_error_handling(),
+            self._run_consumer_with_restart(),
             name="stream_consumer_task",
         )
 
-    async def _stop(self) -> None:
-        """Stop the TaskiqJobManager and clean up all resources."""
-        # Signal the consumer to stop
+        self._reaper_task = asyncio.create_task(self._reap_orphan_sessions(), name="orphan_session_reaper")
+
+    async def stop(self) -> None:
+        """Stop the TaskiqJobManager, cancel workers, and clean up all resources."""
+        # 1. Cancel reaper
+        self._reaper_task.cancel()
+        with contextlib.suppress(asyncio.CancelledError):
+            await self._reaper_task
+
+        # 2. Cancel all running modules (sends cancel signals to workers)
+        await self.stop_all_modules()
+
+        # 3. Clean remaining sessions (releases semaphore slots)
+        for task_id in list(self.tasks_sessions.keys()):
+            session = self.tasks_sessions.get(task_id)
+            if session is not None:
+                await self._task_manager._cleanup_task(task_id, session.mission_id)  # noqa: SLF001
+
+        # 4. Close RStream consumer
         await self.stream_consumer.close()
-        # Cancel the background task
         self.stream_consumer_task.cancel()
         with contextlib.suppress(asyncio.CancelledError):
             await self.stream_consumer_task
 
-        # Clean up job queues
+        # 5. Clear job queues
         queue_count = len(self.job_queues)
         self.job_queues.clear()
-        logger.info("TaskiqJobManager: Cleared %d job queues", queue_count)
+        logger.info("TaskiqJobManager stopped: cleared %d queues", queue_count)
 
-        # Call global cleanup for producer and broker
+        # 6. Close producer and broker
         await TaskiqBrokerConfig.cleanup_global_resources()
 
     def __init__(
@@ -192,11 +287,11 @@ class TaskiqJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
         Raises:
             asyncio.TimeoutError: If waiting for the setup response times out.
         """
-        queue = QueueFactory.create_bounded_queue(maxsize=self.max_queue_size)
-        self.job_queues[job_id] = queue
+        if job_id not in self.job_queues:
+            self.job_queues[job_id] = QueueFactory.create_bounded_queue(maxsize=self.max_queue_size)
+        queue = self.job_queues[job_id]
 
         try:
-            # Add timeout to prevent indefinite blocking
             item = await asyncio.wait_for(queue.get(), timeout=self._config_setup_timeout)
         except asyncio.TimeoutError:
             logger.error(
@@ -207,8 +302,9 @@ class TaskiqJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
             queue.task_done()
             return item
         finally:
-            logger.info("generate_config_setup_module_response: job_id=%s: %s", job_id, self.job_queues[job_id].empty())
             self.job_queues.pop(job_id, None)
+            if (session := self.tasks_sessions.get(job_id)) is not None:
+                await self._task_manager._cleanup_task(job_id, session.mission_id)  # noqa: SLF001
 
     async def create_config_setup_instance_job(
         self,
@@ -260,33 +356,45 @@ class TaskiqJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
 
         job_id = running_task.task_id
 
-        # Create module instance for metadata
-        module = self.module_class(
-            job_id,
-            mission_id=mission_id,
-            setup_id=setup_id,
-            setup_version_id=setup_version_id,
-            request_metadata=request_metadata,
-        )
+        # Pre-create queue to avoid message drop race
+        self.job_queues[job_id] = QueueFactory.create_bounded_queue(maxsize=self.max_queue_size)
 
-        # Register task in TaskManager (remote mode)
-        async def _dummy_coro() -> None:
-            """Dummy coroutine - actual execution happens in worker."""
+        try:
+            # Create module instance for metadata
+            module = self.module_class(
+                job_id,
+                mission_id=mission_id,
+                setup_id=setup_id,
+                setup_version_id=setup_version_id,
+                request_metadata=request_metadata,
+            )
 
-        await self.create_task(
-            job_id,
-            mission_id,
-            module,
-            _dummy_coro(),
-        )
+            # Wire RStream callback so stop() can send EndOfStreamOutput
+            callback = await self.job_specific_callback(TaskiqBrokerConfig.send_message_to_stream, job_id)
+            module.context.callbacks.send_message = callback
+
+            # Register task in TaskManager (remote mode)
+            async def _dummy_coro() -> None:
+                """Dummy coroutine - actual execution happens in worker."""
+
+            await self.create_task(
+                job_id,
+                mission_id,
+                module,
+                _dummy_coro(),
+            )
+        except Exception:
+            self.job_queues.pop(job_id, None)
+            raise
 
-        logger.info("Registered config task: %s, waiting for initial result", job_id)
-        result = await running_task.wait_result(timeout=10)
-        logger.info("Job %s with data %s", job_id, result)
+        logger.info("Registered config task: %s", job_id)
+        if os.environ.get("DIGITALKIN_TASKIQ_RESULT_BACKEND_URL"):
+            result = await running_task.wait_result(timeout=10)
+            logger.debug("Job %s config result: %s", job_id, result)
         return job_id
 
     @asynccontextmanager
-    async def generate_stream_consumer(self, job_id: str) -> AsyncIterator[AsyncGenerator[dict[str, Any], None]]:
+    async def generate_stream_consumer(self, job_id: str) -> AsyncIterator[AsyncGenerator[dict[str, Any], None]]:  # noqa: C901, PLR0915
         """Generate a stream consumer for the RStream stream.
 
         Args:
@@ -295,10 +403,11 @@ class TaskiqJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
         Yields:
             messages: The stream messages from the associated module.
         """
-        queue = QueueFactory.create_bounded_queue(maxsize=self.max_queue_size)
-        self.job_queues[job_id] = queue
+        if job_id not in self.job_queues:
+            self.job_queues[job_id] = QueueFactory.create_bounded_queue(maxsize=self.max_queue_size)
+        queue = self.job_queues[job_id]
 
-        async def _stream() -> AsyncGenerator[dict[str, Any], Any]:
+        async def _stream() -> AsyncGenerator[dict[str, Any], Any]:  # noqa: C901
             """Generate the stream with batch-drain optimization.
 
             Yields:
@@ -352,16 +461,22 @@ class TaskiqJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
                     logger.info("Job %s no longer registered, ending stream", job_id)
                     break
 
-                status = await self.get_module_status(job_id)
+                session = self.tasks_sessions[job_id]
+                if session.stream_closed:
+                    logger.info("Job %s stream closed, draining queue and ending stream", job_id)
+                    while not queue.empty():
+                        item = queue.get_nowait()
+                        queue.task_done()
+                        yield item
+                    break
 
-                if status in {"cancelled", "failed"}:
+                status = await self.get_module_status(job_id)
+                if status in {"cancelled", "failed", "completed"}:
                     logger.info("Job %s has terminal status %s, draining queue and ending stream", job_id, status)
-
                     while not queue.empty():
                         item = queue.get_nowait()
                         queue.task_done()
                         yield item
-
                     break
 
         try:
@@ -417,29 +532,41 @@ class TaskiqJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
         )
         job_id = running_task.task_id
 
-        # Create module instance for metadata
-        module = self.module_class(
-            job_id,
-            mission_id=mission_id,
-            setup_id=setup_id,
-            setup_version_id=setup_version_id,
-            request_metadata=request_metadata,
-        )
+        # Pre-create queue to avoid message drop race
+        self.job_queues[job_id] = QueueFactory.create_bounded_queue(maxsize=self.max_queue_size)
 
-        # Register task in TaskManager (remote mode)
-        async def _dummy_coro() -> None:
-            """Dummy coroutine - actual execution happens in worker."""
+        try:
+            # Create module instance for metadata
+            module = self.module_class(
+                job_id,
+                mission_id=mission_id,
+                setup_id=setup_id,
+                setup_version_id=setup_version_id,
+                request_metadata=request_metadata,
+            )
 
-        await self.create_task(
-            job_id,
-            mission_id,
-            module,
-            _dummy_coro(),
-        )
+            # Wire RStream callback so stop() can send EndOfStreamOutput
+            callback = await self.job_specific_callback(TaskiqBrokerConfig.send_message_to_stream, job_id)
+            module.context.callbacks.send_message = callback
+
+            # Register task in TaskManager (remote mode)
+            async def _dummy_coro() -> None:
+                """Dummy coroutine - actual execution happens in worker."""
+
+            await self.create_task(
+                job_id,
+                mission_id,
+                module,
+                _dummy_coro(),
+            )
+        except Exception:
+            self.job_queues.pop(job_id, None)
+            raise
 
-        logger.info("Registered remote task: %s, waiting for initial result", job_id)
-        result = await running_task.wait_result(timeout=10)
-        logger.debug("Job %s with data %s", job_id, result)
+        logger.info("Registered remote task: %s", job_id)
+        if os.environ.get("DIGITALKIN_TASKIQ_RESULT_BACKEND_URL"):
+            result = await running_task.wait_result(timeout=10)
+            logger.debug("Job %s result: %s", job_id, result)
         return job_id
 
     async def get_module_status(self, job_id: str) -> str:
@@ -458,37 +585,28 @@ class TaskiqJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
         return session.status
 
     async def wait_for_completion(self, job_id: str, max_wait: float = 600.0) -> None:
-        """Wait for a task to complete by polling its status.
+        """Wait for a task to complete via stream-closed event.
 
-        Uses adaptive polling: starts at 50ms for fast jobs, doubles up to 500ms
-        for long-running tasks to reduce CPU overhead while maintaining low latency.
+        Relies on ``_on_message`` setting ``_stream_closed`` when
+        ``end_of_stream`` arrives from RStream. Falls back to ``max_wait``
+        timeout for crash scenarios.
 
         Args:
             job_id: The unique identifier of the job to wait for.
             max_wait: Maximum time in seconds to wait before giving up.
 
         Raises:
-            KeyError: If the job_id is not found in tasks_sessions.
             asyncio.TimeoutError: If max_wait is exceeded.
         """
-        if job_id not in self.tasks_sessions:
-            msg = f"Job {job_id} not found"
-            raise KeyError(msg)
-
-        terminal_states = {"completed", "failed", "cancelled"}
-        poll_interval = 0.05
-        elapsed = 0.0
-        while True:
-            session = self.tasks_sessions.get(job_id)
-            if session is None or session.status in terminal_states:
-                logger.debug("Job %s reached terminal state: %s", job_id, session.status if session else "removed")
-                break
-            if elapsed >= max_wait:
-                logger.error("Job %s: max wait time (%.1fs) exceeded, giving up", job_id, max_wait)
-                raise asyncio.TimeoutError
-            await asyncio.sleep(poll_interval)
-            elapsed += poll_interval
-            poll_interval = min(poll_interval * 2, 0.5)
+        session = self.tasks_sessions.get(job_id)
+        if session is None:
+            return
+        try:
+            await asyncio.wait_for(session._stream_closed.wait(), timeout=max_wait)  # noqa: SLF001
+        except asyncio.TimeoutError:
+            logger.error("Job %s: max wait time (%.1fs) exceeded", job_id, max_wait)
+            raise
+        logger.debug("Job %s: stream closed, completion detected (status=%s)", job_id, session.status)
 
     async def stop_module(self, job_id: str) -> bool:
         """Stop a running module using TaskManager.

{digitalkin-0.3.3.dev6 → digitalkin-0.3.3.dev8}/src/digitalkin/core/task_manager/task_session.py

@@ -31,6 +31,7 @@ class TaskSession:
     task_id: str
     mission_id: str
 
+    created_at: datetime.datetime
     started_at: datetime.datetime | None
     completed_at: datetime.datetime | None
 
@@ -73,6 +74,7 @@ class TaskSession:
         self.task_id = task_id
         self.mission_id = mission_id
 
+        self.created_at = datetime.datetime.now(datetime.timezone.utc)
         self.started_at = None
         self.completed_at = None
 

{digitalkin-0.3.3.dev6 → digitalkin-0.3.3.dev8}/src/digitalkin/grpc_servers/module_server.py

@@ -155,6 +155,16 @@ class ModuleServer(BaseServer):
             except Exception:
                 logger.exception("Failed to shutdown module servicer resources")
 
+            try:
+                await self.module_servicer.job_manager.stop_all_modules()
+            except Exception:
+                logger.exception("Failed to stop all modules during shutdown")
+
+            try:
+                await self.module_servicer.job_manager.stop()
+            except Exception:
+                logger.exception("Failed to stop job manager during shutdown")
+
         # Close server-level registry channel
         if isinstance(self.registry, GrpcRegistry):
             try: