1

video-gen-api/app/tasks/generation_poll_tasks.py

@@ -1,11 +1,14 @@
 from __future__ import annotations
 
+import logging
 import json
+import uuid
 from datetime import datetime, timedelta, timezone
 from typing import Any
 
 from app.config import settings
-from app.enums.celery_queue import CeleryQueue
+from app.enums.celery_queue import CeleryQueue, CeleryTaskName
+from app.enums.celery_runtime import CeleryRuntimeDomain
 from app.enums.generation_status import GenerationRecordPipelineStage
 from app.enums.generation_task import (
     ALLOWED_GENERATION_MODES,
@@ -36,6 +39,7 @@ from app.services.generation.pipeline.owner_service import (
     owner_provider_task_id,
     owner_type_of,
     redis_owner_item_id,
+    renew_generation_owner_claim_lease,
 )
 from app.services.generation.poll_schedule_service import (
     build_default_poll_schedule,
@@ -52,16 +56,19 @@ from app.services.media_token_usage_snapshot_service import (
 )
 from app.services.redis_registry_service import (
     RedisExecutionLockError,
-    RedisExecutionLockLease,
     datetime_to_epoch,
     ensure_aware_utc,
+    redis_get_registry_payloads,
     redis_remove_registry_item,
     redis_upsert_registry_item,
     utc_now,
 )
+from app.services.celery_runtime.runtime_service import CeleryRuntimeLease, RuntimeIdentity
 from app.tasks.async_runner import run_async
 from app.tasks.celery_app import celery_app
 
+logger = logging.getLogger("video_gen")
+
 POLL_QUEUE = CeleryQueue.GEN_PROVIDER_POLL.value
 
 
@@ -183,16 +190,27 @@ async def register_poll_active(
     reason: str,
     next_poll_at: datetime | None = None,
 ) -> None:
+    item_id = _registry_id(owner)
+    payload = _build_poll_active_payload(
+        owner,
+        reason=reason,
+        next_poll_at=next_poll_at,
+        check_at=check_at,
+    )
+    existing = await redis_get_registry_payloads(
+        hash_key=settings.POLL_ACTIVE_REDIS_HASH_KEY,
+        item_ids=[item_id],
+        log_context="poll_active",
+    )
+    if item_id in existing:
+        merged = dict(existing[item_id])
+        merged.update(payload)
+        payload = merged
     await redis_upsert_registry_item(
         hash_key=settings.POLL_ACTIVE_REDIS_HASH_KEY,
         zset_key=settings.POLL_ACTIVE_REDIS_ZSET_KEY,
-        item_id=_registry_id(owner),
-        payload=_build_poll_active_payload(
-            owner,
-            reason=reason,
-            next_poll_at=next_poll_at,
-            check_at=check_at,
-        ),
+        item_id=item_id,
+        payload=payload,
         check_at=check_at,
         log_context="poll_active",
     )
@@ -388,12 +406,34 @@ async def _run(
     if effective_attempt is None:
         return
 
-    lease = await RedisExecutionLockLease.acquire(
+    token = uuid.uuid4().hex
+    lease = await CeleryRuntimeLease.acquire(
+        identity=RuntimeIdentity(
+            domain=CeleryRuntimeDomain.GENERATION_POLL.value,
+            owner_type=normalized_owner_type,
+            owner_id=task_id,
+            attempt_no=effective_attempt,
+            task_name=CeleryTaskName.POLL_GENERATION.value,
+            queue=POLL_QUEUE,
+            registry_item_id=redis_owner_item_id(
+                normalized_owner_type, task_id, effective_attempt
+            ),
+        ),
         lock_key=_lock_key(normalized_owner_type, task_id, effective_attempt),
+        hash_key=settings.POLL_ACTIVE_REDIS_HASH_KEY,
+        zset_key=settings.POLL_ACTIVE_REDIS_ZSET_KEY,
+        token=token,
         ttl_seconds=int(settings.GENERATION_POLL_LOCK_TTL_SECONDS or 300),
-        log_context="generation_poll",
-        renew_interval_seconds=int(
-            settings.REDIS_EXECUTION_LOCK_RENEW_INTERVAL_SECONDS or 30
+        heartbeat_interval_seconds=int(settings.REDIS_EXECUTION_LOCK_RENEW_INTERVAL_SECONDS or 30),
+        pipeline_stage=ChatGenerationPipelineStage.POLLING.value,
+        db_heartbeat=lambda owned_token: renew_generation_owner_claim_lease(
+            owner_type=normalized_owner_type,
+            owner_id=task_id,
+            attempt_no=effective_attempt,
+            claim_field="poll_claim_token",
+            lease_field="poll_lease_until",
+            token=owned_token,
+            lease_seconds=int(settings.POLL_TASK_LEASE_SECONDS or 300),
         ),
     )
     if lease is None: