diff --git a/backend/alembic/versions/0019_task_results_index.py b/backend/alembic/versions/0019_task_results_index.py new file mode 100644 index 0000000..0ea3a95 --- /dev/null +++ b/backend/alembic/versions/0019_task_results_index.py @@ -0,0 +1,25 @@ +"""0019 task results composite index + +Revision ID: 0019 +Revises: 0018 +Create Date: 2026-05-06 +""" + +from alembic import op + +revision = "0019" +down_revision = "0018" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + op.create_index( + "ix_task_results_task_device", + "task_results", + ["task_id", "device_id"], + ) + + +def downgrade() -> None: + op.drop_index("ix_task_results_task_device", table_name="task_results") diff --git a/backend/app/services/task.py b/backend/app/services/task.py index c01128d..c1958cc 100644 --- a/backend/app/services/task.py +++ b/backend/app/services/task.py @@ -87,7 +87,6 @@ async def _resolve_scope(scope: dict) -> list[tuple[Device, Object]]: query = query.where(Device.object_id == scope["object_id"]) elif scope_type == "filter": # Cross-object filter scope (from WorkDevices global page) - from app.models.object import Object as Obj if scope.get("category"): cats = scope["category"] if isinstance(cats, list): @@ -96,10 +95,12 @@ async def _resolve_scope(scope: dict) -> list[tuple[Device, Object]]: query = query.where(Device.category == cats) if scope.get("object_id"): query = query.where(Device.object_id == scope["object_id"]) - if scope.get("city"): - query = query.where(Obj.city == scope["city"]) - if scope.get("client"): - query = query.where(Obj.client == scope["client"]) + if scope.get("city") or scope.get("client"): + query = query.join(Object, Device.object_id == Object.id) + if scope.get("city"): + query = query.where(Object.city == scope["city"]) + if scope.get("client"): + query = query.where(Object.client == scope["client"]) else: return [] @@ -117,20 +118,14 @@ async def _run_on_device( sem: asyncio.Semaphore, ) -> bool | None: async with sem: - # Bail out after acquiring the semaphore so large target lists don't - # stampede the DB pool with cancellation checks. - async with AsyncSessionLocal() as db: - task_row = await db.get(Task, task_id) - if task_row and task_row.cancel_requested: - return None # None = skipped due to cancellation - - emit = lambda msg, level="info": _emit(task_id, msg, level) - - # Create task_result row - result_id: int | None = None + # Session 1: cancellation check + create TaskResult started_at = datetime.now(timezone.utc) + result_id: int | None = None async with AsyncSessionLocal() as db: async with db.begin(): + task_row = await db.get(Task, task_id) + if task_row and task_row.cancel_requested: + return None # None = skipped due to cancellation tr = TaskResult( task_id=task_id, device_id=device.id, @@ -141,6 +136,8 @@ async def _run_on_device( await db.flush() result_id = tr.id + emit = lambda msg, level="info": _emit(task_id, msg, level) + try: action_result = await action.execute(device, obj, params, emit, db=None) except Exception as exc: @@ -161,6 +158,7 @@ async def _run_on_device( finished_at = datetime.now(timezone.utc) duration_ms = int((finished_at - started_at).total_seconds() * 1000) + # Session 2: update TaskResult + device status (ping) in one transaction async with AsyncSessionLocal() as db: async with db.begin(): tr = await db.get(TaskResult, result_id) @@ -172,11 +170,8 @@ async def _run_on_device( tr.data = {**(action_data or {}), "error": action_error} if action_error else (action_data or None) tr.duration_ms = duration_ms - # For ping action: update device status, last_seen, last_ping_ms - if action.name == "ping": - new_status = "online" if action_result_success else "offline" - async with AsyncSessionLocal() as db: - async with db.begin(): + if action.name == "ping": + new_status = "online" if action_result_success else "offline" device_row = await db.get(Device, device.id) if device_row: device_row.status = new_status @@ -186,6 +181,10 @@ async def _run_on_device( ping_ms = (action_data or {}).get("ping_ms") if ping_ms is not None: device_row.last_ping_ms = ping_ms + + # For ping action: publish SSE device_status after DB commit + if action.name == "ping": + new_status = "online" if action_result_success else "offline" await sse_manager.publish(SSEEvent( type="device_status", task_id="__monitor__", @@ -217,6 +216,17 @@ async def _run_on_device( return action_result_success +async def _heartbeat_writer(task_id: str, interval: int = 30) -> None: + """Write task.heartbeat_at every `interval` seconds while task is running.""" + while True: + await asyncio.sleep(interval) + async with AsyncSessionLocal() as db: + async with db.begin(): + task = await db.get(Task, task_id) + if task and task.status == "running": + task.heartbeat_at = datetime.now(timezone.utc) + + async def _run_task(task_id: str) -> None: await _emit(task_id, "Task started", "info") @@ -227,6 +237,7 @@ async def _run_task(task_id: str) -> None: return task.status = "running" task.started_at = datetime.now(timezone.utc) + task.heartbeat_at = datetime.now(timezone.utc) scope = dict(task.target_scope) action_type = task.type params = dict(task.params) @@ -275,10 +286,20 @@ async def _run_task(task_id: str) -> None: await _emit(task_id, f"Running on {len(targets)} device(s)...", "info") sem = asyncio.Semaphore(_TASK_DEVICE_CONCURRENCY) - results = await asyncio.gather( - *[_run_on_device(task_id, device, obj, action, params, sem) for device, obj in targets], - return_exceptions=True, - ) + # Start heartbeat writer + heartbeat_task = asyncio.create_task(_heartbeat_writer(task_id)) + + try: + results = await asyncio.gather( + *[_run_on_device(task_id, device, obj, action, params, sem) for device, obj in targets], + return_exceptions=True, + ) + finally: + heartbeat_task.cancel() + try: + await heartbeat_task + except asyncio.CancelledError: + pass ok = sum(1 for r in results if r is True) skipped = sum(1 for r in results if r is None) diff --git a/backend/app/workers/startup.py b/backend/app/workers/startup.py index e57455b..20580d0 100644 --- a/backend/app/workers/startup.py +++ b/backend/app/workers/startup.py @@ -6,6 +6,7 @@ Runs once at app startup: from datetime import datetime, timedelta, timezone +import sqlalchemy as sa from sqlalchemy import update from app.models.task import Task @@ -20,7 +21,15 @@ async def recover_stale_tasks() -> int: async with db.begin(): result = await db.execute( update(Task) - .where(Task.status == "running", Task.started_at < cutoff) + .where( + Task.status == "running", + # Use heartbeat_at when available (more accurate for long tasks), + # fall back to started_at for tasks that never wrote a heartbeat. + sa.or_( + sa.and_(Task.heartbeat_at.isnot(None), Task.heartbeat_at < cutoff), + sa.and_(Task.heartbeat_at.is_(None), Task.started_at < cutoff), + ), + ) .values(status="stale") .returning(Task.id) ) diff --git a/frontend/src/api/tasks.ts b/frontend/src/api/tasks.ts index 539d7c7..66acf03 100644 --- a/frontend/src/api/tasks.ts +++ b/frontend/src/api/tasks.ts @@ -1,5 +1,6 @@ import { useQuery, useMutation, useQueryClient } from '@tanstack/react-query' import { api } from './client' +import { useSSE } from '../hooks/useSSE' export interface TaskEvent { id: number @@ -78,11 +79,50 @@ export const useTask = (id: string | undefined) => queryKey: ['tasks', id], queryFn: () => api.get(`/api/v1/tasks/${id}`).then((r) => r.data), enabled: !!id, - refetchInterval: (query) => { - const status = query.state.data?.status - return status === 'running' || status === 'pending' ? 2000 : false + }) + +// Subscribe to SSE for a running task — merges results in cache, invalidates on completion. +export const useTaskSSE = (id: string | undefined) => { + const qc = useQueryClient() + useSSE(id, { + onTaskResult: (data) => { + qc.setQueryData(['tasks', id], (prev) => { + if (!prev) return prev + const existing = prev.device_results.find((r) => r.device_id === data.device_id) + const updated: TaskResult = existing + ? { ...existing, status: data.status, stdout: data.stdout, stderr: data.stderr, duration_ms: data.duration_ms } + : { + id: 0, + device_id: data.device_id, + device_ip: data.device_ip, + device_hostname: data.device_hostname, + object_id: null, + object_name: null, + rack_id: null, + rack_name: null, + status: data.status, + stdout: data.stdout, + stderr: data.stderr, + exit_code: null, + data: null, + duration_ms: data.duration_ms, + executed_at: null, + } + return { + ...prev, + device_results: existing + ? prev.device_results.map((r) => (r.device_id === data.device_id ? updated : r)) + : [...prev.device_results, updated], + } + }) + }, + onTaskStatus: () => { + // Task finished — do a full refetch to get final status and complete results + qc.invalidateQueries({ queryKey: ['tasks', id] }) + qc.invalidateQueries({ queryKey: ['tasks'] }) }, }) +} export const useActions = () => useQuery({ @@ -133,4 +173,5 @@ export const useDeviceTasks = (deviceId: number) => api .get('/api/v1/tasks', { params: { device_id: deviceId, limit: 100 } }) .then((r) => r.data), + refetchInterval: 10_000, }) diff --git a/frontend/src/hooks/useSSE.ts b/frontend/src/hooks/useSSE.ts index 756bdf8..c246ea5 100644 --- a/frontend/src/hooks/useSSE.ts +++ b/frontend/src/hooks/useSSE.ts @@ -39,8 +39,9 @@ export function useSSE(taskId: string | null | undefined, handlers: SSEHandlers) // ignore parse errors } }, - onerror() { + onerror(err) { ctrl.abort() + throw err }, }) diff --git a/frontend/src/pages/ObjectDetail.tsx b/frontend/src/pages/ObjectDetail.tsx index 658f8e0..5f6fe7e 100644 --- a/frontend/src/pages/ObjectDetail.tsx +++ b/frontend/src/pages/ObjectDetail.tsx @@ -45,7 +45,7 @@ import { } from '../api/objects' import { useRacks, useCreateRack, useUpdateRack, useDeleteRack, type RackItem } from '../api/racks' import { useZones, useCreateZone, useDeleteZone, type ZoneItem } from '../api/zones' -import { useActions, useCreateTask, useTask } from '../api/tasks' +import { useActions, useCreateTask, useTask, useTaskSSE } from '../api/tasks' import { CategoryTag } from '../components/StatusBadge' import { DeviceEditModal } from '../components/DeviceEditModal' import { DiscoveryDrawer } from '../components/DiscoveryDrawer' @@ -90,8 +90,11 @@ function PingButton({ const [taskId, setTaskId] = useState(null) const queryClient = useQueryClient() const createTask = useCreateTask() - const { data: task } = useTask(taskId ?? undefined) + useTaskSSE(taskId ?? undefined) + + // When task finishes (SSE invalidates ['tasks', taskId]), refresh devices and clear taskId + const { data: task } = useTask(taskId ?? undefined) useEffect(() => { if (task && ['success', 'failed', 'partial', 'stale', 'cancelled'].includes(task.status)) { queryClient.invalidateQueries({ queryKey: ['devices', objId] })