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:
parent
73d19e29d5
commit
8559c4b461
11 changed files with 774 additions and 72 deletions
|
|
@ -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>
|
||||
);
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue