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

@ -0,0 +1,50 @@
"""Publishes task progress updates to Redis pub/sub channels.
Celery workers call publish_task_progress() alongside task.update_state() so
that the FastAPI WebSocket endpoint can forward real-time updates to connected
browsers without polling.
Channel naming: ``task_progress:{task_id}``
"""
import json
import logging
import os
from typing import Any
import redis
logger = logging.getLogger(__name__)
_redis_client: redis.Redis | None = None # type: ignore[type-arg]
def _get_redis_client() -> redis.Redis: # type: ignore[type-arg]
"""Lazy-initialised Redis client derived from CELERY_BROKER_URL."""
global _redis_client
if _redis_client is None:
broker_url = os.getenv("CELERY_BROKER_URL", "redis://redis:6379/0")
_redis_client = redis.Redis.from_url(broker_url, decode_responses=True)
return _redis_client
def publish_task_progress(task_id: str, state: str, meta: dict[str, Any]) -> None:
"""Publish a task progress update to Redis pub/sub.
Args:
task_id: Celery task ID.
state: Celery state string (e.g. 'PROGRESS', 'SUCCESS').
meta: Metadata dict (progress, phase, logs, counters, etc.).
Failures are caught and logged at DEBUG level so they never break the
critical task execution path.
"""
try:
client = _get_redis_client()
payload = json.dumps({
"task_id": task_id,
"status": state,
**meta,
})
client.publish(f"task_progress:{task_id}", payload)
except Exception:
logger.debug("Failed to publish task progress for %s", task_id, exc_info=True)

View file

@ -177,11 +177,19 @@ def cancel_task(task_id: str, user_email: str | None = None) -> bool:
"""
# Lazy import: celery_app bootstraps the broker connection.
from celery_app import app as celery_app
from services.task_progress_publisher import publish_task_progress
logger.info("Cancelling task %s (user=%s)", task_id, user_email)
# Revoke the task in Celery
celery_app.control.revoke(task_id, terminate=True)
# Publish REVOKED status via pub/sub so WebSocket clients learn immediately
publish_task_progress(task_id, "REVOKED", {
"phase": "completed",
"progress": 0,
"message": "Task cancelled",
})
# Also remove from user's task list if user_email provided
if user_email:
remove_task_from_user(user_email, task_id)
@ -222,6 +230,7 @@ def clear_all_tasks(user_email: str, revoke: bool = True) -> int:
# Lazy imports: see get_user_tasks and cancel_task for rationale.
from redis_repository import RedisRepository
from celery_app import app as celery_app
from services.task_progress_publisher import publish_task_progress
redis_repo = RedisRepository.instance()
user = _make_system_user(user_email)
@ -238,5 +247,11 @@ def clear_all_tasks(user_email: str, revoke: bool = True) -> int:
logger.warning(
"Failed to revoke task %s: %s", task_id, e
)
# Publish REVOKED via pub/sub so WebSocket clients learn immediately
publish_task_progress(task_id, "REVOKED", {
"phase": "completed",
"progress": 0,
"message": "Task cancelled",
})
return redis_repo.clear_tasks_for_user(user)