Add real-time WebSocket task progress with multi-job drawer

Replace 5s HTTP polling with WebSocket-based real-time updates for task
progress. Celery workers publish progress to Redis pub/sub channels;
a FastAPI WebSocket endpoint subscribes and forwards to the browser.
Polling is kept as a 30s fallback when WebSocket is unavailable.

The task progress drawer now supports multiple concurrent jobs with a
tab bar for switching between scrape and POI distance tasks.

Backend:
- Add services/task_progress_publisher.py (Redis pub/sub bridge)
- Add api/ws_routes.py (WebSocket endpoint with JWT auth)
- Publish progress from listing_tasks and poi_tasks
- Publish REVOKED via pub/sub on cancel/clear to fix stuck UI

Frontend:
- Add useTaskWebSocket hook with reconnection and keepalive
- Add TaskState and WS message types
- TaskIndicator: WS-driven updates with polling fallback
- TaskProgressDrawer: multi-job tabs, POI phase timeline
- Guard against WS overwriting local cancel state
This commit is contained in:
Viktor Barzin 2026-02-09 21:31:45 +00:00
parent 73d19e29d5
commit 8559c4b461
No known key found for this signature in database
GPG key ID: 0EB088298288D958
11 changed files with 774 additions and 72 deletions

View file

@ -3,8 +3,8 @@ import { getStoredPasskeyUser } from '@/auth/passkeyService';
import { fromOidcUser, type AuthUser } from '@/auth/types';
import { POLLING_INTERVALS } from '@/constants';
import { fetchTaskStatus, cancelTask, clearAllTasks } from '@/services';
import { TaskStatus, type TaskResult } from '@/types';
import { useEffect, useState, useRef } from 'react';
import { TaskStatus, type TaskResult, type TaskState } from '@/types';
import { useEffect, useState, useRef, useMemo } from 'react';
import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from './ui/tooltip';
import { Button } from './ui/button';
import { Loader2, CheckCircle2, XCircle, X, Trash2 } from 'lucide-react';
@ -14,9 +14,48 @@ interface TaskIndicatorProps {
taskID: string | null;
onTaskCancelled?: () => void;
onTaskCompleted?: () => void;
wsTasks?: Record<string, TaskState>;
wsConnected?: boolean;
}
export function TaskIndicator({ taskID, onTaskCancelled, onTaskCompleted }: TaskIndicatorProps) {
/** Convert a TaskState (from WS) into a TaskResult (for the drawer). */
function taskStateToResult(ts: TaskState): TaskResult {
return {
progress: ts.progress ?? 0,
processed: ts.processed,
total: ts.total,
phase: ts.phase,
message: ts.message,
subqueries_probed: ts.subqueries_probed,
subqueries_initial: ts.subqueries_initial,
estimated_results: ts.estimated_results,
subqueries_total: ts.subqueries_total,
subqueries_completed: ts.subqueries_completed,
ids_collected: ts.ids_collected,
pages_fetched: ts.pages_fetched,
fetching_done: ts.fetching_done,
details_fetched: ts.details_fetched,
images_downloaded: ts.images_downloaded,
ocr_completed: ts.ocr_completed,
failed: ts.failed,
elapsed_seconds: ts.elapsed_seconds,
rate_per_second: ts.rate_per_second,
eta_seconds: ts.eta_seconds,
logs: ts.logs,
};
}
function isTerminalStatus(status: string): boolean {
return status === 'SUCCESS' || status === 'FAILURE' || status === 'REVOKED';
}
export function TaskIndicator({
taskID,
onTaskCancelled,
onTaskCompleted,
wsTasks,
wsConnected,
}: TaskIndicatorProps) {
const [user, setUser] = useState<AuthUser | null>(null);
const [progressPercentage, setProgressPercentage] = useState<number>(0);
const [processed, setProcessed] = useState<number | null>(null);
@ -26,6 +65,11 @@ export function TaskIndicator({ taskID, onTaskCancelled, onTaskCompleted }: Task
const [isCancelling, setIsCancelling] = useState(false);
const [isClearing, setIsClearing] = useState(false);
const [drawerOpen, setDrawerOpen] = useState(false);
const [selectedTaskId, setSelectedTaskId] = useState<string | null>(null);
// Prevents WS effect from overwriting local cancel/clear state before
// the parent's setTaskID(null) propagates. Reset when taskID changes.
const cancelledRef = useRef(false);
const onTaskCompletedRef = useRef(onTaskCompleted);
useEffect(() => {
@ -43,7 +87,58 @@ export function TaskIndicator({ taskID, onTaskCancelled, onTaskCompleted }: Task
}
}, []);
// Track the currently-viewed task in the drawer; default to the externally-provided taskID
useEffect(() => {
if (taskID) {
setSelectedTaskId(taskID);
cancelledRef.current = false; // new task, reset cancelled guard
}
}, [taskID]);
// Count active (non-terminal) tasks from WS
const activeWsTaskCount = useMemo(() => {
if (!wsTasks) return 0;
return Object.values(wsTasks).filter(
(t) => !isTerminalStatus(t.status),
).length;
}, [wsTasks]);
// ----- WebSocket-driven state updates -----
// When wsConnected, derive taskStatus/taskResult/progress from wsTasks
useEffect(() => {
if (!wsConnected || !wsTasks || !taskID) return;
// Don't let WS overwrite local cancel/clear state
if (cancelledRef.current) return;
const wsTask = wsTasks[taskID];
if (!wsTask) return;
const status = wsTask.status as TaskStatus;
setTaskStatus(status);
if (wsTask.phase) {
setTaskResult(taskStateToResult(wsTask));
}
if (wsTask.progress !== undefined) {
setProgressPercentage(wsTask.progress * 100);
}
if (wsTask.processed !== undefined) {
setProcessed(wsTask.processed);
}
if (wsTask.total !== undefined) {
setTotal(wsTask.total);
}
if (status === TaskStatus.SUCCESS) {
setProgressPercentage(100);
onTaskCompletedRef.current?.();
}
}, [wsConnected, wsTasks, taskID]);
// ----- Polling fallback (only when WS is not connected) -----
useEffect(() => {
// If WS is connected, skip polling
if (wsConnected) return;
if (!user || !taskID) {
setTaskStatus(null);
setTaskResult(null);
@ -65,10 +160,6 @@ export function TaskIndicator({ taskID, onTaskCancelled, onTaskCompleted }: Task
if (status === TaskStatus.SUCCESS) {
setProgressPercentage(100);
// Parse final result for the drawer to show completed state.
// Only update taskResult if the new result has phase info;
// otherwise keep the last in-progress result which has richer data
// than the bare SUCCESS return value.
if (data.result) {
try {
const parsedResult: TaskResult = JSON.parse(data.result);
@ -80,26 +171,19 @@ export function TaskIndicator({ taskID, onTaskCancelled, onTaskCompleted }: Task
}
}
onTaskCompletedRef.current?.();
return true; // Stop polling
return true;
}
if (status === TaskStatus.FAILURE || status === TaskStatus.REVOKED) {
return true; // Stop polling
return true;
}
// Parse progress for in-progress tasks
if (data.result) {
try {
const parsedResult: TaskResult = JSON.parse(data.result);
// Only update taskResult if the parsed data has a phase field.
// This prevents blanking the drawer when the backend sends a
// state update without phase info (e.g. during brief transitions).
if (parsedResult.phase) {
setTaskResult(parsedResult);
}
// Only update progress/processed/total when the fields are
// actually present — otherwise keep the previous values so
// the UI doesn't flash back to 0 during phase transitions.
if (parsedResult.progress !== undefined) {
setProgressPercentage(parsedResult.progress * 100);
}
@ -113,14 +197,13 @@ export function TaskIndicator({ taskID, onTaskCancelled, onTaskCompleted }: Task
// Ignore parsing errors
}
}
return false; // Continue polling
return false;
} catch {
setTaskStatus(TaskStatus.FAILURE);
return true; // Stop polling on error
return true;
}
};
// Initial poll
pollTaskStatus();
const interval = setInterval(async () => {
@ -131,7 +214,7 @@ export function TaskIndicator({ taskID, onTaskCancelled, onTaskCompleted }: Task
}, POLLING_INTERVALS.TASK_STATUS_MS);
return () => clearInterval(interval);
}, [taskID, user]);
}, [taskID, user, wsConnected]);
const handleCancel = async () => {
if (!user || !taskID || isCancelling) return;
@ -140,6 +223,7 @@ export function TaskIndicator({ taskID, onTaskCancelled, onTaskCompleted }: Task
try {
const result = await cancelTask(user, taskID);
if (result.success) {
cancelledRef.current = true;
setTaskStatus(TaskStatus.REVOKED);
onTaskCancelled?.();
}
@ -157,6 +241,7 @@ export function TaskIndicator({ taskID, onTaskCancelled, onTaskCompleted }: Task
try {
const result = await clearAllTasks(user);
if (result.success) {
cancelledRef.current = true;
setTaskStatus(null);
setTaskResult(null);
onTaskCancelled?.();
@ -245,6 +330,11 @@ export function TaskIndicator({ taskID, onTaskCancelled, onTaskCompleted }: Task
{taskStatus}
</span>
)}
{activeWsTaskCount > 1 && (
<span className="inline-flex items-center justify-center h-4 min-w-[16px] rounded-full bg-blue-500 text-[10px] font-medium text-white px-1">
{activeWsTaskCount}
</span>
)}
</div>
</TooltipTrigger>
<TooltipContent side="bottom">
@ -292,9 +382,12 @@ export function TaskIndicator({ taskID, onTaskCancelled, onTaskCompleted }: Task
onOpenChange={setDrawerOpen}
taskResult={taskResult}
taskStatus={taskStatus}
taskID={taskID}
taskID={selectedTaskId ?? taskID}
onCancel={handleCancel}
isCancelling={isCancelling}
wsTasks={wsTasks}
selectedTaskId={selectedTaskId}
onSelectTask={setSelectedTaskId}
/>
</TooltipProvider>
);