digitalkin 0.3.1__py3-none-any.whl → 0.3.1.dev0__py3-none-any.whl

digitalkin/__version__.py

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

digitalkin/core/job_manager/taskiq_broker.py

@@ -208,7 +208,7 @@ async def run_start_module(
         # Reconstruct Pydantic models from dicts for type safety
         try:
             input_model = module_class.create_input_model(input_data)
-            setup_model = await module_class.create_setup_model(setup_data)
+            setup_model = 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

digitalkin/core/job_manager/taskiq_job_manager.py

@@ -140,7 +140,7 @@ class TaskiqJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
         services_mode: ServicesMode,
         default_timeout: float = 10.0,
         max_concurrent_tasks: int = 100,
-        stream_timeout: float = 30.0,
+        stream_timeout: float = 15.0,
     ) -> None:
         """Initialize the Taskiq job manager.
 
@@ -298,6 +298,7 @@ class TaskiqJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
             while True:
                 try:
                     # Block for first item with timeout to allow termination checks
+                    # Configurable timeout (default 15s) to account for distributed system latencies
                     item = await asyncio.wait_for(queue.get(), timeout=self.stream_timeout)
                     queue.task_done()
                     yield item

digitalkin/core/task_manager/base_task_manager.py

@@ -11,7 +11,6 @@ from typing import Any
 from digitalkin.core.task_manager.surrealdb_repository import SurrealDBConnection
 from digitalkin.core.task_manager.task_session import TaskSession
 from digitalkin.logger import logger
-from digitalkin.models.core.task_monitor import CancellationReason
 from digitalkin.modules._base_module import BaseModule
 
 
@@ -85,32 +84,13 @@ class BaseTaskManager(ABC):
             task_id: The ID of the task to clean up
             mission_id: The ID of the mission associated with the task
         """
-        session = self.tasks_sessions.get(task_id)
-        cancellation_reason = session.cancellation_reason.value if session else "no_session"
-        final_status = session.status.value if session else "unknown"
-
         logger.debug(
-            "Cleaning up resources",
-            extra={
-                "mission_id": mission_id,
-                "task_id": task_id,
-                "final_status": final_status,
-                "cancellation_reason": cancellation_reason,
-            },
+            "Cleaning up resources for task: '%s'", task_id, extra={"mission_id": mission_id, "task_id": task_id}
         )
-
-        if session:
+        if task_id in self.tasks_sessions:
+            session = self.tasks_sessions[task_id]
             await session.cleanup()
             self.tasks_sessions.pop(task_id, None)
-            logger.debug(
-                "Task session cleanup completed",
-                extra={
-                    "mission_id": mission_id,
-                    "task_id": task_id,
-                    "final_status": final_status,
-                    "cancellation_reason": cancellation_reason,
-                },
-            )
 
         self.tasks.pop(task_id, None)
 
@@ -281,21 +261,10 @@ class BaseTaskManager(ABC):
             )
 
         except asyncio.TimeoutError:
-            # Set timeout as cancellation reason
-            if task_id in self.tasks_sessions:
-                session = self.tasks_sessions[task_id]
-                if session.cancellation_reason == CancellationReason.UNKNOWN:
-                    session.cancellation_reason = CancellationReason.TIMEOUT
-
             logger.warning(
                 "Graceful cancellation timed out for task: '%s', forcing cancellation",
                 task_id,
-                extra={
-                    "mission_id": mission_id,
-                    "task_id": task_id,
-                    "timeout": timeout,
-                    "cancellation_reason": CancellationReason.TIMEOUT.value,
-                },
+                extra={"mission_id": mission_id, "task_id": task_id, "timeout": timeout},
             )
 
             # Phase 2: Force cancellation
@@ -303,16 +272,8 @@ class BaseTaskManager(ABC):
             with contextlib.suppress(asyncio.CancelledError):
                 await task
 
-            logger.warning(
-                "Task force-cancelled: '%s', reason: %s",
-                task_id,
-                CancellationReason.TIMEOUT.value,
-                extra={
-                    "mission_id": mission_id,
-                    "task_id": task_id,
-                    "cancellation_reason": CancellationReason.TIMEOUT.value,
-                },
-            )
+            logger.warning("Task force-cancelled: '%s'", task_id, extra={"mission_id": mission_id, "task_id": task_id})
+            await self._cleanup_task(task_id, mission_id)
             return True
 
         except Exception as e:
@@ -322,9 +283,10 @@ class BaseTaskManager(ABC):
                 extra={"mission_id": mission_id, "task_id": task_id, "error": str(e)},
                 exc_info=True,
             )
-            return False
-        finally:
             await self._cleanup_task(task_id, mission_id)
+            return False
+
+        await self._cleanup_task(task_id, mission_id)
         return True
 
     async def clean_session(self, task_id: str, mission_id: str) -> bool:
@@ -420,16 +382,7 @@ class BaseTaskManager(ABC):
         results: dict[str, bool | BaseException] = {}
         for task_id, result in zip(task_ids, results_list):
             if isinstance(result, Exception):
-                logger.error(
-                    "Exception cancelling task: '%s', error: %s",
-                    task_id,
-                    result,
-                    extra={
-                        "mission_id": mission_id,
-                        "task_id": task_id,
-                        "error": str(result),
-                    },
-                )
+                logger.error("Exception cancelling task %s: %s", task_id, result)
                 results[task_id] = False
             else:
                 results[task_id] = result
@@ -450,21 +403,6 @@ class BaseTaskManager(ABC):
         )
 
         self._shutdown_event.set()
-
-        # Mark all sessions with shutdown reason before cancellation
-        for task_id, session in self.tasks_sessions.items():
-            if session.cancellation_reason == CancellationReason.UNKNOWN:
-                session.cancellation_reason = CancellationReason.SHUTDOWN
-                logger.debug(
-                    "Marking task for shutdown: '%s'",
-                    task_id,
-                    extra={
-                        "mission_id": mission_id,
-                        "task_id": task_id,
-                        "cancellation_reason": CancellationReason.SHUTDOWN.value,
-                    },
-                )
-
         results = await self.cancel_all_tasks(mission_id, timeout)
 
         failed_tasks = [task_id for task_id, success in results.items() if not success]
@@ -473,26 +411,13 @@ class BaseTaskManager(ABC):
                 "Failed to cancel %d tasks during shutdown: %s",
                 len(failed_tasks),
                 failed_tasks,
-                extra={
-                    "mission_id": mission_id,
-                    "failed_tasks": failed_tasks,
-                    "failed_count": len(failed_tasks),
-                    "cancellation_reason": CancellationReason.SHUTDOWN.value,
-                },
+                extra={"mission_id": mission_id, "failed_tasks": failed_tasks, "failed_count": len(failed_tasks)},
             )
 
         # Clean up any remaining sessions (in case cancellation didn't clean them)
         remaining_sessions = list(self.tasks_sessions.keys())
         if remaining_sessions:
-            logger.info(
-                "Cleaning up %d remaining task sessions after shutdown",
-                len(remaining_sessions),
-                extra={
-                    "mission_id": mission_id,
-                    "remaining_sessions": remaining_sessions,
-                    "remaining_count": len(remaining_sessions),
-                },
-            )
+            logger.info("Cleaning up %d remaining task sessions", len(remaining_sessions))
             cleanup_coros = [self._cleanup_task(task_id, mission_id) for task_id in remaining_sessions]
             await asyncio.gather(*cleanup_coros, return_exceptions=True)
 

digitalkin/core/task_manager/task_executor.py

@@ -8,12 +8,7 @@ from typing import Any
 from digitalkin.core.task_manager.surrealdb_repository import SurrealDBConnection
 from digitalkin.core.task_manager.task_session import TaskSession
 from digitalkin.logger import logger
-from digitalkin.models.core.task_monitor import (
-    CancellationReason,
-    SignalMessage,
-    SignalType,
-    TaskStatus,
-)
+from digitalkin.models.core.task_monitor import SignalMessage, SignalType, TaskStatus
 
 
 class TaskExecutor:
@@ -87,7 +82,7 @@ class TaskExecutor:
             finally:
                 logger.info("Heartbeat task ended", extra={"mission_id": mission_id, "task_id": task_id})
 
-        async def supervisor() -> None:  # noqa: C901, PLR0912, PLR0915
+        async def supervisor() -> None:
             """Supervise the three concurrent tasks and handle outcomes.
 
             Raises:
@@ -101,7 +96,6 @@ class TaskExecutor:
             main_task = None
             hb_task = None
             sig_task = None
-            cleanup_reason = CancellationReason.UNKNOWN
 
             try:
                 main_task = asyncio.create_task(coro, name=f"{task_id}_main")
@@ -112,37 +106,12 @@ class TaskExecutor:
                     return_when=asyncio.FIRST_COMPLETED,
                 )
 
-                # Determine cleanup reason based on which task completed first
-                completed = next(iter(done))
-
-                if completed is main_task:
-                    # Main task finished - cleanup is due to success
-                    cleanup_reason = CancellationReason.SUCCESS_CLEANUP
-                elif completed is sig_task or (completed is hb_task and sig_task.done()):
-                    # Signal task finished - external cancellation
-                    cleanup_reason = CancellationReason.SIGNAL
-                elif completed is hb_task:
-                    # Heartbeat stopped - failure cleanup
-                    cleanup_reason = CancellationReason.FAILURE_CLEANUP
-
-                # Cancel pending tasks with proper reason logging
-                if pending:
-                    pending_names = [t.get_name() for t in pending]
-                    logger.debug(
-                        "Cancelling pending tasks: %s, reason: %s",
-                        pending_names,
-                        cleanup_reason.value,
-                        extra={
-                            "mission_id": mission_id,
-                            "task_id": task_id,
-                            "pending_tasks": pending_names,
-                            "cancellation_reason": cleanup_reason.value,
-                        },
-                    )
-                    for t in pending:
-                        t.cancel()
+                # 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
@@ -153,95 +122,50 @@ class TaskExecutor:
                         extra={"mission_id": mission_id, "task_id": task_id},
                     )
                 elif completed is sig_task or (completed is hb_task and sig_task.done()):
-                    session.status = TaskStatus.CANCELLED
-                    session.cancellation_reason = CancellationReason.SIGNAL
-                    logger.info(
-                        "Task cancelled via external signal",
-                        extra={
-                            "mission_id": mission_id,
-                            "task_id": task_id,
-                            "cancellation_reason": CancellationReason.SIGNAL.value,
-                        },
+                    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
-                    session.cancellation_reason = CancellationReason.HEARTBEAT_FAILURE
                     logger.error(
-                        "Heartbeat stopped unexpectedly for task: '%s'",
-                        task_id,
-                        extra={
-                            "mission_id": mission_id,
-                            "task_id": task_id,
-                            "cancellation_reason": CancellationReason.HEARTBEAT_FAILURE.value,
-                        },
+                        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
-                # Only set reason if not already set (preserve original reason)
-                logger.info(
-                    "Task cancelled externally: '%s', reason: %s",
-                    task_id,
-                    session.cancellation_reason.value,
-                    extra={
-                        "mission_id": mission_id,
-                        "task_id": task_id,
-                        "cancellation_reason": session.cancellation_reason.value,
-                    },
-                )
-                cleanup_reason = CancellationReason.FAILURE_CLEANUP
+                logger.info("Task cancelled", extra={"mission_id": mission_id, "task_id": task_id})
                 raise
             except Exception:
                 session.status = TaskStatus.FAILED
-                cleanup_reason = CancellationReason.FAILURE_CLEANUP
-                logger.exception(
-                    "Task failed with exception: '%s'",
-                    task_id,
-                    extra={"mission_id": mission_id, "task_id": task_id},
-                )
+                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 with proper reason
-                tasks_to_cleanup = [t for t in [main_task, hb_task, sig_task] if t is not None and not t.done()]
-                if tasks_to_cleanup:
-                    cleanup_names = [t.get_name() for t in tasks_to_cleanup]
-                    logger.debug(
-                        "Final cleanup of %d remaining tasks: %s, reason: %s",
-                        len(tasks_to_cleanup),
-                        cleanup_names,
-                        cleanup_reason.value,
-                        extra={
-                            "mission_id": mission_id,
-                            "task_id": task_id,
-                            "cleanup_count": len(tasks_to_cleanup),
-                            "cleanup_tasks": cleanup_names,
-                            "cancellation_reason": cleanup_reason.value,
-                        },
-                    )
-                    for t in tasks_to_cleanup:
+                # 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)
 
-                duration = (
-                    (session.completed_at - session.started_at).total_seconds()
-                    if session.started_at and session.completed_at
-                    else None
-                )
                 logger.info(
-                    "Task execution completed: '%s', status: %s, reason: %s, duration: %.2fs",
-                    task_id,
-                    session.status.value,
-                    session.cancellation_reason.value if session.status == TaskStatus.CANCELLED else "n/a",
-                    duration or 0,
+                    "Task execution completed with status: %s",
+                    session.status,
                     extra={
                         "mission_id": mission_id,
                         "task_id": task_id,
-                        "status": session.status.value,
-                        "cancellation_reason": session.cancellation_reason.value,
-                        "duration": duration,
+                        "status": session.status,
+                        "duration": (
+                            (session.completed_at - session.started_at).total_seconds()
+                            if session.started_at and session.completed_at
+                            else None
+                        ),
                     },
                 )
 

digitalkin/core/task_manager/task_session.py

@@ -6,13 +6,7 @@ from collections.abc import AsyncGenerator
 
 from digitalkin.core.task_manager.surrealdb_repository import SurrealDBConnection
 from digitalkin.logger import logger
-from digitalkin.models.core.task_monitor import (
-    CancellationReason,
-    HeartbeatMessage,
-    SignalMessage,
-    SignalType,
-    TaskStatus,
-)
+from digitalkin.models.core.task_monitor import HeartbeatMessage, SignalMessage, SignalType, TaskStatus
 from digitalkin.modules._base_module import BaseModule
 
 
@@ -37,7 +31,6 @@ class TaskSession:
     completed_at: datetime.datetime | None
 
     is_cancelled: asyncio.Event
-    cancellation_reason: CancellationReason
     _paused: asyncio.Event
     _heartbeat_interval: datetime.timedelta
     _last_heartbeat: datetime.datetime
@@ -49,24 +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
         # Bounded queue to prevent unbounded memory growth (max 1000 items)
-        self.queue: asyncio.Queue = asyncio.Queue(maxsize=queue_maxsize)
+        self.queue: asyncio.Queue = asyncio.Queue(maxsize=1000)
 
         self.task_id = task_id
         self.mission_id = mission_id
@@ -79,7 +62,6 @@ class TaskSession:
         self.heartbeat_record_id = None
 
         self.is_cancelled = asyncio.Event()
-        self.cancellation_reason = CancellationReason.UNKNOWN
         self._paused = asyncio.Event()
         self._heartbeat_interval = heartbeat_interval
 
@@ -161,26 +143,17 @@ class TaskSession:
 
     async def generate_heartbeats(self) -> None:
         """Periodic heartbeat generator with cancellation support."""
-        logger.debug(
-            "Heartbeat generator started for task: '%s'",
-            self.task_id,
-            extra={"task_id": self.task_id, "mission_id": self.mission_id},
-        )
+        logger.debug("Heartbeat started")
         while not self.cancelled:
-            logger.debug(
-                "Heartbeat tick for task: '%s', cancelled=%s",
-                self.task_id,
-                self.cancelled,
-                extra={"task_id": self.task_id, "mission_id": self.mission_id},
-            )
+            logger.debug(f"Heartbeat tick for task: '{self.task_id}' | {self.cancelled=}")
             success = await self.send_heartbeat()
             if not success:
                 logger.error(
                     "Heartbeat failed, cancelling task: '%s'",
                     self.task_id,
-                    extra={"task_id": self.task_id, "mission_id": self.mission_id},
+                    extra={"task_id": self.task_id},
                 )
-                await self._handle_cancel(CancellationReason.HEARTBEAT_FAILURE)
+                await self._handle_cancel()
                 break
             await asyncio.sleep(self._heartbeat_interval.total_seconds())
 
@@ -219,7 +192,7 @@ class TaskSession:
                     continue
 
                 if signal["action"] == "cancel":
-                    await self._handle_cancel(CancellationReason.SIGNAL)
+                    await self._handle_cancel()
                 elif signal["action"] == "pause":
                     await self._handle_pause()
                 elif signal["action"] == "resume":
@@ -249,55 +222,26 @@ class TaskSession:
                 extra={"task_id": self.task_id},
             )
 
-    async def _handle_cancel(self, reason: CancellationReason = CancellationReason.UNKNOWN) -> None:
-        """Idempotent cancellation with acknowledgment and reason tracking.
-
-        Args:
-            reason: The reason for cancellation (signal, heartbeat failure, cleanup, etc.)
-        """
+    async def _handle_cancel(self) -> None:
+        """Idempotent cancellation with acknowledgment."""
+        logger.debug("Handle cancel called")
         if self.is_cancelled.is_set():
             logger.debug(
-                "Cancel ignored - task already cancelled: '%s' (existing reason: %s, new reason: %s)",
+                "Cancel signal ignored - task already cancelled: '%s'",
                 self.task_id,
-                self.cancellation_reason.value,
-                reason.value,
-                extra={
-                    "task_id": self.task_id,
-                    "mission_id": self.mission_id,
-                    "existing_reason": self.cancellation_reason.value,
-                    "new_reason": reason.value,
-                },
+                extra={"task_id": self.task_id},
             )
             return
 
-        self.cancellation_reason = reason
+        logger.info(
+            "Cancelling task: '%s'",
+            self.task_id,
+            extra={"task_id": self.task_id},
+        )
+
         self.status = TaskStatus.CANCELLED
         self.is_cancelled.set()
 
-        # Log with appropriate level based on reason
-        if reason in {CancellationReason.SUCCESS_CLEANUP, CancellationReason.FAILURE_CLEANUP}:
-            logger.debug(
-                "Task cancelled (cleanup): '%s', reason: %s",
-                self.task_id,
-                reason.value,
-                extra={
-                    "task_id": self.task_id,
-                    "mission_id": self.mission_id,
-                    "cancellation_reason": reason.value,
-                },
-            )
-        else:
-            logger.info(
-                "Task cancelled: '%s', reason: %s",
-                self.task_id,
-                reason.value,
-                extra={
-                    "task_id": self.task_id,
-                    "mission_id": self.mission_id,
-                    "cancellation_reason": reason.value,
-                },
-            )
-
         # Resume if paused so cancellation can proceed
         if self._paused.is_set():
             self._paused.set()

digitalkin/grpc_servers/module_servicer.py

@@ -112,7 +112,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         # TODO: Secret should be used here as well
         setup_version = request.setup_version
         config_setup_data = self.module_class.create_config_setup_model(json_format.MessageToDict(request.content))
-        setup_version_data = await self.module_class.create_setup_model(
+        setup_version_data = self.module_class.create_setup_model(
             json_format.MessageToDict(request.setup_version.content),
             config_fields=True,
         )
@@ -185,7 +185,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
             msg = "No setup data returned."
             raise ServicerError(msg)
 
-        setup_data = await self.module_class.create_setup_model(setup_data_class.current_setup_version.content)
+        setup_data = self.module_class.create_setup_model(setup_data_class.current_setup_version.content)
 
         # create a task to run the module in background
         job_id = await self.job_manager.create_module_instance_job(
@@ -220,13 +220,9 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
                         break
 
                     if message.get("code", None) is not None and message.get("code") == "__END_OF_STREAM__":
-                        logger.info(
-                            "End of stream via __END_OF_STREAM__",
-                            extra={"job_id": job_id, "mission_id": request.mission_id},
-                        )
+                        yield lifecycle_pb2.StartModuleResponse(success=True, job_id=job_id)
                         break
 
-                    logger.info("Yielding message from job %s: %s", job_id, message)
                     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:
@@ -350,9 +346,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         # Get input schema if available
         try:
             # Convert schema to proto format
-            input_schema_proto = await self.module_class.get_input_format(
-                llm_format=request.llm_format,
-            )
+            input_schema_proto = self.module_class.get_input_format(llm_format=request.llm_format)
             input_format_struct = json_format.Parse(
                 text=input_schema_proto,
                 message=struct_pb2.Struct(),  # pylint: disable=no-member
@@ -388,9 +382,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         # Get output schema if available
         try:
             # Convert schema to proto format
-            output_schema_proto = await self.module_class.get_output_format(
-                llm_format=request.llm_format,
-            )
+            output_schema_proto = self.module_class.get_output_format(llm_format=request.llm_format)
             output_format_struct = json_format.Parse(
                 text=output_schema_proto,
                 message=struct_pb2.Struct(),  # pylint: disable=no-member
@@ -426,7 +418,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         # Get setup schema if available
         try:
             # Convert schema to proto format
-            setup_schema_proto = await self.module_class.get_setup_format(llm_format=request.llm_format)
+            setup_schema_proto = self.module_class.get_setup_format(llm_format=request.llm_format)
             setup_format_struct = json_format.Parse(
                 text=setup_schema_proto,
                 message=struct_pb2.Struct(),  # pylint: disable=no-member
@@ -443,7 +435,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
             setup_schema=setup_format_struct,
         )
 
-    async def GetModuleSecret(  # noqa: N802
+    def GetModuleSecret(  # noqa: N802
         self,
         request: information_pb2.GetModuleSecretRequest,
         context: grpc.ServicerContext,
@@ -462,7 +454,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         # Get secret schema if available
         try:
             # Convert schema to proto format
-            secret_schema_proto = await self.module_class.get_secret_format(llm_format=request.llm_format)
+            secret_schema_proto = self.module_class.get_secret_format(llm_format=request.llm_format)
             secret_format_struct = json_format.Parse(
                 text=secret_schema_proto,
                 message=struct_pb2.Struct(),  # pylint: disable=no-member
@@ -498,7 +490,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         # Get setup schema if available
         try:
             # Convert schema to proto format
-            config_setup_schema_proto = await self.module_class.get_config_setup_format(llm_format=request.llm_format)
+            config_setup_schema_proto = self.module_class.get_config_setup_format(llm_format=request.llm_format)
             config_setup_format_struct = json_format.Parse(
                 text=config_setup_schema_proto,
                 message=struct_pb2.Struct(),  # pylint: disable=no-member

digitalkin/models/core/task_monitor.py

@@ -17,23 +17,6 @@ class TaskStatus(Enum):
     FAILED = "failed"
 
 
-class CancellationReason(Enum):
-    """Reason for task cancellation - helps distinguish cleanup vs real cancellation."""
-
-    # Cleanup cancellations (not errors)
-    SUCCESS_CLEANUP = "success_cleanup"  # Main task completed, cleaning up helper tasks
-    FAILURE_CLEANUP = "failure_cleanup"  # Main task failed, cleaning up helper tasks
-
-    # Real cancellations
-    SIGNAL = "signal"  # External signal requested cancellation
-    HEARTBEAT_FAILURE = "heartbeat_failure"  # Heartbeat stopped working
-    TIMEOUT = "timeout"  # Task timed out
-    SHUTDOWN = "shutdown"  # Manager is shutting down
-
-    # Unknown/unset
-    UNKNOWN = "unknown"  # Reason not determined
-
-
 class SignalType(Enum):
     """Signal type enumeration."""