456 lines
20 KiB
Python
456 lines
20 KiB
Python
# #region Core.Manager.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 -> [Core.PluginLoader]
|
|
# @RELATION DEPENDS_ON -> [Core.Persistence.TaskPersistenceService]
|
|
# @RELATION DEPENDS_ON -> [Core.Persistence.TaskLogPersistenceService]
|
|
# @RELATION DEPENDS_ON -> [Core.Context.TaskContext]
|
|
# @RELATION DEPENDS_ON -> [Core.Graph.TaskGraph]
|
|
# @RELATION DEPENDS_ON -> [Core.Lifecycle.JobLifecycle]
|
|
# @RELATION DEPENDS_ON -> [Core.EventBus]
|
|
# @INVARIANT Task IDs are unique.
|
|
# @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 Core.Manager.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 -> [Core.Persistence.TaskPersistenceService]
|
|
# @RELATION DEPENDS_ON -> [Core.Persistence.TaskLogPersistenceService]
|
|
# @RELATION DEPENDS_ON -> [Core.PluginLoader]
|
|
# @RELATION DEPENDS_ON -> [Core.Context.TaskContext]
|
|
# @RELATION DEPENDS_ON -> [Core.Graph.TaskGraph]
|
|
# @RELATION DEPENDS_ON -> [Core.Lifecycle.JobLifecycle]
|
|
# @RELATION DEPENDS_ON -> [Core.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 Core.Manager.Init [C:4] [TYPE Function]
|
|
# @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 Core.Manager.Init
|
|
|
|
# #region Core.Manager.MakeAddLogCallback [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=None, message=None, source="system", metadata=None, context=None, event=None, **_kwargs):
|
|
task = self.graph.get_task(task_id)
|
|
if not task:
|
|
return
|
|
if event is None:
|
|
from ss_tools.shared.cot_logger import build_cot_event
|
|
marker = "EXPLORE" if str(level or "INFO").upper() in {"WARNING", "ERROR"} else "REASON"
|
|
event = build_cot_event(
|
|
src=source if source and "." in source else f"task.{source or 'system'}",
|
|
marker=marker,
|
|
intent=message or "Task event",
|
|
payload=metadata or context,
|
|
error=(message or "Task event") if marker == "EXPLORE" else None,
|
|
level=level or "INFO",
|
|
)
|
|
await self.event_bus.add_log(
|
|
task_id, level=event["level"], event=event,
|
|
task_logs_list=task.logs,
|
|
)
|
|
return _add_log
|
|
# #endregion Core.Manager.MakeAddLogCallback
|
|
|
|
# ── Deprecated legacy aliases (delegate to EventBus) ──
|
|
|
|
# #region Core.Manager.FlusherLoop [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 Core.Manager.FlusherLoop
|
|
|
|
# #region Core.Manager.FlushLogs [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 Core.Manager.FlushLogs
|
|
|
|
# #region Core.Manager.FlushTaskLogs [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 Core.Manager.FlushTaskLogs
|
|
|
|
# ── Task CRUD delegates to TaskGraph ──
|
|
|
|
# #region Core.Manager.GetTask [C:2] [TYPE Function]
|
|
# @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 Core.Manager.GetTask
|
|
|
|
# #region Core.Manager.GetAllTasks [TYPE Function] [C:1]
|
|
# @BRIEF Retrieves all registered tasks.
|
|
def get_all_tasks(self) -> list[Task]:
|
|
return self.graph.get_all_tasks()
|
|
# #endregion Core.Manager.GetAllTasks
|
|
|
|
# #region Core.Manager.GetTasks [C:3] [TYPE Function]
|
|
# @ingroup TaskManager
|
|
# @BRIEF Retrieves tasks with pagination and optional status/plugin/search filters.
|
|
def get_tasks(
|
|
self,
|
|
limit: int = 10,
|
|
offset: int = 0,
|
|
status: TaskStatus | None = None,
|
|
plugin_ids: list[str] | None = None,
|
|
completed_only: bool = False,
|
|
search: str | None = None,
|
|
) -> list[Task]:
|
|
return self.graph.get_tasks(limit, offset, status, plugin_ids, completed_only, search)
|
|
# #endregion Core.Manager.GetTasks
|
|
|
|
# #region Core.Manager.LoadPersistedTasks [C:2] [TYPE Function]
|
|
# @ingroup TaskManager
|
|
# @BRIEF Load persisted tasks using persistence service.
|
|
def load_persisted_tasks(self) -> None:
|
|
self.graph.load_persisted_tasks(limit=100)
|
|
# #endregion Core.Manager.LoadPersistedTasks
|
|
|
|
# #region Core.Manager.ClearTasks [C:4] [TYPE Function]
|
|
# @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.reason("Cleared tasks from registry", payload={"removed_count": removed})
|
|
return removed
|
|
# #endregion Core.Manager.ClearTasks
|
|
|
|
# ── Log delegates to EventBus ──
|
|
|
|
# #region Core.Manager.GetTaskLogs [C:3] [TYPE Function]
|
|
# @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 Core.Manager.GetTaskLogs
|
|
|
|
# #region Core.Manager.GetTaskLogStats [C:2] [TYPE Function]
|
|
# @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 Core.Manager.GetTaskLogStats
|
|
|
|
# #region Core.Manager.GetTaskLogSources [C:2] [TYPE Function]
|
|
# @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 Core.Manager.GetTaskLogSources
|
|
|
|
# ── Subscription delegates to EventBus ──
|
|
|
|
# #region Core.Manager.SubscribeLogs [C:2] [TYPE Function]
|
|
# @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 Core.Manager.SubscribeLogs
|
|
|
|
# #region Core.Manager.UnsubscribeLogs [C:2] [TYPE Function]
|
|
# @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 Core.Manager.UnsubscribeLogs
|
|
|
|
# ── Status subscribers ──
|
|
|
|
# #region Core.Manager.SubscribeStatus [C:2] [TYPE Function]
|
|
# @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 Core.Manager.SubscribeStatus
|
|
|
|
# #region Core.Manager.UnsubscribeStatus [C:2] [TYPE Function]
|
|
# @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 Core.Manager.UnsubscribeStatus
|
|
|
|
# #region Core.Manager.SubscribeTaskEvents [C:2] [TYPE Function]
|
|
# @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 Core.Manager.SubscribeTaskEvents
|
|
|
|
# #region Core.Manager.UnsubscribeTaskEvents [C:2] [TYPE Function]
|
|
# @ingroup TaskManager
|
|
# @BRIEF Unsubscribes from global task events.
|
|
def unsubscribe_task_events(self, queue: asyncio.Queue):
|
|
self.event_bus.unsubscribe_task_events(queue)
|
|
# #endregion Core.Manager.UnsubscribeTaskEvents
|
|
|
|
# ── Lifecycle delegates to JobLifecycle ──
|
|
|
|
# #region Core.Manager.CreateTask [C:4] [TYPE Function]
|
|
# @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
|
|
async_task.add_done_callback(lambda _: self._async_tasks.pop(task.id, None))
|
|
return task
|
|
# #endregion Core.Manager.CreateTask
|
|
|
|
# #region Core.Manager.RunTask [C:4] [TYPE Function]
|
|
# @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 Core.Manager.RunTask
|
|
|
|
# #region Core.Manager.CancelTask [C:3] [TYPE Function]
|
|
# @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 Core.Manager.CancelTask
|
|
|
|
# #region Core.Manager.ResolveTask [C:3] [TYPE Function]
|
|
# @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 Core.Manager.ResolveTask
|
|
|
|
# #region Core.Manager.WaitForResolution [C:3] [TYPE Function]
|
|
# @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 Core.Manager.WaitForResolution
|
|
|
|
# #region Core.Manager.WaitForInput [C:3] [TYPE Function]
|
|
# @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 Core.Manager.WaitForInput
|
|
|
|
# #region Core.Manager.AwaitInput [C:3] [TYPE Function]
|
|
# @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 Core.Manager.AwaitInput
|
|
|
|
# #region Core.Manager.ResumeTaskWithPassword [C:3] [TYPE Function]
|
|
# @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],
|
|
requester_user_id: str | None = None,
|
|
allow_task_override: bool = False,
|
|
) -> None:
|
|
await self.lifecycle.resume_task_with_password(
|
|
task_id, passwords,
|
|
add_log_callback=self._add_log,
|
|
requester_user_id=requester_user_id,
|
|
allow_task_override=allow_task_override,
|
|
)
|
|
# #endregion Core.Manager.ResumeTaskWithPassword
|
|
|
|
# #region Core.Manager.RetryTask [C:3] [TYPE Function]
|
|
# @ingroup TaskManager
|
|
# @BRIEF Retry a failed task by resetting state and re-queuing execution.
|
|
async def retry_task(self, task_id: str) -> Task:
|
|
task = await self.lifecycle.retry_task(
|
|
task_id,
|
|
add_log_callback=self._add_log,
|
|
)
|
|
# Schedule re-execution and track (mirrors create_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
|
|
async_task.add_done_callback(lambda _: self._async_tasks.pop(task_id, None))
|
|
return task
|
|
# #endregion Core.Manager.RetryTask
|
|
|
|
# ── Maintenance event delegates to EventBus ──
|
|
|
|
# #region Core.Manager.SubscribeMaintenanceEvents [C:2] [TYPE Function]
|
|
# @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 Core.Manager.SubscribeMaintenanceEvents
|
|
|
|
# #region Core.Manager.UnsubscribeMaintenanceEvents [C:2] [TYPE Function]
|
|
# @ingroup TaskManager
|
|
# @BRIEF Unsubscribes from global maintenance events.
|
|
def unsubscribe_maintenance_events(self, queue: asyncio.Queue):
|
|
self.event_bus.unsubscribe_maintenance_events(queue)
|
|
# #endregion Core.Manager.UnsubscribeMaintenanceEvents
|
|
|
|
# #region Core.Manager.BroadcastMaintenanceEvent [C:2] [TYPE Function]
|
|
# @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 Core.Manager.BroadcastMaintenanceEvent
|
|
|
|
# ── Dataset event delegates to JobLifecycle ──
|
|
|
|
# #region Core.Manager.SubscribeDatasetEvents [C:2] [TYPE Function]
|
|
# @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 Core.Manager.SubscribeDatasetEvents
|
|
|
|
# #region Core.Manager.UnsubscribeDatasetEvents [C:2] [TYPE Function]
|
|
# @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 Core.Manager.UnsubscribeDatasetEvents
|
|
# #endregion Core.Manager.TaskManager
|
|
# #endregion Core.Manager.TaskManagerModule
|