830 lines
36 KiB
Python
830 lines
36 KiB
Python
# [DEF:TaskManagerModule:Module]
|
|
# @COMPLEXITY: 5
|
|
# @SEMANTICS: task, manager, lifecycle, execution, state
|
|
# @PURPOSE: Manages the lifecycle of tasks, including their creation, execution, and state tracking. It uses a thread pool to run plugins asynchronously.
|
|
# @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.
|
|
# @CONSTRAINT: Must use belief_scope for logging.
|
|
# @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]
|
|
|
|
# [SECTION: IMPORTS]
|
|
import asyncio
|
|
import threading
|
|
import inspect
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from datetime import datetime, timezone
|
|
from typing import Dict, Any, List, Optional
|
|
|
|
from .models import Task, TaskStatus, LogEntry, LogFilter, LogStats
|
|
from .persistence import TaskPersistenceService, TaskLogPersistenceService
|
|
from .context import TaskContext
|
|
from ..logger import logger, belief_scope, should_log_task_level
|
|
# [/SECTION]
|
|
|
|
|
|
# [DEF:TaskManager:Class]
|
|
# @COMPLEXITY: 5
|
|
# @SEMANTICS: task, manager, lifecycle, execution, state
|
|
# @PURPOSE: Manages the lifecycle of tasks, including their creation, execution, and state tracking.
|
|
# @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 for task state hydration.
|
|
# @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]
|
|
class TaskManager:
|
|
"""
|
|
Manages the lifecycle of tasks, including their creation, execution, and state tracking.
|
|
"""
|
|
|
|
# Log flush interval in seconds
|
|
LOG_FLUSH_INTERVAL = 2.0
|
|
|
|
# [DEF:TaskGraph:Block]
|
|
# @COMPLEXITY: 5
|
|
# @PURPOSE: Represents the in-memory task dependency graph spanning task registry nodes, paused futures, and persistence-backed hydration.
|
|
# @RELATION: [DEPENDS_ON] ->[Task]
|
|
# @RELATION: [DEPENDS_ON] ->[TaskPersistenceService]
|
|
# @PRE: Task ids are generated before insertion and persisted tasks can be reconstructed into Task models.
|
|
# @POST: Each live task id resolves to at most one active Task node and optional pause future.
|
|
# @SIDE_EFFECT: Mutates the in-memory task registry and loads persisted state during manager startup.
|
|
# @DATA_CONTRACT: Input[Task lifecycle events] -> Output[tasks:Dict[str, Task], task_futures:Dict[str, asyncio.Future]]
|
|
# @INVARIANT: Registry membership is keyed by unique task id and survives log streaming side channels.
|
|
# [/DEF:TaskGraph:Block]
|
|
|
|
# [DEF:EventBus:Block]
|
|
# @COMPLEXITY: 5
|
|
# @PURPOSE: Coordinates task-scoped log buffering, persistence flushes, and subscriber fan-out for real-time observers.
|
|
# @RELATION: [DEPENDS_ON] ->[LogEntry]
|
|
# @RELATION: [DEPENDS_ON] ->[TaskLogPersistenceService]
|
|
# @PRE: Task log records accept structured LogEntry payloads and subscribers consume asyncio queues.
|
|
# @POST: Accepted task logs are buffered, optionally persisted, and dispatched to active subscribers in emission order.
|
|
# @SIDE_EFFECT: Writes task logs to persistence, mutates in-memory buffers, and pushes log entries into subscriber queues.
|
|
# @DATA_CONTRACT: Input[task_id, LogEntry, Queue subscribers] -> Output[persisted log rows, streamed log events]
|
|
# @INVARIANT: Buffered logs are retried on persistence failure and every subscriber receives only task-scoped events.
|
|
# [/DEF:EventBus:Block]
|
|
|
|
# [DEF:JobLifecycle:Block]
|
|
# @COMPLEXITY: 5
|
|
# @PURPOSE: Encodes task creation, execution, pause/resume, and completion transitions for plugin-backed jobs.
|
|
# @RELATION: [DEPENDS_ON] ->[TaskGraph]
|
|
# @RELATION: [DEPENDS_ON] ->[EventBus]
|
|
# @RELATION: [DEPENDS_ON] ->[TaskContext]
|
|
# @RELATION: [DEPENDS_ON] ->[PluginLoader]
|
|
# @PRE: Requested plugin ids resolve to executable plugins and task state transitions target known TaskStatus values.
|
|
# @POST: Every scheduled task transitions through a valid lifecycle path ending in persisted terminal or waiting state.
|
|
# @SIDE_EFFECT: Schedules async execution, pauses on input/mapping requests, and mutates persisted task lifecycle markers.
|
|
# @DATA_CONTRACT: Input[plugin_id, params, task_id, resolution payloads] -> Output[Task status transitions, task result]
|
|
# @INVARIANT: A task cannot be resumed from a waiting state unless a matching future exists or a new wait future is created.
|
|
# [/DEF:JobLifecycle:Block]
|
|
|
|
# [DEF:__init__:Function]
|
|
# @COMPLEXITY: 5
|
|
# @PURPOSE: Initialize the TaskManager with dependencies.
|
|
# @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.
|
|
# @RELATION: [CALLS] ->[TaskPersistenceService.load_tasks]
|
|
# @RELATION: [DEPENDS_ON] ->[TaskGraph]
|
|
# @RELATION: [DEPENDS_ON] ->[EventBus]
|
|
# @PARAM: plugin_loader - The plugin loader instance.
|
|
def __init__(self, plugin_loader):
|
|
with belief_scope("TaskManager.__init__"):
|
|
logger.reason("Initializing task manager runtime services")
|
|
self.plugin_loader = plugin_loader
|
|
self.tasks: Dict[str, Task] = {}
|
|
self.subscribers: Dict[str, List[asyncio.Queue]] = {}
|
|
self.executor = ThreadPoolExecutor(
|
|
max_workers=5
|
|
) # For CPU-bound plugin execution
|
|
self.persistence_service = TaskPersistenceService()
|
|
self.log_persistence_service = TaskLogPersistenceService()
|
|
|
|
# Log buffer: task_id -> List[LogEntry]
|
|
self._log_buffer: Dict[str, List[LogEntry]] = {}
|
|
self._log_buffer_lock = threading.Lock()
|
|
|
|
# Flusher thread for batch writing logs
|
|
self._flusher_stop_event = threading.Event()
|
|
self._flusher_thread = threading.Thread(
|
|
target=self._flusher_loop, daemon=True
|
|
)
|
|
self._flusher_thread.start()
|
|
|
|
try:
|
|
self.loop = asyncio.get_running_loop()
|
|
except RuntimeError:
|
|
self.loop = asyncio.get_event_loop()
|
|
self.task_futures: Dict[str, asyncio.Future] = {}
|
|
|
|
# Load persisted tasks on startup
|
|
self.load_persisted_tasks()
|
|
logger.reflect(
|
|
"Task manager runtime initialized",
|
|
extra={"task_count": len(self.tasks)},
|
|
)
|
|
|
|
# [/DEF:__init__:Function]
|
|
|
|
# [DEF:_flusher_loop:Function]
|
|
# @COMPLEXITY: 3
|
|
# @PURPOSE: Background thread that periodically flushes log buffer to database.
|
|
# @PRE: TaskManager is initialized.
|
|
# @POST: Logs are batch-written to database every LOG_FLUSH_INTERVAL seconds.
|
|
# @RELATION: [CALLS] ->[TaskManager._flush_logs]
|
|
# @RELATION: [DEPENDS_ON] ->[EventBus]
|
|
def _flusher_loop(self):
|
|
"""Background thread that flushes log buffer to database."""
|
|
while not self._flusher_stop_event.is_set():
|
|
self._flush_logs()
|
|
self._flusher_stop_event.wait(self.LOG_FLUSH_INTERVAL)
|
|
|
|
# [/DEF:_flusher_loop:Function]
|
|
|
|
# [DEF:_flush_logs:Function]
|
|
# @COMPLEXITY: 3
|
|
# @PURPOSE: Flush all buffered logs to the database.
|
|
# @PRE: None.
|
|
# @POST: All buffered logs are written to task_logs table.
|
|
# @RELATION: [CALLS] ->[TaskLogPersistenceService.add_logs]
|
|
# @RELATION: [DEPENDS_ON] ->[EventBus]
|
|
def _flush_logs(self):
|
|
"""Flush all buffered logs to the database."""
|
|
with self._log_buffer_lock:
|
|
task_ids = list(self._log_buffer.keys())
|
|
|
|
for task_id in task_ids:
|
|
with self._log_buffer_lock:
|
|
logs = self._log_buffer.pop(task_id, [])
|
|
|
|
if logs:
|
|
try:
|
|
self.log_persistence_service.add_logs(task_id, logs)
|
|
logger.debug(f"Flushed {len(logs)} logs for task {task_id}")
|
|
except Exception as e:
|
|
logger.error(f"Failed to flush logs for task {task_id}: {e}")
|
|
# Re-add logs to buffer on failure
|
|
with self._log_buffer_lock:
|
|
if task_id not in self._log_buffer:
|
|
self._log_buffer[task_id] = []
|
|
self._log_buffer[task_id].extend(logs)
|
|
|
|
# [/DEF:_flush_logs:Function]
|
|
|
|
# [DEF:_flush_task_logs:Function]
|
|
# @COMPLEXITY: 3
|
|
# @PURPOSE: Flush logs for a specific task immediately.
|
|
# @PRE: task_id exists.
|
|
# @POST: Task's buffered logs are written to database.
|
|
# @PARAM: task_id (str) - The task ID.
|
|
# @RELATION: [CALLS] ->[TaskLogPersistenceService.add_logs]
|
|
# @RELATION: [DEPENDS_ON] ->[EventBus]
|
|
def _flush_task_logs(self, task_id: str):
|
|
"""Flush logs for a specific task immediately."""
|
|
with belief_scope("_flush_task_logs"):
|
|
with self._log_buffer_lock:
|
|
logs = self._log_buffer.pop(task_id, [])
|
|
|
|
if logs:
|
|
try:
|
|
self.log_persistence_service.add_logs(task_id, logs)
|
|
except Exception as e:
|
|
logger.error(f"Failed to flush logs for task {task_id}: {e}")
|
|
|
|
# [/DEF:_flush_task_logs:Function]
|
|
|
|
# [DEF:create_task:Function]
|
|
# @COMPLEXITY: 3
|
|
# @PURPOSE: Creates and queues a new task for execution.
|
|
# @PRE: Plugin with plugin_id exists. Params are valid.
|
|
# @POST: Task is created, added to registry, and scheduled for execution.
|
|
# @PARAM: plugin_id (str) - The ID of the plugin to run.
|
|
# @PARAM: params (Dict[str, Any]) - Parameters for the plugin.
|
|
# @PARAM: user_id (Optional[str]) - ID of the user requesting the task.
|
|
# @RETURN: Task - The created task instance.
|
|
# @THROWS: ValueError if plugin not found or params invalid.
|
|
# @RELATION: [CALLS] ->[TaskPersistenceService.persist_task]
|
|
# @RELATION: [DEPENDS_ON] ->[JobLifecycle]
|
|
# @RELATION: [DEPENDS_ON] ->[TaskGraph]
|
|
async def create_task(
|
|
self, plugin_id: str, params: Dict[str, Any], user_id: Optional[str] = None
|
|
) -> Task:
|
|
with belief_scope("TaskManager.create_task", f"plugin_id={plugin_id}"):
|
|
if not self.plugin_loader.has_plugin(plugin_id):
|
|
logger.error(f"Plugin with ID '{plugin_id}' not found.")
|
|
raise ValueError(f"Plugin with ID '{plugin_id}' not found.")
|
|
|
|
self.plugin_loader.get_plugin(plugin_id)
|
|
|
|
if not isinstance(params, dict):
|
|
logger.error("Task parameters must be a dictionary.")
|
|
raise ValueError("Task parameters must be a dictionary.")
|
|
|
|
logger.reason("Creating task node and scheduling execution")
|
|
task = Task(plugin_id=plugin_id, params=params, user_id=user_id)
|
|
self.tasks[task.id] = task
|
|
self.persistence_service.persist_task(task)
|
|
logger.info(f"Task {task.id} created and scheduled for execution")
|
|
self.loop.create_task(
|
|
self._run_task(task.id)
|
|
) # Schedule task for execution
|
|
logger.reflect(
|
|
"Task creation persisted and execution scheduled",
|
|
extra={"task_id": task.id, "plugin_id": plugin_id},
|
|
)
|
|
return task
|
|
|
|
# [/DEF:create_task:Function]
|
|
|
|
# [DEF:_run_task:Function]
|
|
# @COMPLEXITY: 3
|
|
# @PURPOSE: Internal method to execute a task with TaskContext support.
|
|
# @PRE: Task exists in registry.
|
|
# @POST: Task is executed, status updated to SUCCESS or FAILED.
|
|
# @PARAM: task_id (str) - The ID of the task to run.
|
|
# @RELATION: [CALLS] ->[TaskPersistenceService.persist_task]
|
|
# @RELATION: [DEPENDS_ON] ->[JobLifecycle]
|
|
# @RELATION: [DEPENDS_ON] ->[TaskContext]
|
|
# @RELATION: [DEPENDS_ON] ->[EventBus]
|
|
async def _run_task(self, task_id: str):
|
|
with belief_scope("TaskManager._run_task", f"task_id={task_id}"):
|
|
task = self.tasks[task_id]
|
|
plugin = self.plugin_loader.get_plugin(task.plugin_id)
|
|
|
|
logger.reason(
|
|
"Transitioning task to running state",
|
|
extra={"task_id": task_id, "plugin_id": task.plugin_id},
|
|
)
|
|
logger.info(
|
|
f"Starting execution of task {task_id} for plugin '{plugin.name}'"
|
|
)
|
|
task.status = TaskStatus.RUNNING
|
|
task.started_at = datetime.utcnow()
|
|
self.persistence_service.persist_task(task)
|
|
self._add_log(
|
|
task_id,
|
|
"INFO",
|
|
f"Task started for plugin '{plugin.name}'",
|
|
source="system",
|
|
)
|
|
|
|
try:
|
|
# Prepare params and check if plugin supports new TaskContext
|
|
params = {**task.params, "_task_id": task_id}
|
|
|
|
# Check if plugin's execute method accepts 'context' parameter
|
|
sig = inspect.signature(plugin.execute)
|
|
accepts_context = "context" in sig.parameters
|
|
|
|
if accepts_context:
|
|
# Create TaskContext for new-style plugins
|
|
context = TaskContext(
|
|
task_id=task_id,
|
|
add_log_fn=self._add_log,
|
|
params=params,
|
|
default_source="plugin",
|
|
background_tasks=None,
|
|
)
|
|
|
|
if asyncio.iscoroutinefunction(plugin.execute):
|
|
task.result = await plugin.execute(params, context=context)
|
|
else:
|
|
task.result = await self.loop.run_in_executor(
|
|
self.executor,
|
|
lambda: plugin.execute(params, context=context),
|
|
)
|
|
else:
|
|
# Backward compatibility: old-style plugins without context
|
|
if asyncio.iscoroutinefunction(plugin.execute):
|
|
task.result = await plugin.execute(params)
|
|
else:
|
|
task.result = await self.loop.run_in_executor(
|
|
self.executor, plugin.execute, params
|
|
)
|
|
|
|
logger.info(f"Task {task_id} completed successfully")
|
|
task.status = TaskStatus.SUCCESS
|
|
self._add_log(
|
|
task_id,
|
|
"INFO",
|
|
f"Task completed successfully for plugin '{plugin.name}'",
|
|
source="system",
|
|
)
|
|
except Exception as e:
|
|
logger.error(f"Task {task_id} failed: {e}")
|
|
task.status = TaskStatus.FAILED
|
|
self._add_log(
|
|
task_id,
|
|
"ERROR",
|
|
f"Task failed: {e}",
|
|
source="system",
|
|
metadata={"error_type": type(e).__name__},
|
|
)
|
|
finally:
|
|
task.finished_at = datetime.utcnow()
|
|
# Flush any remaining buffered logs before persisting task
|
|
self._flush_task_logs(task_id)
|
|
self.persistence_service.persist_task(task)
|
|
logger.info(
|
|
f"Task {task_id} execution finished with status: {task.status}"
|
|
)
|
|
logger.reflect(
|
|
"Task lifecycle reached persisted terminal state",
|
|
extra={"task_id": task_id, "status": str(task.status)},
|
|
)
|
|
|
|
# [/DEF:_run_task:Function]
|
|
|
|
# [DEF:resolve_task:Function]
|
|
# @COMPLEXITY: 3
|
|
# @PURPOSE: Resumes a task that is awaiting mapping.
|
|
# @PRE: Task exists and is in AWAITING_MAPPING state.
|
|
# @POST: Task status updated to RUNNING, params updated, execution resumed.
|
|
# @PARAM: task_id (str) - The ID of the task.
|
|
# @PARAM: resolution_params (Dict[str, Any]) - Params to resolve the wait.
|
|
# @THROWS: ValueError if task not found or not awaiting mapping.
|
|
# @RELATION: [CALLS] ->[TaskPersistenceService.persist_task]
|
|
# @RELATION: [DEPENDS_ON] ->[JobLifecycle]
|
|
async def resolve_task(self, task_id: str, resolution_params: Dict[str, Any]):
|
|
with belief_scope("TaskManager.resolve_task", f"task_id={task_id}"):
|
|
task = self.tasks.get(task_id)
|
|
if not task or task.status != TaskStatus.AWAITING_MAPPING:
|
|
raise ValueError("Task is not awaiting mapping.")
|
|
|
|
# Update task params with resolution
|
|
task.params.update(resolution_params)
|
|
task.status = TaskStatus.RUNNING
|
|
self.persistence_service.persist_task(task)
|
|
self._add_log(task_id, "INFO", "Task resumed after mapping resolution.")
|
|
|
|
# Signal the future to continue
|
|
if task_id in self.task_futures:
|
|
self.task_futures[task_id].set_result(True)
|
|
|
|
# [/DEF:resolve_task:Function]
|
|
|
|
# [DEF:wait_for_resolution:Function]
|
|
# @COMPLEXITY: 3
|
|
# @PURPOSE: Pauses execution and waits for a resolution signal.
|
|
# @PRE: Task exists.
|
|
# @POST: Execution pauses until future is set.
|
|
# @PARAM: task_id (str) - The ID of the task.
|
|
# @RELATION: [CALLS] ->[TaskPersistenceService.persist_task]
|
|
# @RELATION: [DEPENDS_ON] ->[JobLifecycle]
|
|
async def wait_for_resolution(self, task_id: str):
|
|
with belief_scope("TaskManager.wait_for_resolution", f"task_id={task_id}"):
|
|
task = self.tasks.get(task_id)
|
|
if not task:
|
|
return
|
|
|
|
task.status = TaskStatus.AWAITING_MAPPING
|
|
self.persistence_service.persist_task(task)
|
|
self.task_futures[task_id] = self.loop.create_future()
|
|
|
|
try:
|
|
await self.task_futures[task_id]
|
|
finally:
|
|
if task_id in self.task_futures:
|
|
del self.task_futures[task_id]
|
|
|
|
# [/DEF:wait_for_resolution:Function]
|
|
|
|
# [DEF:wait_for_input:Function]
|
|
# @COMPLEXITY: 3
|
|
# @PURPOSE: Pauses execution and waits for user input.
|
|
# @PRE: Task exists.
|
|
# @POST: Execution pauses until future is set via resume_task_with_password.
|
|
# @PARAM: task_id (str) - The ID of the task.
|
|
# @RELATION: [DEPENDS_ON] ->[JobLifecycle]
|
|
async def wait_for_input(self, task_id: str):
|
|
with belief_scope("TaskManager.wait_for_input", f"task_id={task_id}"):
|
|
task = self.tasks.get(task_id)
|
|
if not task:
|
|
return
|
|
|
|
# Status is already set to AWAITING_INPUT by await_input()
|
|
self.task_futures[task_id] = self.loop.create_future()
|
|
|
|
try:
|
|
await self.task_futures[task_id]
|
|
finally:
|
|
if task_id in self.task_futures:
|
|
del self.task_futures[task_id]
|
|
|
|
# [/DEF:wait_for_input:Function]
|
|
|
|
# [DEF:get_task:Function]
|
|
# @COMPLEXITY: 3
|
|
# @PURPOSE: Retrieves a task by its ID.
|
|
# @PRE: task_id is a string.
|
|
# @POST: Returns Task object or None.
|
|
# @PARAM: task_id (str) - ID of the task.
|
|
# @RETURN: Optional[Task] - The task or None.
|
|
# @RELATION: [DEPENDS_ON] ->[TaskGraph]
|
|
def get_task(self, task_id: str) -> Optional[Task]:
|
|
with belief_scope("TaskManager.get_task", f"task_id={task_id}"):
|
|
return self.tasks.get(task_id)
|
|
|
|
# [/DEF:get_task:Function]
|
|
|
|
# [DEF:get_all_tasks:Function]
|
|
# @COMPLEXITY: 3
|
|
# @PURPOSE: Retrieves all registered tasks.
|
|
# @PRE: None.
|
|
# @POST: Returns list of all Task objects.
|
|
# @RETURN: List[Task] - All tasks.
|
|
# @RELATION: [DEPENDS_ON] ->[TaskGraph]
|
|
def get_all_tasks(self) -> List[Task]:
|
|
with belief_scope("TaskManager.get_all_tasks"):
|
|
return list(self.tasks.values())
|
|
|
|
# [/DEF:get_all_tasks:Function]
|
|
|
|
# [DEF:get_tasks:Function]
|
|
# @COMPLEXITY: 3
|
|
# @PURPOSE: Retrieves tasks with pagination and optional status filter.
|
|
# @PRE: limit and offset are non-negative integers.
|
|
# @POST: Returns a list of tasks sorted by start_time descending.
|
|
# @PARAM: limit (int) - Maximum number of tasks to return.
|
|
# @PARAM: offset (int) - Number of tasks to skip.
|
|
# @PARAM: status (Optional[TaskStatus]) - Filter by task status.
|
|
# @RETURN: List[Task] - List of tasks matching criteria.
|
|
# @RELATION: [DEPENDS_ON] ->[TaskGraph]
|
|
def get_tasks(
|
|
self,
|
|
limit: int = 10,
|
|
offset: int = 0,
|
|
status: Optional[TaskStatus] = None,
|
|
plugin_ids: Optional[List[str]] = None,
|
|
completed_only: bool = False,
|
|
) -> List[Task]:
|
|
with belief_scope("TaskManager.get_tasks"):
|
|
tasks = list(self.tasks.values())
|
|
if status:
|
|
tasks = [t for t in tasks if t.status == status]
|
|
if plugin_ids:
|
|
plugin_id_set = set(plugin_ids)
|
|
tasks = [t for t in tasks if t.plugin_id in plugin_id_set]
|
|
if completed_only:
|
|
tasks = [
|
|
t for t in tasks if t.status in [TaskStatus.SUCCESS, TaskStatus.FAILED]
|
|
]
|
|
|
|
# Sort by started_at descending with tolerant handling of mixed tz-aware/naive values.
|
|
def sort_key(task: Task) -> float:
|
|
started_at = task.started_at
|
|
if started_at is None:
|
|
return float("-inf")
|
|
if not isinstance(started_at, datetime):
|
|
return float("-inf")
|
|
if started_at.tzinfo is None:
|
|
return started_at.replace(tzinfo=timezone.utc).timestamp()
|
|
return started_at.timestamp()
|
|
|
|
tasks.sort(key=sort_key, reverse=True)
|
|
return tasks[offset : offset + limit]
|
|
|
|
# [/DEF:get_tasks:Function]
|
|
|
|
# [DEF:get_task_logs:Function]
|
|
# @COMPLEXITY: 3
|
|
# @PURPOSE: Retrieves logs for a specific task (from memory for running, persistence for completed).
|
|
# @PRE: task_id is a string.
|
|
# @POST: Returns list of LogEntry or TaskLog objects.
|
|
# @PARAM: task_id (str) - ID of the task.
|
|
# @PARAM: log_filter (Optional[LogFilter]) - Filter parameters.
|
|
# @RETURN: List[LogEntry] - List of log entries.
|
|
# @RELATION: [CALLS] ->[TaskLogPersistenceService.get_logs]
|
|
# @RELATION: [DEPENDS_ON] ->[EventBus]
|
|
def get_task_logs(
|
|
self, task_id: str, log_filter: Optional[LogFilter] = None
|
|
) -> List[LogEntry]:
|
|
with belief_scope("TaskManager.get_task_logs", f"task_id={task_id}"):
|
|
task = self.tasks.get(task_id)
|
|
|
|
# For completed tasks, fetch from persistence
|
|
if task and task.status in [TaskStatus.SUCCESS, TaskStatus.FAILED]:
|
|
if log_filter is None:
|
|
log_filter = LogFilter()
|
|
task_logs = self.log_persistence_service.get_logs(task_id, log_filter)
|
|
# Convert TaskLog to LogEntry for backward compatibility
|
|
return [
|
|
LogEntry(
|
|
timestamp=log.timestamp,
|
|
level=log.level,
|
|
message=log.message,
|
|
source=log.source,
|
|
metadata=log.metadata,
|
|
)
|
|
for log in task_logs
|
|
]
|
|
|
|
# For running/pending tasks, return from memory
|
|
return task.logs if task else []
|
|
|
|
# [/DEF:get_task_logs:Function]
|
|
|
|
# [DEF:get_task_log_stats:Function]
|
|
# @COMPLEXITY: 3
|
|
# @PURPOSE: Get statistics about logs for a task.
|
|
# @PRE: task_id is a valid task ID.
|
|
# @POST: Returns LogStats with counts by level and source.
|
|
# @PARAM: task_id (str) - The task ID.
|
|
# @RETURN: LogStats - Statistics about task logs.
|
|
# @RELATION: [CALLS] ->[TaskLogPersistenceService.get_log_stats]
|
|
# @RELATION: [DEPENDS_ON] ->[EventBus]
|
|
def get_task_log_stats(self, task_id: str) -> LogStats:
|
|
with belief_scope("TaskManager.get_task_log_stats", f"task_id={task_id}"):
|
|
return self.log_persistence_service.get_log_stats(task_id)
|
|
|
|
# [/DEF:get_task_log_stats:Function]
|
|
|
|
# [DEF:get_task_log_sources:Function]
|
|
# @COMPLEXITY: 3
|
|
# @PURPOSE: Get unique sources for a task's logs.
|
|
# @PRE: task_id is a valid task ID.
|
|
# @POST: Returns list of unique source strings.
|
|
# @PARAM: task_id (str) - The task ID.
|
|
# @RETURN: List[str] - Unique source names.
|
|
# @RELATION: [CALLS] ->[TaskLogPersistenceService.get_sources]
|
|
# @RELATION: [DEPENDS_ON] ->[EventBus]
|
|
def get_task_log_sources(self, task_id: str) -> List[str]:
|
|
with belief_scope("TaskManager.get_task_log_sources", f"task_id={task_id}"):
|
|
return self.log_persistence_service.get_sources(task_id)
|
|
|
|
# [/DEF:get_task_log_sources:Function]
|
|
|
|
# [DEF:_add_log:Function]
|
|
# @COMPLEXITY: 3
|
|
# @PURPOSE: Adds a log entry to a task buffer and notifies subscribers.
|
|
# @PRE: Task exists.
|
|
# @POST: Log added to buffer and pushed to queues (if level meets task_log_level filter).
|
|
# @PARAM: task_id (str) - ID of the task.
|
|
# @PARAM: level (str) - Log level.
|
|
# @PARAM: message (str) - Log message.
|
|
# @PARAM: source (str) - Source component (default: "system").
|
|
# @PARAM: metadata (Optional[Dict]) - Additional structured data.
|
|
# @PARAM: context (Optional[Dict]) - Legacy context (for backward compatibility).
|
|
# @RELATION: [CALLS] ->[should_log_task_level]
|
|
# @RELATION: [DISPATCHES] ->[EventBus]
|
|
def _add_log(
|
|
self,
|
|
task_id: str,
|
|
level: str,
|
|
message: str,
|
|
source: str = "system",
|
|
metadata: Optional[Dict[str, Any]] = None,
|
|
context: Optional[Dict[str, Any]] = None,
|
|
):
|
|
with belief_scope("TaskManager._add_log", f"task_id={task_id}"):
|
|
task = self.tasks.get(task_id)
|
|
if not task:
|
|
return
|
|
|
|
# Filter logs based on task_log_level configuration
|
|
if not should_log_task_level(level):
|
|
return
|
|
|
|
# Create log entry with new fields
|
|
log_entry = LogEntry(
|
|
level=level,
|
|
message=message,
|
|
source=source,
|
|
metadata=metadata,
|
|
context=context, # Keep for backward compatibility
|
|
)
|
|
|
|
# Add to in-memory logs (for backward compatibility with legacy JSON field)
|
|
task.logs.append(log_entry)
|
|
|
|
# Add to buffer for batch persistence
|
|
with self._log_buffer_lock:
|
|
if task_id not in self._log_buffer:
|
|
self._log_buffer[task_id] = []
|
|
self._log_buffer[task_id].append(log_entry)
|
|
|
|
# Notify subscribers (for real-time WebSocket updates)
|
|
if task_id in self.subscribers:
|
|
for queue in self.subscribers[task_id]:
|
|
self.loop.call_soon_threadsafe(queue.put_nowait, log_entry)
|
|
|
|
# [/DEF:_add_log:Function]
|
|
|
|
# [DEF:subscribe_logs:Function]
|
|
# @COMPLEXITY: 3
|
|
# @PURPOSE: Subscribes to real-time logs for a task.
|
|
# @PRE: task_id is a string.
|
|
# @POST: Returns an asyncio.Queue for log entries.
|
|
# @PARAM: task_id (str) - ID of the task.
|
|
# @RETURN: asyncio.Queue - Queue for log entries.
|
|
# @RELATION: [DEPENDS_ON] ->[EventBus]
|
|
async def subscribe_logs(self, task_id: str) -> asyncio.Queue:
|
|
with belief_scope("TaskManager.subscribe_logs", f"task_id={task_id}"):
|
|
queue = asyncio.Queue()
|
|
if task_id not in self.subscribers:
|
|
self.subscribers[task_id] = []
|
|
self.subscribers[task_id].append(queue)
|
|
return queue
|
|
|
|
# [/DEF:subscribe_logs:Function]
|
|
|
|
# [DEF:unsubscribe_logs:Function]
|
|
# @COMPLEXITY: 3
|
|
# @PURPOSE: Unsubscribes from real-time logs for a task.
|
|
# @PRE: task_id is a string, queue is asyncio.Queue.
|
|
# @POST: Queue removed from subscribers.
|
|
# @PARAM: task_id (str) - ID of the task.
|
|
# @PARAM: queue (asyncio.Queue) - Queue to remove.
|
|
# @RELATION: [DEPENDS_ON] ->[EventBus]
|
|
def unsubscribe_logs(self, task_id: str, queue: asyncio.Queue):
|
|
with belief_scope("TaskManager.unsubscribe_logs", f"task_id={task_id}"):
|
|
if task_id in self.subscribers:
|
|
if queue in self.subscribers[task_id]:
|
|
self.subscribers[task_id].remove(queue)
|
|
if not self.subscribers[task_id]:
|
|
del self.subscribers[task_id]
|
|
|
|
# [/DEF:unsubscribe_logs:Function]
|
|
|
|
# [DEF:load_persisted_tasks:Function]
|
|
# @COMPLEXITY: 3
|
|
# @PURPOSE: Load persisted tasks using persistence service.
|
|
# @PRE: None.
|
|
# @POST: Persisted tasks loaded into self.tasks.
|
|
# @RELATION: [CALLS] ->[TaskPersistenceService.load_tasks]
|
|
# @RELATION: [DEPENDS_ON] ->[TaskGraph]
|
|
def load_persisted_tasks(self) -> None:
|
|
with belief_scope("TaskManager.load_persisted_tasks"):
|
|
loaded_tasks = self.persistence_service.load_tasks(limit=100)
|
|
for task in loaded_tasks:
|
|
if task.id not in self.tasks:
|
|
self.tasks[task.id] = task
|
|
|
|
# [/DEF:load_persisted_tasks:Function]
|
|
|
|
# [DEF:await_input:Function]
|
|
# @COMPLEXITY: 3
|
|
# @PURPOSE: Transition a task to AWAITING_INPUT state with input request.
|
|
# @PRE: Task exists and is in RUNNING state.
|
|
# @POST: Task status changed to AWAITING_INPUT, input_request set, persisted.
|
|
# @PARAM: task_id (str) - ID of the task.
|
|
# @PARAM: input_request (Dict) - Details about required input.
|
|
# @THROWS: ValueError if task not found or not RUNNING.
|
|
# @RELATION: [CALLS] ->[TaskPersistenceService.persist_task]
|
|
# @RELATION: [DEPENDS_ON] ->[JobLifecycle]
|
|
def await_input(self, task_id: str, input_request: Dict[str, Any]) -> None:
|
|
with belief_scope("TaskManager.await_input", f"task_id={task_id}"):
|
|
task = self.tasks.get(task_id)
|
|
if not task:
|
|
raise ValueError(f"Task {task_id} not found")
|
|
if task.status != TaskStatus.RUNNING:
|
|
raise ValueError(
|
|
f"Task {task_id} is not RUNNING (current: {task.status})"
|
|
)
|
|
|
|
task.status = TaskStatus.AWAITING_INPUT
|
|
task.input_required = True
|
|
task.input_request = input_request
|
|
self.persistence_service.persist_task(task)
|
|
self._add_log(
|
|
task_id,
|
|
"INFO",
|
|
"Task paused for user input",
|
|
metadata={"input_request": input_request},
|
|
)
|
|
|
|
# [/DEF:await_input:Function]
|
|
|
|
# [DEF:resume_task_with_password:Function]
|
|
# @COMPLEXITY: 3
|
|
# @PURPOSE: Resume a task that is awaiting input with provided passwords.
|
|
# @PRE: Task exists and is in AWAITING_INPUT state.
|
|
# @POST: Task status changed to RUNNING, passwords injected, task resumed.
|
|
# @PARAM: task_id (str) - ID of the task.
|
|
# @PARAM: passwords (Dict[str, str]) - Mapping of database name to password.
|
|
# @THROWS: ValueError if task not found, not awaiting input, or passwords invalid.
|
|
# @RELATION: [CALLS] ->[TaskPersistenceService.persist_task]
|
|
# @RELATION: [DEPENDS_ON] ->[JobLifecycle]
|
|
def resume_task_with_password(
|
|
self, task_id: str, passwords: Dict[str, str]
|
|
) -> None:
|
|
with belief_scope(
|
|
"TaskManager.resume_task_with_password", f"task_id={task_id}"
|
|
):
|
|
task = self.tasks.get(task_id)
|
|
if not task:
|
|
raise ValueError(f"Task {task_id} not found")
|
|
if task.status != TaskStatus.AWAITING_INPUT:
|
|
raise ValueError(
|
|
f"Task {task_id} is not AWAITING_INPUT (current: {task.status})"
|
|
)
|
|
|
|
if not isinstance(passwords, dict) or not passwords:
|
|
raise ValueError("Passwords must be a non-empty dictionary")
|
|
|
|
task.params["passwords"] = passwords
|
|
task.input_required = False
|
|
task.input_request = None
|
|
task.status = TaskStatus.RUNNING
|
|
self.persistence_service.persist_task(task)
|
|
self._add_log(
|
|
task_id,
|
|
"INFO",
|
|
"Task resumed with passwords",
|
|
metadata={"databases": list(passwords.keys())},
|
|
)
|
|
|
|
if task_id in self.task_futures:
|
|
self.task_futures[task_id].set_result(True)
|
|
|
|
# [/DEF:resume_task_with_password:Function]
|
|
|
|
# [DEF:clear_tasks:Function]
|
|
# @COMPLEXITY: 3
|
|
# @PURPOSE: Clears tasks based on status filter (also deletes associated logs).
|
|
# @PRE: status is Optional[TaskStatus].
|
|
# @POST: Tasks matching filter (or all non-active) cleared from registry and database.
|
|
# @PARAM: status (Optional[TaskStatus]) - Filter by task status.
|
|
# @RETURN: int - Number of tasks cleared.
|
|
# @RELATION: [CALLS] ->[TaskPersistenceService.delete_tasks]
|
|
# @RELATION: [DEPENDS_ON] ->[TaskGraph]
|
|
# @RELATION: [DEPENDS_ON] ->[EventBus]
|
|
def clear_tasks(self, status: Optional[TaskStatus] = None) -> int:
|
|
with belief_scope("TaskManager.clear_tasks"):
|
|
tasks_to_remove = []
|
|
for task_id, task in list(self.tasks.items()):
|
|
# If status is provided, match it.
|
|
# If status is None, match everything EXCEPT RUNNING (unless they are awaiting input/mapping which are technically running but paused?)
|
|
# Actually, AWAITING_INPUT and AWAITING_MAPPING are distinct statuses in TaskStatus enum.
|
|
# RUNNING is active execution.
|
|
|
|
should_remove = False
|
|
if status:
|
|
if task.status == status:
|
|
should_remove = True
|
|
else:
|
|
# Clear all non-active tasks (keep RUNNING, AWAITING_INPUT, AWAITING_MAPPING)
|
|
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)
|
|
|
|
for tid in tasks_to_remove:
|
|
# Cancel future if exists (e.g. for AWAITING_INPUT/MAPPING)
|
|
if tid in self.task_futures:
|
|
self.task_futures[tid].cancel()
|
|
del self.task_futures[tid]
|
|
|
|
del self.tasks[tid]
|
|
|
|
# Remove from persistence (task_records and task_logs via CASCADE)
|
|
self.persistence_service.delete_tasks(tasks_to_remove)
|
|
|
|
# Also explicitly delete logs (in case CASCADE is not set up)
|
|
if tasks_to_remove:
|
|
self.log_persistence_service.delete_logs_for_tasks(tasks_to_remove)
|
|
|
|
logger.info(f"Cleared {len(tasks_to_remove)} tasks.")
|
|
return len(tasks_to_remove)
|
|
|
|
# [/DEF:clear_tasks:Function]
|
|
|
|
|
|
# [/DEF:TaskManager:Class]
|
|
# [/DEF:TaskManagerModule:Module]
|