websocket add
This commit is contained in:
@@ -350,6 +350,17 @@ async def start_maintenance(
|
||||
extra={"event_id": event.id, "task_id": task.id},
|
||||
)
|
||||
|
||||
# ── Broadcast maintenance event ──
|
||||
task_manager.broadcast_maintenance_event({
|
||||
"type": "maintenance.event_created",
|
||||
"maintenance_id": event.id,
|
||||
"tables": sorted_tables,
|
||||
"start_time": event.start_time.isoformat() if event.start_time else None,
|
||||
"end_time": event.end_time.isoformat() if event.end_time else None,
|
||||
"message": message,
|
||||
"environment_id": env_id,
|
||||
})
|
||||
|
||||
return MaintenanceStartResponse(
|
||||
task_id=task.id,
|
||||
maintenance_id=event.id,
|
||||
@@ -421,6 +432,12 @@ async def end_maintenance(
|
||||
extra={"event_id": maintenance_id, "task_id": task.id},
|
||||
)
|
||||
|
||||
# ── Broadcast maintenance event ──
|
||||
task_manager.broadcast_maintenance_event({
|
||||
"type": "maintenance.event_ended",
|
||||
"maintenance_id": maintenance_id,
|
||||
})
|
||||
|
||||
return MaintenanceEndResponse(
|
||||
task_id=task.id,
|
||||
status="pending",
|
||||
@@ -485,6 +502,12 @@ async def end_all_maintenance(
|
||||
extra={"task_id": task.id, "environment_id": effective_env_id},
|
||||
)
|
||||
|
||||
# ── Broadcast maintenance event ──
|
||||
task_manager.broadcast_maintenance_event({
|
||||
"type": "maintenance.events_ended_all",
|
||||
"environment_id": effective_env_id or "",
|
||||
})
|
||||
|
||||
return MaintenanceEndResponse(
|
||||
task_id=task.id,
|
||||
status="pending",
|
||||
|
||||
@@ -470,7 +470,10 @@ async def websocket_endpoint(
|
||||
websocket: WebSocket, task_id: str, source: str = None, level: str = None
|
||||
):
|
||||
"""
|
||||
WebSocket endpoint for real-time log streaming with optional server-side filtering.
|
||||
WebSocket endpoint for real-time log streaming AND task status updates.
|
||||
Sends two message types:
|
||||
- Log entries: plain dicts with level/message/timestamp (backward compatible, no type field)
|
||||
- Status updates: dict with type="task_status" and nested task dict
|
||||
Query Parameters:
|
||||
source: Filter logs by source component (e.g., "plugin", "superset_api")
|
||||
level: Filter logs by minimum level (DEBUG, INFO, WARNING, ERROR)
|
||||
@@ -490,7 +493,7 @@ async def websocket_endpoint(
|
||||
level_hierarchy = {"DEBUG": 0, "INFO": 1, "WARNING": 2, "ERROR": 3}
|
||||
min_level = level_hierarchy.get(level_filter, 0) if level_filter else 0
|
||||
logger.reason(
|
||||
"Accepted WebSocket log stream connection",
|
||||
"Accepted WebSocket log+status stream connection",
|
||||
extra={
|
||||
"task_id": task_id,
|
||||
"source_filter": source_filter,
|
||||
@@ -499,11 +502,13 @@ async def websocket_endpoint(
|
||||
},
|
||||
)
|
||||
task_manager = get_task_manager()
|
||||
queue = await task_manager.subscribe_logs(task_id)
|
||||
log_queue = await task_manager.subscribe_logs(task_id)
|
||||
status_queue = await task_manager.subscribe_status(task_id)
|
||||
logger.reason(
|
||||
"Subscribed WebSocket client to task log queue",
|
||||
"Subscribed WebSocket client to task log and status queues",
|
||||
extra={"task_id": task_id},
|
||||
)
|
||||
|
||||
def matches_filters(log_entry) -> bool:
|
||||
"""Check if log entry matches the filter criteria."""
|
||||
log_source = getattr(log_entry, "source", None)
|
||||
@@ -514,7 +519,37 @@ async def websocket_endpoint(
|
||||
if log_level < min_level:
|
||||
return False
|
||||
return True
|
||||
|
||||
async def send_status_update(task: Any) -> None:
|
||||
"""Send a structured task status update over the WebSocket."""
|
||||
status_dict = {
|
||||
"id": task.id,
|
||||
"plugin_id": task.plugin_id,
|
||||
"status": task.status.value if hasattr(task.status, "value") else str(task.status),
|
||||
"started_at": task.started_at.isoformat() if task.started_at else None,
|
||||
"finished_at": task.finished_at.isoformat() if task.finished_at else None,
|
||||
"user_id": task.user_id,
|
||||
"result": task.result,
|
||||
"input_required": task.input_required,
|
||||
"input_request": task.input_request,
|
||||
}
|
||||
await websocket.send_json({
|
||||
"type": "task_status",
|
||||
"task_id": task.id,
|
||||
"task": status_dict,
|
||||
})
|
||||
|
||||
try:
|
||||
# ── Send initial task status ──
|
||||
task = task_manager.get_task(task_id)
|
||||
if task:
|
||||
await send_status_update(task)
|
||||
logger.reason(
|
||||
"Sent initial task status",
|
||||
extra={"task_id": task_id, "status": str(task.status)},
|
||||
)
|
||||
|
||||
# ── Replay initial logs ──
|
||||
logger.reason(
|
||||
"Starting task log stream replay and live forwarding",
|
||||
extra={"task_id": task_id},
|
||||
@@ -535,6 +570,8 @@ async def websocket_endpoint(
|
||||
"total_available_logs": len(initial_logs),
|
||||
},
|
||||
)
|
||||
|
||||
# ── Send synthetic AWAITING_INPUT prompt if needed ──
|
||||
task = task_manager.get_task(task_id)
|
||||
if task and task.status == "AWAITING_INPUT" and task.input_request:
|
||||
synthetic_log = {
|
||||
@@ -550,47 +587,167 @@ async def websocket_endpoint(
|
||||
"Replayed awaiting-input prompt to restored WebSocket client",
|
||||
extra={"task_id": task_id, "task_status": task.status},
|
||||
)
|
||||
|
||||
# ── Main loop: listen on both log and status queues ──
|
||||
while True:
|
||||
log_entry = await queue.get()
|
||||
if not matches_filters(log_entry):
|
||||
continue
|
||||
log_dict = log_entry.model_dump()
|
||||
log_dict["timestamp"] = log_dict["timestamp"].isoformat()
|
||||
await websocket.send_json(log_dict)
|
||||
logger.reflect(
|
||||
"Forwarded task log entry to WebSocket client",
|
||||
extra={
|
||||
"task_id": task_id,
|
||||
"level": log_dict.get("level"),
|
||||
},
|
||||
done, _ = await asyncio.wait(
|
||||
[log_queue.get(), status_queue.get()],
|
||||
return_when=asyncio.FIRST_COMPLETED,
|
||||
)
|
||||
if (
|
||||
"Task completed successfully" in log_entry.message
|
||||
or "Task failed" in log_entry.message
|
||||
):
|
||||
logger.reason(
|
||||
"Observed terminal task log entry; delaying to preserve client visibility",
|
||||
extra={"task_id": task_id, "message": log_entry.message},
|
||||
for coro in done:
|
||||
result = coro.result()
|
||||
|
||||
# ── Status update ──
|
||||
if isinstance(result, dict) and result.get("type") == "task_status":
|
||||
await websocket.send_json(result)
|
||||
task_status = result.get("task", {}).get("status", "")
|
||||
if task_status in ("SUCCESS", "FAILED"):
|
||||
logger.reason(
|
||||
"Task reached terminal state via status broadcast; closing stream",
|
||||
extra={"task_id": task_id, "status": task_status},
|
||||
)
|
||||
await asyncio.sleep(2)
|
||||
raise StopIteration # exit the while loop
|
||||
continue
|
||||
|
||||
# ── Log entry ──
|
||||
if not matches_filters(result):
|
||||
continue
|
||||
log_dict = result.model_dump()
|
||||
log_dict["timestamp"] = log_dict["timestamp"].isoformat()
|
||||
await websocket.send_json(log_dict)
|
||||
logger.reflect(
|
||||
"Forwarded task log entry to WebSocket client",
|
||||
extra={
|
||||
"task_id": task_id,
|
||||
"level": log_dict.get("level"),
|
||||
},
|
||||
)
|
||||
await asyncio.sleep(2)
|
||||
except WebSocketDisconnect:
|
||||
if (
|
||||
"Task completed successfully" in result.message
|
||||
or "Task failed" in result.message
|
||||
):
|
||||
logger.reason(
|
||||
"Observed terminal task log entry; delaying to preserve client visibility",
|
||||
extra={"task_id": task_id, "message": result.message},
|
||||
)
|
||||
await asyncio.sleep(2)
|
||||
except (WebSocketDisconnect, StopIteration):
|
||||
if isinstance(StopIteration):
|
||||
pass # normal termination
|
||||
logger.reason(
|
||||
"WebSocket client disconnected from task log stream",
|
||||
"WebSocket client disconnected or stream ended",
|
||||
extra={"task_id": task_id},
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.explore(
|
||||
"WebSocket log streaming encountered an unexpected failure",
|
||||
"WebSocket log+status streaming encountered an unexpected failure",
|
||||
extra={"task_id": task_id, "error": str(exc)},
|
||||
)
|
||||
raise
|
||||
finally:
|
||||
task_manager.unsubscribe_logs(task_id, queue)
|
||||
task_manager.unsubscribe_logs(task_id, log_queue)
|
||||
task_manager.unsubscribe_status(task_id, status_queue)
|
||||
logger.reflect(
|
||||
"Released WebSocket log queue subscription",
|
||||
"Released WebSocket log and status queue subscriptions",
|
||||
extra={"task_id": task_id},
|
||||
)
|
||||
# #endregion websocket_endpoint
|
||||
|
||||
# #region task_events_websocket [C:4] [TYPE Function]
|
||||
# @BRIEF WebSocket endpoint for global task events (status changes for ALL tasks).
|
||||
# @RELATION CALLS -> [TaskManagerPackage]
|
||||
# @PRE WebSocket must be authenticated via `token` query param.
|
||||
# @POST WebSocket streams task status events until disconnect.
|
||||
@app.websocket("/ws/task-events")
|
||||
async def task_events_websocket(websocket: WebSocket):
|
||||
"""
|
||||
WebSocket endpoint for global task events.
|
||||
Streams {type: "task_status", task_id: ..., task: {...}} for ALL task status changes.
|
||||
Query Parameters:
|
||||
token: JWT or API key for authentication (required)
|
||||
"""
|
||||
seed_trace_id()
|
||||
with belief_scope("task_events_websocket"):
|
||||
if not await _authenticate_websocket(websocket, "ws/task-events"):
|
||||
await websocket.close(code=4001, reason="Authentication required")
|
||||
return
|
||||
|
||||
await websocket.accept()
|
||||
logger.reason("Accepted global task events WebSocket connection")
|
||||
|
||||
task_manager = get_task_manager()
|
||||
event_queue = await task_manager.subscribe_task_events()
|
||||
logger.reason("Subscribed to global task events")
|
||||
|
||||
try:
|
||||
while True:
|
||||
event = await event_queue.get()
|
||||
await websocket.send_json(event)
|
||||
logger.reflect(
|
||||
"Forwarded task event to global client",
|
||||
extra={"task_id": event.get("task_id")},
|
||||
)
|
||||
except WebSocketDisconnect:
|
||||
logger.reason("Global task events client disconnected")
|
||||
except Exception as exc:
|
||||
logger.explore(
|
||||
"Global task events streaming failed",
|
||||
extra={"error": str(exc)},
|
||||
)
|
||||
raise
|
||||
finally:
|
||||
task_manager.unsubscribe_task_events(event_queue)
|
||||
logger.reflect("Released global task events subscription")
|
||||
# #endregion task_events_websocket
|
||||
|
||||
# #region maintenance_events_websocket [C:4] [TYPE Function]
|
||||
# @BRIEF WebSocket endpoint for maintenance events (created/ended/banner changes).
|
||||
# @RELATION CALLS -> [TaskManagerPackage]
|
||||
# @PRE WebSocket must be authenticated via `token` query param.
|
||||
# @POST WebSocket streams maintenance events until disconnect.
|
||||
@app.websocket("/ws/maintenance/events")
|
||||
async def maintenance_events_websocket(websocket: WebSocket):
|
||||
"""
|
||||
WebSocket endpoint for maintenance events.
|
||||
Streams {type: "maintenance.event_created", ...} / {type: "maintenance.event_ended"}.
|
||||
Query Parameters:
|
||||
token: JWT or API key for authentication (required)
|
||||
"""
|
||||
seed_trace_id()
|
||||
with belief_scope("maintenance_events_websocket"):
|
||||
if not await _authenticate_websocket(websocket, "ws/maintenance/events"):
|
||||
await websocket.close(code=4001, reason="Authentication required")
|
||||
return
|
||||
|
||||
await websocket.accept()
|
||||
logger.reason("Accepted maintenance events WebSocket connection")
|
||||
|
||||
task_manager = get_task_manager()
|
||||
event_queue = await task_manager.subscribe_maintenance_events()
|
||||
logger.reason("Subscribed to maintenance events")
|
||||
|
||||
try:
|
||||
while True:
|
||||
event = await event_queue.get()
|
||||
await websocket.send_json(event)
|
||||
logger.reflect(
|
||||
"Forwarded maintenance event to client",
|
||||
extra={"event_type": event.get("type"), "maintenance_id": event.get("maintenance_id")},
|
||||
)
|
||||
except WebSocketDisconnect:
|
||||
logger.reason("Maintenance events client disconnected")
|
||||
except Exception as exc:
|
||||
logger.explore(
|
||||
"Maintenance events streaming failed",
|
||||
extra={"error": str(exc)},
|
||||
)
|
||||
raise
|
||||
finally:
|
||||
task_manager.unsubscribe_maintenance_events(event_queue)
|
||||
logger.reflect("Released maintenance events subscription")
|
||||
# #endregion maintenance_events_websocket
|
||||
|
||||
# #region dataset_websocket_endpoint [C:4] [TYPE Function]
|
||||
# @BRIEF WebSocket endpoint for dataset.updated events — auto-refresh on task completion.
|
||||
# @RELATION CALLS -> [TaskManagerPackage]
|
||||
|
||||
@@ -49,11 +49,17 @@ class EventBus:
|
||||
LOG_FLUSH_INTERVAL = 2.0
|
||||
|
||||
# #region __init__ [TYPE Function]
|
||||
# @BRIEF Initialize empty log buffer, subscriber registry, and start background flusher.
|
||||
# @BRIEF Initialize empty log buffer, subscriber registry, status subscribers, and start background flusher.
|
||||
def __init__(self, log_persistence_service: TaskLogPersistenceService, loop: asyncio.AbstractEventLoop):
|
||||
self.log_persistence_service = log_persistence_service
|
||||
self.loop = loop
|
||||
self.subscribers: dict[str, list[asyncio.Queue]] = {}
|
||||
# Per-task status subscribers (same pattern as log subscribers)
|
||||
self._status_subscribers: dict[str, list[asyncio.Queue]] = {}
|
||||
# Global task event subscribers (all tasks)
|
||||
self._task_event_subscribers: list[asyncio.Queue] = []
|
||||
# Maintenance event subscribers
|
||||
self._maintenance_subscribers: list[asyncio.Queue] = []
|
||||
self._log_buffer: dict[str, list[LogEntry]] = {}
|
||||
self._log_buffer_lock = threading.Lock()
|
||||
self._flusher_stop_event = threading.Event()
|
||||
@@ -178,6 +184,101 @@ class EventBus:
|
||||
del self.subscribers[task_id]
|
||||
# #endregion unsubscribe_logs
|
||||
|
||||
# ── Task Status Subscribers ──
|
||||
|
||||
# #region subscribe_status [C:2] [TYPE Function] [SEMANTICS subscribe,status,realtime]
|
||||
# @BRIEF Subscribes to real-time status updates for a specific task.
|
||||
# @POST Returns an asyncio.Queue for Task status dicts.
|
||||
async def subscribe_status(self, task_id: str) -> asyncio.Queue:
|
||||
queue = asyncio.Queue()
|
||||
if task_id not in self._status_subscribers:
|
||||
self._status_subscribers[task_id] = []
|
||||
self._status_subscribers[task_id].append(queue)
|
||||
return queue
|
||||
# #endregion subscribe_status
|
||||
|
||||
# #region unsubscribe_status [C:2] [TYPE Function] [SEMANTICS unsubscribe,status,remove]
|
||||
# @BRIEF Unsubscribes from real-time status updates for a task.
|
||||
# @POST Queue removed from status subscribers.
|
||||
def unsubscribe_status(self, task_id: str, queue: asyncio.Queue):
|
||||
if task_id in self._status_subscribers:
|
||||
if queue in self._status_subscribers[task_id]:
|
||||
self._status_subscribers[task_id].remove(queue)
|
||||
if not self._status_subscribers[task_id]:
|
||||
del self._status_subscribers[task_id]
|
||||
# #endregion unsubscribe_status
|
||||
|
||||
# #region broadcast_status [C:2] [TYPE Function] [SEMANTICS broadcast,status,notify]
|
||||
# @BRIEF Broadcast a task status update to all per-task and global subscribers.
|
||||
# @POST Status event pushed to all matching queues.
|
||||
def broadcast_status(self, task_id: str, task_dict: dict):
|
||||
from ..logger import logger as _lg
|
||||
event = {"type": "task_status", "task_id": task_id, "task": task_dict}
|
||||
# Per-task subscribers
|
||||
if task_id in self._status_subscribers:
|
||||
for queue in self._status_subscribers[task_id]:
|
||||
try:
|
||||
self.loop.call_soon_threadsafe(queue.put_nowait, event)
|
||||
except Exception as _e:
|
||||
_lg.warning("Failed to send status to per-task subscriber", extra={"task_id": task_id, "error": str(_e)})
|
||||
# Global task event subscribers
|
||||
for queue in self._task_event_subscribers:
|
||||
try:
|
||||
self.loop.call_soon_threadsafe(queue.put_nowait, event)
|
||||
except Exception as _e:
|
||||
_lg.warning("Failed to send status to global subscriber", extra={"task_id": task_id, "error": str(_e)})
|
||||
# #endregion broadcast_status
|
||||
|
||||
# ── Maintenance Event Subscribers ──
|
||||
|
||||
# #region subscribe_maintenance_events [C:2] [TYPE Function] [SEMANTICS subscribe,maintenance,events]
|
||||
# @BRIEF Subscribes to maintenance event updates (created/ended/banner changes).
|
||||
# @POST Returns an asyncio.Queue for maintenance event dicts.
|
||||
async def subscribe_maintenance_events(self) -> asyncio.Queue:
|
||||
queue = asyncio.Queue()
|
||||
self._maintenance_subscribers.append(queue)
|
||||
return queue
|
||||
# #endregion subscribe_maintenance_events
|
||||
|
||||
# #region unsubscribe_maintenance_events [C:2] [TYPE Function] [SEMANTICS unsubscribe,maintenance,events]
|
||||
# @BRIEF Unsubscribes from maintenance events.
|
||||
# @POST Queue removed from maintenance subscribers.
|
||||
def unsubscribe_maintenance_events(self, queue: asyncio.Queue):
|
||||
if queue in self._maintenance_subscribers:
|
||||
self._maintenance_subscribers.remove(queue)
|
||||
# #endregion unsubscribe_maintenance_events
|
||||
|
||||
# #region broadcast_maintenance_event [C:2] [TYPE Function] [SEMANTICS broadcast,maintenance,notify]
|
||||
# @BRIEF Broadcast a maintenance event to all subscribers.
|
||||
# @POST Event pushed to all maintenance subscriber queues.
|
||||
def broadcast_maintenance_event(self, event: dict):
|
||||
from ..logger import logger as _lg
|
||||
for queue in self._maintenance_subscribers:
|
||||
try:
|
||||
self.loop.call_soon_threadsafe(queue.put_nowait, event)
|
||||
except Exception as _e:
|
||||
_lg.warning("Failed to send maintenance event to subscriber", extra={"event_type": event.get("type"), "error": str(_e)})
|
||||
# #endregion broadcast_maintenance_event
|
||||
|
||||
# ── Global Task Event Subscribers ──
|
||||
|
||||
# #region subscribe_task_events [C:2] [TYPE Function] [SEMANTICS subscribe,global,events]
|
||||
# @BRIEF Subscribes to ALL task events (status changes) globally.
|
||||
# @POST Returns an asyncio.Queue for task event dicts.
|
||||
async def subscribe_task_events(self) -> asyncio.Queue:
|
||||
queue = asyncio.Queue()
|
||||
self._task_event_subscribers.append(queue)
|
||||
return queue
|
||||
# #endregion subscribe_task_events
|
||||
|
||||
# #region unsubscribe_task_events [C:2] [TYPE Function] [SEMANTICS unsubscribe,global,events]
|
||||
# @BRIEF Unsubscribes from global task events.
|
||||
# @POST Queue removed from global subscribers.
|
||||
def unsubscribe_task_events(self, queue: asyncio.Queue):
|
||||
if queue in self._task_event_subscribers:
|
||||
self._task_event_subscribers.remove(queue)
|
||||
# #endregion unsubscribe_task_events
|
||||
|
||||
# #region get_task_logs [C:3] [TYPE Function] [SEMANTICS logs,read,persistence,backfill]
|
||||
# @BRIEF Retrieves logs for a specific task (from memory for running, persistence for completed).
|
||||
# @RELATION CALLS -> [EXT:method:TaskLogPersistenceService.get_logs]
|
||||
|
||||
@@ -34,6 +34,22 @@ from .context import TaskContext
|
||||
from .graph import TaskGraph
|
||||
from .models import Task, TaskStatus
|
||||
|
||||
def _task_to_dict(task: Task) -> dict:
|
||||
"""Serialize a Task to a dict for WebSocket broadcast."""
|
||||
return {
|
||||
"id": task.id,
|
||||
"plugin_id": task.plugin_id,
|
||||
"status": task.status.value if isinstance(task.status, TaskStatus) else task.status,
|
||||
"started_at": task.started_at.isoformat() if task.started_at else None,
|
||||
"finished_at": task.finished_at.isoformat() if task.finished_at else None,
|
||||
"user_id": task.user_id,
|
||||
"params": task.params,
|
||||
"result": task.result,
|
||||
"input_required": task.input_required,
|
||||
"input_request": task.input_request,
|
||||
"error": str(task.result) if task.status == TaskStatus.FAILED and task.result else None,
|
||||
}
|
||||
|
||||
|
||||
# #region JobLifecycle [C:5] [TYPE Class] [SEMANTICS task,lifecycle,execution,state,machine]
|
||||
# @BRIEF Encodes task creation, execution, pause/resume, and completion transitions for
|
||||
@@ -129,6 +145,7 @@ class JobLifecycle:
|
||||
task.status = TaskStatus.RUNNING
|
||||
task.started_at = datetime.now(UTC)
|
||||
self.persistence_service.persist_task(task)
|
||||
self._broadcast_task_status(task)
|
||||
if add_log_callback:
|
||||
add_log_callback(
|
||||
task_id, "INFO",
|
||||
@@ -164,6 +181,7 @@ class JobLifecycle:
|
||||
)
|
||||
logger.info(f"Task {task_id} completed successfully")
|
||||
task.status = TaskStatus.SUCCESS
|
||||
self._broadcast_task_status(task)
|
||||
if add_log_callback:
|
||||
add_log_callback(
|
||||
task_id, "INFO",
|
||||
@@ -173,6 +191,7 @@ class JobLifecycle:
|
||||
except Exception as e:
|
||||
logger.error(f"Task {task_id} failed: {e}")
|
||||
task.status = TaskStatus.FAILED
|
||||
self._broadcast_task_status(task)
|
||||
if add_log_callback:
|
||||
add_log_callback(
|
||||
task_id, "ERROR",
|
||||
@@ -214,6 +233,7 @@ class JobLifecycle:
|
||||
task.params.update(resolution_params)
|
||||
task.status = TaskStatus.RUNNING
|
||||
self.persistence_service.persist_task(task)
|
||||
self._broadcast_task_status(task)
|
||||
self.graph.resolve_future(task_id, True)
|
||||
# #endregion resolve_task
|
||||
|
||||
@@ -228,6 +248,7 @@ class JobLifecycle:
|
||||
return
|
||||
task.status = TaskStatus.AWAITING_MAPPING
|
||||
self.persistence_service.persist_task(task)
|
||||
self._broadcast_task_status(task)
|
||||
future = self.loop.create_future()
|
||||
self.graph.create_future(task_id, future)
|
||||
try:
|
||||
@@ -274,6 +295,7 @@ class JobLifecycle:
|
||||
task.input_required = True
|
||||
task.input_request = input_request
|
||||
self.persistence_service.persist_task(task)
|
||||
self._broadcast_task_status(task)
|
||||
if add_log_callback:
|
||||
add_log_callback(
|
||||
task_id, "INFO",
|
||||
@@ -308,6 +330,7 @@ class JobLifecycle:
|
||||
task.input_request = None
|
||||
task.status = TaskStatus.RUNNING
|
||||
self.persistence_service.persist_task(task)
|
||||
self._broadcast_task_status(task)
|
||||
if add_log_callback:
|
||||
add_log_callback(
|
||||
task_id, "INFO",
|
||||
@@ -337,6 +360,13 @@ class JobLifecycle:
|
||||
del self._dataset_subscribers[env_id]
|
||||
# #endregion unsubscribe_dataset_events
|
||||
|
||||
# #region _broadcast_task_status [C:2] [TYPE Function] [SEMANTICS broadcast,status,websocket]
|
||||
# @BRIEF Broadcast task status update to all subscribers via EventBus.
|
||||
# @POST EventBus.broadcast_status called with serialized task dict.
|
||||
def _broadcast_task_status(self, task: Task):
|
||||
self.event_bus.broadcast_status(task.id, _task_to_dict(task))
|
||||
# #endregion _broadcast_task_status
|
||||
|
||||
# #region _broadcast_dataset_updated [C:3] [TYPE Function] [SEMANTICS dataset,broadcast,event,websocket]
|
||||
# @BRIEF Broadcast dataset.updated event to all subscribers for a given environment.
|
||||
def _broadcast_dataset_updated(self, env_id: str, dataset_ids: list[int]):
|
||||
|
||||
@@ -261,6 +261,32 @@ class TaskManager:
|
||||
self.event_bus.unsubscribe_logs(task_id, queue)
|
||||
# #endregion unsubscribe_logs
|
||||
|
||||
# ── Status subscribers ──
|
||||
|
||||
# #region subscribe_status [TYPE Function] [C:2]
|
||||
# @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]
|
||||
# @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]
|
||||
# @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]
|
||||
# @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]
|
||||
@@ -318,6 +344,26 @@ class TaskManager:
|
||||
)
|
||||
# #endregion resume_task_with_password
|
||||
|
||||
# ── Maintenance event delegates to EventBus ──
|
||||
|
||||
# #region subscribe_maintenance_events [TYPE Function] [C:2]
|
||||
# @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]
|
||||
# @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]
|
||||
# @BRIEF Broadcast a maintenance event to all subscribers.
|
||||
def broadcast_maintenance_event(self, event: dict):
|
||||
self.event_bus.broadcast_maintenance_event(event)
|
||||
# #endregion broadcast_maintenance_event
|
||||
|
||||
# ── Dataset event delegates to JobLifecycle ──
|
||||
|
||||
# #region subscribe_dataset_events [TYPE Function] [C:2]
|
||||
|
||||
Reference in New Issue
Block a user