busya
7760214170
- test_agent_handler: additional edge cases for streaming, HITL, file upload - test_confirmations: HITL confirm/deny lifecycle coverage - test_conversation_api: conversation save/load persistence tests - test_langchain_tools: tool registration, dual-auth header propagation - ConversationList.test.ts (frontend): conversation list component tests - conftest: shared fixtures for agent tests - task_manager/manager: minor fixes from test coverage - tasks.md/test-documentation.md: spec and test documentation updates - speckit.test.md: speckit workflow documentation update
431 lines
18 KiB
Python
431 lines
18 KiB
Python
# #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
|