# #region TaskManagerModule [C:5] [TYPE Module] [SEMANTICS task, schedule, execution, task-manager] # @defgroup TaskManager Module group. # @BRIEF Thin facade composing TaskGraph (registry), EventBus (log/pub-sub), and JobLifecycle # (state machine) into a single TaskManager interface for backward compatibility. # @LAYER Core # @PRE Plugin loader and database sessions are initialized. # @POST Orchestrates task execution and persistence. # @SIDE_EFFECT Spawns worker threads and flushes logs to DB. # @DATA_CONTRACT Input[plugin_id, params] -> Model[Task, LogEntry] # @RELATION DEPENDS_ON -> [PluginLoader] # @RELATION DEPENDS_ON -> [TaskPersistenceService] # @RELATION DEPENDS_ON -> [TaskLogPersistenceService] # @RELATION DEPENDS_ON -> [TaskContext] # @RELATION DEPENDS_ON -> [TaskGraph] # @RELATION DEPENDS_ON -> [JobLifecycle] # @RELATION DEPENDS_ON -> [EventBus] # @INVARIANT Task IDs are unique. # @TEST_CONTRACT TaskManagerRuntime -> { # required_fields: {plugin_loader: PluginLoader}, # optional_fields: {}, # invariants: ["Must use belief_scope for logging"] # } # @TEST_FIXTURE valid_module -> {"manager_initialized": true} # @TEST_EDGE missing_required_field -> {"plugin_loader": null} # @TEST_EDGE empty_response -> {"tasks": []} # @TEST_EDGE invalid_type -> {"plugin_loader": "string_instead_of_object"} # @TEST_EDGE external_failure -> {"db_unavailable": true} # @TEST_INVARIANT logger_compliance -> verifies: [valid_module] # @RATIONALE Decomposed from 708-line monolithic module into four focused modules (TaskGraph, # EventBus, JobLifecycle, and this facade) to satisfy INV_7. TaskManager now delegates # to sub-services while preserving the public API contract. # @REJECTED Keeping all five concerns (registry, lifecycle, log buffer, subscriptions, execution) # in one module was rejected — it violated INV_7 (708 lines vs 400 max), mixed threading # with async execution, and created an unmaintainable god class. import asyncio from typing import Any from src.core.logger import belief_scope, logger from src.core.task_manager.event_bus import EventBus from src.core.task_manager.graph import TaskGraph from src.core.task_manager.lifecycle import JobLifecycle from src.core.task_manager.models import LogFilter, LogStats, Task, TaskStatus from src.core.task_manager.persistence import TaskLogPersistenceService, TaskPersistenceService # #region TaskManager [C:5] [TYPE Class] [SEMANTICS task, manager, lifecycle, execution, state] # @defgroup TaskManager Module group. # @BRIEF Facade composing TaskGraph, EventBus, and JobLifecycle into a single interface. # @LAYER Core # @RELATION DEPENDS_ON -> [TaskPersistenceService] # @RELATION DEPENDS_ON -> [TaskLogPersistenceService] # @RELATION DEPENDS_ON -> [PluginLoader] # @RELATION DEPENDS_ON -> [TaskContext] # @RELATION DEPENDS_ON -> [TaskGraph] # @RELATION DEPENDS_ON -> [JobLifecycle] # @RELATION DEPENDS_ON -> [EventBus] # @PRE Plugin loader resolves plugin ids and persistence services are available. # @POST In-memory task graph, lifecycle scheduler, and log event bus stay consistent with # persisted task state. # @INVARIANT Task IDs are unique within the registry. # @INVARIANT Each task has exactly one status at any time. # @INVARIANT Log entries are never deleted after being added to a task. # @SIDE_EFFECT Spawns worker threads, flushes logs to database, and mutates task states. # @DATA_CONTRACT Input[plugin_id, params] -> Output[Task] # @RATIONALE Thin facade — all delegating methods are under 15 lines. Actual business logic # (registry CRUD, log buffering, lifecycle state machine) lives in extracted modules. # @REJECTED Keeping all five concerns in one class was rejected — it violated INV_7 (708-line # manager.py) and made the god class impossible to test or maintain independently. class TaskManager: """ Facade composing TaskGraph (registry), EventBus (log/pub-sub), and JobLifecycle (state machine) into a single TaskManager interface. """ # #region __init__ [TYPE Function] [C:5] # @BRIEF Initialize sub-services, create add_log callback, start background flusher. # @PRE plugin_loader is initialized. # @POST TaskManager is ready to accept tasks. # @SIDE_EFFECT Starts background flusher thread and loads persisted task state into memory. def __init__(self, plugin_loader): with belief_scope("TaskManager.__init__"): logger.reason("Initializing task manager runtime services") self.plugin_loader = plugin_loader # Track running asyncio tasks for management/cancellation self._async_tasks: dict[str, asyncio.Task] = {} self.persistence_service = TaskPersistenceService() self.log_persistence_service = TaskLogPersistenceService() # Create sub-services self.graph = TaskGraph(self.persistence_service) self.event_bus = EventBus(self.log_persistence_service) # Create add_log callback that lifecycle and plugins use self._add_log = self._make_add_log_callback() self.lifecycle = JobLifecycle( plugin_loader=self.plugin_loader, graph=self.graph, event_bus=self.event_bus, persistence_service=self.persistence_service, ) # Load persisted tasks on startup self.graph.load_persisted_tasks() # Start the async flusher if an event loop is available try: self.event_bus.start() except RuntimeError: # No running event loop (e.g., sync test context) — skip auto-start pass # Backward-compatible property aliases for tests self.tasks = self.graph.tasks self.task_futures = self.graph.task_futures logger.reflect( "Task manager runtime initialized", extra={"task_count": len(self.tasks)}, ) # #endregion __init__ # #region _make_add_log_callback [C:3] [TYPE Function] # @BRIEF Create an async closure for adding logs that looks up the task and delegates to EventBus. def _make_add_log_callback(self): async def _add_log(task_id, level, message, source="system", metadata=None, context=None): task = self.graph.get_task(task_id) if not task: return await self.event_bus.add_log( task_id, level, message, source, metadata, context, task_logs_list=task.logs, ) return _add_log # #endregion _make_add_log_callback # ── Deprecated legacy aliases (delegate to EventBus) ── # #region _flusher_loop [C:3] [TYPE Function] [SEMANTICS flush,background,async] # @BRIEF Legacy alias delegating to EventBus.async_flusher_loop. async def _flusher_loop(self): await self.event_bus.async_flusher_loop() # #endregion _flusher_loop # #region _flush_logs [C:3] [TYPE Function] [SEMANTICS flush,batch,persistence] # @BRIEF Legacy alias delegating to EventBus._flush_logs. async def _flush_logs(self): await self.event_bus._flush_logs() # #endregion _flush_logs # #region _flush_task_logs [C:3] [TYPE Function] [SEMANTICS flush,single,persistence] # @BRIEF Legacy alias delegating to EventBus.flush_task_logs. async def _flush_task_logs(self, task_id: str): await self.event_bus.flush_task_logs(task_id) # #endregion _flush_task_logs # ── Task CRUD delegates to TaskGraph ── # #region get_task [TYPE Function] [C:2] # @ingroup TaskManager # @BRIEF Retrieves a task by its ID. def get_task(self, task_id: str) -> Task | None: return self.graph.get_task(task_id) # #endregion get_task # #region get_all_tasks [TYPE Function] [C:1] # @BRIEF Retrieves all registered tasks. def get_all_tasks(self) -> list[Task]: return self.graph.get_all_tasks() # #endregion get_all_tasks # #region get_tasks [TYPE Function] [C:3] # @ingroup TaskManager # @BRIEF Retrieves tasks with pagination and optional status filter. def get_tasks( self, limit: int = 10, offset: int = 0, status: TaskStatus | None = None, plugin_ids: list[str] | None = None, completed_only: bool = False, ) -> list[Task]: return self.graph.get_tasks(limit, offset, status, plugin_ids, completed_only) # #endregion get_tasks # #region load_persisted_tasks [TYPE Function] [C:2] # @ingroup TaskManager # @BRIEF Load persisted tasks using persistence service. def load_persisted_tasks(self) -> None: self.graph.load_persisted_tasks(limit=100) # #endregion load_persisted_tasks # #region clear_tasks [TYPE Function] [C:4] # @ingroup TaskManager # @BRIEF Clears tasks based on status filter (also deletes associated logs). # @SIDE_EFFECT Removes tasks from registry and persistence; cancels futures. def clear_tasks(self, status: TaskStatus | None = None) -> int: with belief_scope("TaskManager.clear_tasks"): tasks_to_remove = [] for task_id, task in list(self.tasks.items()): should_remove = False if status: if task.status == status: should_remove = True else: if task.status not in [ TaskStatus.RUNNING, TaskStatus.AWAITING_INPUT, TaskStatus.AWAITING_MAPPING, ]: should_remove = True if should_remove: tasks_to_remove.append(task_id) # Delete logs first, then remove tasks if tasks_to_remove: self.event_bus.delete_logs_for_tasks(tasks_to_remove) removed = self.graph.remove_tasks(tasks_to_remove) logger.info(f"Cleared {removed} tasks.") return removed # #endregion clear_tasks # ── Log delegates to EventBus ── # #region get_task_logs [TYPE Function] [C:3] # @ingroup TaskManager # @BRIEF Retrieves logs for a specific task (from memory or persistence). def get_task_logs( self, task_id: str, log_filter: LogFilter | None = None ) -> list: task = self.graph.get_task(task_id) task_status = task.status if task else None task_logs = task.logs if task else [] return self.event_bus.get_task_logs( task_id, log_filter, task_status=task_status, task_logs=task_logs ) # #endregion get_task_logs # #region get_task_log_stats [TYPE Function] [C:2] # @ingroup TaskManager # @BRIEF Get statistics about logs for a task. def get_task_log_stats(self, task_id: str) -> LogStats: return self.event_bus.get_task_log_stats(task_id) # #endregion get_task_log_stats # #region get_task_log_sources [TYPE Function] [C:2] # @ingroup TaskManager # @BRIEF Get unique sources for a task's logs. def get_task_log_sources(self, task_id: str) -> list[str]: return self.event_bus.get_task_log_sources(task_id) # #endregion get_task_log_sources # ── Subscription delegates to EventBus ── # #region subscribe_logs [TYPE Function] [C:2] # @ingroup TaskManager # @BRIEF Subscribes to real-time logs for a task. async def subscribe_logs(self, task_id: str) -> asyncio.Queue: return await self.event_bus.subscribe_logs(task_id) # #endregion subscribe_logs # #region unsubscribe_logs [TYPE Function] [C:2] # @ingroup TaskManager # @BRIEF Unsubscribes from real-time logs for a task. def unsubscribe_logs(self, task_id: str, queue: asyncio.Queue): self.event_bus.unsubscribe_logs(task_id, queue) # #endregion unsubscribe_logs # ── Status subscribers ── # #region subscribe_status [TYPE Function] [C:2] # @ingroup TaskManager # @BRIEF Subscribes to real-time status updates for a task. async def subscribe_status(self, task_id: str) -> asyncio.Queue: return await self.event_bus.subscribe_status(task_id) # #endregion subscribe_status # #region unsubscribe_status [TYPE Function] [C:2] # @ingroup TaskManager # @BRIEF Unsubscribes from status updates for a task. def unsubscribe_status(self, task_id: str, queue: asyncio.Queue): self.event_bus.unsubscribe_status(task_id, queue) # #endregion unsubscribe_status # #region subscribe_task_events [TYPE Function] [C:2] # @ingroup TaskManager # @BRIEF Subscribes to global task events (all task status changes). async def subscribe_task_events(self) -> asyncio.Queue: return await self.event_bus.subscribe_task_events() # #endregion subscribe_task_events # #region unsubscribe_task_events [TYPE Function] [C:2] # @ingroup TaskManager # @BRIEF Unsubscribes from global task events. def unsubscribe_task_events(self, queue: asyncio.Queue): self.event_bus.unsubscribe_task_events(queue) # #endregion unsubscribe_task_events # ── Lifecycle delegates to JobLifecycle ── # #region create_task [TYPE Function] [C:4] # @ingroup TaskManager # @BRIEF Creates and queues a new task for execution. async def create_task( self, plugin_id: str, params: dict[str, Any], user_id: str | None = None ) -> Task: task = await self.lifecycle.create_task( plugin_id, params, user_id, add_log_callback=self._add_log, ) # Schedule execution and track the asyncio task async_task = asyncio.create_task( self.lifecycle._run_task(task.id, add_log_callback=self._add_log) ) self._async_tasks[task.id] = async_task return task # #endregion create_task # #region _run_task [TYPE Function] [C:4] # @BRIEF Internal method to execute a task with TaskContext support (delegates to lifecycle). # Tracks the asyncio.Task for management/cancellation. async def _run_task(self, task_id: str): async_task = asyncio.create_task( self.lifecycle._run_task(task_id, add_log_callback=self._add_log) ) self._async_tasks[task_id] = async_task try: await async_task finally: self._async_tasks.pop(task_id, None) # #endregion _run_task # #region cancel_task [TYPE Function] [C:3] # @ingroup TaskManager # @BRIEF Cancel a running task by ID. # @PRE Task must be currently tracked as running. # @POST Task is cancelled and removed from tracking dict. async def cancel_task(self, task_id: str) -> bool: async_task = self._async_tasks.get(task_id) if async_task is None or async_task.done(): return False async_task.cancel() self._async_tasks.pop(task_id, None) return True # #endregion cancel_task # #region resolve_task [TYPE Function] [C:3] # @ingroup TaskManager # @BRIEF Resumes a task that is awaiting mapping. async def resolve_task(self, task_id: str, resolution_params: dict[str, Any]): await self.lifecycle.resolve_task(task_id, resolution_params) # #endregion resolve_task # #region wait_for_resolution [TYPE Function] [C:3] # @ingroup TaskManager # @BRIEF Pauses execution and waits for a resolution signal. async def wait_for_resolution(self, task_id: str): await self.lifecycle.wait_for_resolution(task_id) # #endregion wait_for_resolution # #region wait_for_input [TYPE Function] [C:3] # @ingroup TaskManager # @BRIEF Pauses execution and waits for user input. async def wait_for_input(self, task_id: str): await self.lifecycle.wait_for_input(task_id) # #endregion wait_for_input # #region await_input [TYPE Function] [C:3] # @ingroup TaskManager # @BRIEF Transition a task to AWAITING_INPUT state with input request. async def await_input(self, task_id: str, input_request: dict[str, Any]) -> None: await self.lifecycle.await_input( task_id, input_request, add_log_callback=self._add_log, ) # #endregion await_input # #region resume_task_with_password [TYPE Function] [C:3] # @ingroup TaskManager # @BRIEF Resume a task that is awaiting input with provided passwords. async def resume_task_with_password( self, task_id: str, passwords: dict[str, str] ) -> None: await self.lifecycle.resume_task_with_password( task_id, passwords, add_log_callback=self._add_log, ) # #endregion resume_task_with_password # ── Maintenance event delegates to EventBus ── # #region subscribe_maintenance_events [TYPE Function] [C:2] # @ingroup TaskManager # @BRIEF Subscribes to global maintenance events. async def subscribe_maintenance_events(self) -> asyncio.Queue: return await self.event_bus.subscribe_maintenance_events() # #endregion subscribe_maintenance_events # #region unsubscribe_maintenance_events [TYPE Function] [C:2] # @ingroup TaskManager # @BRIEF Unsubscribes from global maintenance events. def unsubscribe_maintenance_events(self, queue: asyncio.Queue): self.event_bus.unsubscribe_maintenance_events(queue) # #endregion unsubscribe_maintenance_events # #region broadcast_maintenance_event [TYPE Function] [C:2] # @ingroup TaskManager # @BRIEF Broadcast a maintenance event to all subscribers. async def broadcast_maintenance_event(self, event: dict): await self.event_bus.broadcast_maintenance_event(event) # #endregion broadcast_maintenance_event # ── Dataset event delegates to JobLifecycle ── # #region subscribe_dataset_events [TYPE Function] [C:2] # @ingroup TaskManager # @BRIEF Subscribe to dataset.updated events for an environment. async def subscribe_dataset_events(self, env_id: str) -> asyncio.Queue: return await self.lifecycle.subscribe_dataset_events(env_id) # #endregion subscribe_dataset_events # #region unsubscribe_dataset_events [TYPE Function] [C:2] # @ingroup TaskManager # @BRIEF Unsubscribe from dataset.updated events. def unsubscribe_dataset_events(self, env_id: str, queue: asyncio.Queue): self.lifecycle.unsubscribe_dataset_events(env_id, queue) # #endregion unsubscribe_dataset_events # #endregion TaskManager # #endregion TaskManagerModule