feat(queue): add redis worker orchestration

Author: haasonsaas

agent_pm/observability/metrics.py

@@ -65,6 +65,25 @@
     labelnames=("source",),
 )
 
+task_queue_enqueued_total = Counter(
+    "task_queue_enqueued_total",
+    "Tasks enqueued to background queue",
+    labelnames=("queue",),
+)
+
+task_queue_completed_total = Counter(
+    "task_queue_completed_total",
+    "Tasks completed grouped by status",
+    labelnames=("queue", "status"),
+)
+
+task_queue_latency_seconds = Histogram(
+    "task_queue_latency_seconds",
+    "Task execution latency in seconds",
+    labelnames=("queue",),
+    buckets=(0.05, 0.1, 0.25, 0.5, 1, 2, 5, 10, 30, 60),
+)
+
 plugin_hook_invocations_total = Counter(
     "plugin_hook_invocations_total",
     "Plugin hook invocation counts",
@@ -158,6 +177,18 @@ def record_client_call(client: str):
         client_latency_seconds.labels(client=client).observe(perf_counter() - start)
 
 
+def record_task_enqueued(queue: str) -> None:
+    task_queue_enqueued_total.labels(queue=queue).inc()
+
+
+def record_task_completion(queue: str, status: str) -> None:
+    task_queue_completed_total.labels(queue=queue, status=status).inc()
+
+
+def record_task_latency(queue: str, duration: float) -> None:
+    task_queue_latency_seconds.labels(queue=queue).observe(duration)
+
+
 def latest_metrics() -> bytes:
     return generate_latest()
 
@@ -175,5 +206,8 @@ def latest_metrics() -> bytes:
     "record_plugin_hook_invocation",
     "record_plugin_hook_failure",
     "record_client_call",
+    "record_task_enqueued",
+    "record_task_completion",
+    "record_task_latency",
     "latest_metrics",
 ]

agent_pm/settings.py

@@ -71,6 +71,11 @@ def _parse_google_scopes(cls, value):
     log_format: str = Field("json", alias="LOG_FORMAT")  # json or text
     task_queue_workers: int = Field(5, alias="TASK_QUEUE_WORKERS")
     task_queue_backend: Literal["memory", "redis"] = Field("memory", alias="TASK_QUEUE_BACKEND")
+    task_queue_poll_interval: float = Field(0.2, alias="TASK_QUEUE_POLL_INTERVAL")
+    task_queue_task_timeout: int = Field(300, alias="TASK_QUEUE_TASK_TIMEOUT")
+    task_queue_retry_backoff_base: float = Field(2.0, alias="TASK_QUEUE_RETRY_BACKOFF_BASE")
+    task_queue_retry_backoff_max: float = Field(60.0, alias="TASK_QUEUE_RETRY_BACKOFF_MAX")
+    task_queue_worker_heartbeat_ttl: int = Field(60, alias="TASK_QUEUE_WORKER_HEARTBEAT_TTL")
     database_url: str | None = Field("sqlite+aiosqlite:///./data/agent_pm.db", alias="DATABASE_URL")
     database_echo: bool = Field(False, alias="DATABASE_ECHO")
     redis_url: str = Field("redis://localhost:6379", alias="REDIS_URL")

agent_pm/storage/redis.py

@@ -15,6 +15,8 @@
 
 QUEUE_KEY = "agent_pm:tasks"
 RESULT_KEY = "agent_pm:results"
+DEAD_LETTER_KEY = "agent_pm:dead_letter"
+HEARTBEAT_KEY = "agent_pm:worker_heartbeats"
 
 
 def _queue_key() -> str:
@@ -25,6 +27,14 @@ def _result_key() -> str:
     return RESULT_KEY
 
 
+def _dead_letter_key() -> str:
+    return DEAD_LETTER_KEY
+
+
+def _heartbeat_key() -> str:
+    return HEARTBEAT_KEY
+
+
 async def get_redis_client() -> redis.Redis:
     return redis.from_url(settings.redis_url, decode_responses=True)
 
@@ -57,3 +67,43 @@ async def get_task_result(client: redis.Redis, task_id: str) -> dict[str, Any] |
     if not item:
         return None
     return json.loads(item)
+
+
+async def record_dead_letter(client: redis.Redis, payload: dict[str, Any]) -> None:
+    task_id = payload.get("task_id", uuid.uuid4().hex)
+    await client.hset(_dead_letter_key(), task_id, json.dumps(payload))
+
+
+async def fetch_dead_letters(client: redis.Redis, limit: int = 100) -> list[dict[str, Any]]:
+    items = await client.hgetall(_dead_letter_key())
+    tasks: list[dict[str, Any]] = []
+    for task_id, value in items.items():
+        if len(tasks) >= limit:
+            break
+        try:
+            data = json.loads(value)
+            data.setdefault("task_id", task_id)
+            tasks.append(data)
+        except json.JSONDecodeError:
+            tasks.append({"task_id": task_id, "raw": value})
+    return tasks
+
+
+async def clear_dead_letter(client: redis.Redis, task_id: str) -> None:
+    await client.hdel(_dead_letter_key(), task_id)
+
+
+async def write_heartbeat(client: redis.Redis, worker_id: str, payload: dict[str, Any], ttl: int) -> None:
+    await client.hset(_heartbeat_key(), worker_id, json.dumps(payload))
+    await client.expire(_heartbeat_key(), ttl)
+
+
+async def list_heartbeats(client: redis.Redis) -> dict[str, dict[str, Any]]:
+    items = await client.hgetall(_heartbeat_key())
+    heartbeats: dict[str, dict[str, Any]] = {}
+    for worker_id, value in items.items():
+        try:
+            heartbeats[worker_id] = json.loads(value)
+        except json.JSONDecodeError:
+            heartbeats[worker_id] = {"raw": value}
+    return heartbeats

agent_pm/storage/tasks.py

@@ -12,10 +12,24 @@
 from enum import Enum
 from typing import Any
 
+from ..observability.metrics import (
+    record_task_completion,
+    record_task_enqueued,
+    record_task_latency,
+)
 from ..settings import settings
 from ..utils.datetime import utc_now
-from .redis import enqueue_task as redis_enqueue_task
-from .redis import get_redis_client
+from .redis import (
+    clear_dead_letter,
+    enqueue_task as redis_enqueue_task,
+    fetch_dead_letters,
+    get_redis_client,
+    list_heartbeats,
+    pop_task,
+    record_dead_letter,
+    set_task_result,
+    write_heartbeat,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -61,6 +75,7 @@ def __init__(self, max_workers: int = 5):
         self.workers: list[asyncio.Task[None]] = []
         self.running = False
         self._lock = asyncio.Lock()
+        self.queue_name = "memory"
 
     async def start(self):
         """Start background workers."""
@@ -102,6 +117,7 @@ async def enqueue(
         async with self._lock:
             self.queue.append(task)
             self.tasks[task_id] = task
+        record_task_enqueued(self.queue_name)
         logger.info("Task enqueued: %s (id=%s)", name, task_id)
         return task_id
 
@@ -153,6 +169,8 @@ async def _execute_task(self, task: Task):
             task.status = TaskStatus.COMPLETED
             task.result = result
             task.completed_at = utc_now()
+            record_task_completion(self.queue_name, task.status.value)
+            record_task_latency(self.queue_name, (task.completed_at - task.started_at).total_seconds())
             logger.info("Task completed: %s (id=%s)", task.name, task.task_id)
         except Exception as exc:
             task.error = str(exc)
@@ -175,8 +193,19 @@ async def _execute_task(self, task: Task):
             else:
                 task.status = TaskStatus.FAILED
                 task.completed_at = utc_now()
+                record_task_completion(self.queue_name, task.status.value)
+                record_task_latency(self.queue_name, (task.completed_at - task.started_at).total_seconds())
                 logger.error("Task permanently failed: %s (id=%s)", task.name, task.task_id)
 
+    async def list_dead_letters(self, limit: int = 100) -> list[dict[str, Any]]:
+        return []
+
+    async def delete_dead_letter(self, task_id: str) -> None:
+        return None
+
+    async def worker_heartbeats(self) -> dict[str, Any]:
+        return {}
+
 
 # Global task queue instance
 _task_queue: TaskQueue | None = None
@@ -191,6 +220,20 @@ async def get_task_queue() -> TaskQueue:
             client = await get_redis_client()
 
             class RedisTaskQueue(TaskQueue):
+                def __init__(self, max_workers: int = 5):
+                    super().__init__(max_workers=max_workers)
+                    self.queue_name = "redis"
+                    self._redis = client
+                    self._registry: dict[str, Callable[..., Coroutine[Any, Any, Any]]] = {}
+                    self._poll_interval = settings.task_queue_poll_interval
+                    self._task_timeout = settings.task_queue_task_timeout
+                    self._backoff_base = settings.task_queue_retry_backoff_base
+                    self._backoff_max = settings.task_queue_retry_backoff_max
+                    self._heartbeat_ttl = settings.task_queue_worker_heartbeat_ttl
+
+                def register(self, name: str, coro_fn: Callable[..., Coroutine[Any, Any, Any]]) -> None:
+                    self._registry[name] = coro_fn
+
                 async def enqueue(  # type: ignore[override]
                     self,
                     name: str,
@@ -199,14 +242,101 @@ async def enqueue(  # type: ignore[override]
                     max_retries: int = 3,
                     **kwargs: Any,
                 ) -> str:
+                    self.register(name, coro_fn)
                     payload = {
                         "task_id": uuid.uuid4().hex,
                         "name": name,
                         "args": args,
                         "kwargs": kwargs,
                         "max_retries": max_retries,
+                        "enqueued_at": utc_now().isoformat(),
                     }
-                    return await redis_enqueue_task(client, name, payload)
+                    await redis_enqueue_task(self._redis, name, payload)
+                    record_task_enqueued(self.queue_name)
+                    logger.info("Redis task enqueued: %s (id=%s)", name, payload["task_id"])
+                    return payload["task_id"]
+
+                async def pop(self) -> dict[str, Any] | None:
+                    payload = await pop_task(self._redis)
+                    if not payload:
+                        return None
+                    return payload
+
+                async def _worker(self, worker_id: int):
+                    logger.info("Redis worker %d started", worker_id)
+                    while self.running:
+                        payload = await self.pop()
+                        if not payload:
+                            await asyncio.sleep(self._poll_interval)
+                            continue
+
+                        task_id = payload.get("task_id", "unknown")
+                        name = payload.get("name")
+                        coro_fn = self._registry.get(name)
+                        if coro_fn is None:
+                            logger.error("No registered task callable for %s", name)
+                            await record_dead_letter(
+                                self._redis,
+                                {**payload, "error": "missing_callable", "worker_id": worker_id},
+                            )
+                            record_task_completion(self.queue_name, TaskStatus.FAILED.value)
+                            continue
+
+                        retries = payload.get("retry_count", 0)
+                        max_retries = payload.get("max_retries", 3)
+
+                        start = utc_now()
+                        try:
+                            result = await asyncio.wait_for(
+                                coro_fn(*payload.get("args", ()), **payload.get("kwargs", {})),
+                                timeout=self._task_timeout,
+                            )
+                        except asyncio.TimeoutError:
+                            payload["retry_count"] = retries + 1
+                            payload["last_error"] = "timeout"
+                            payload["worker_id"] = worker_id
+                            await record_dead_letter(self._redis, payload)
+                            record_task_completion(self.queue_name, TaskStatus.FAILED.value)
+                            record_task_latency(self.queue_name, (utc_now() - start).total_seconds())
+                            continue
+                        except Exception as exc:  # pylint: disable=broad-except
+                            retries += 1
+                            payload["retry_count"] = retries
+                            payload["last_error"] = str(exc)
+                            if retries >= max_retries:
+                                payload["worker_id"] = worker_id
+                                await record_dead_letter(self._redis, payload)
+                                record_task_completion(self.queue_name, TaskStatus.FAILED.value)
+                                record_task_latency(self.queue_name, (utc_now() - start).total_seconds())
+                                continue
+
+                            backoff = min(self._backoff_base**retries, self._backoff_max)
+                            await asyncio.sleep(backoff)
+                            await redis_enqueue_task(self._redis, name, payload)
+                            continue
+
+                        await set_task_result(self._redis, task_id, {"status": "completed", "result": result})
+                        record_task_completion(self.queue_name, TaskStatus.COMPLETED.value)
+                        record_task_latency(self.queue_name, (utc_now() - start).total_seconds())
+
+                        heartbeat_payload = {
+                            "worker_id": worker_id,
+                            "task_id": task_id,
+                            "name": name,
+                            "completed_at": utc_now().isoformat(),
+                        }
+                        await write_heartbeat(self._redis, f"worker:{worker_id}", heartbeat_payload, self._heartbeat_ttl)
+
+                    logger.info("Redis worker %d stopped", worker_id)
+
+                async def list_dead_letters(self, limit: int = 100) -> list[dict[str, Any]]:
+                    return await fetch_dead_letters(self._redis, limit)
+
+                async def delete_dead_letter(self, task_id: str) -> None:
+                    await clear_dead_letter(self._redis, task_id)
+
+                async def worker_heartbeats(self) -> dict[str, dict[str, Any]]:
+                    return await list_heartbeats(self._redis)
 
             _task_queue = RedisTaskQueue(max_workers=settings.task_queue_workers)
         else:

app.py

@@ -41,7 +41,10 @@
 )
 from agent_pm.observability.export import schedule_trace_export
 from agent_pm.observability.logging import configure_logging
-from agent_pm.observability.metrics import latest_metrics, record_alignment_export
+from agent_pm.observability.metrics import (
+    latest_metrics,
+    record_alignment_export,
+)
 from agent_pm.observability.structured import (
     configure_structured_logging,
     get_correlation_id,
@@ -86,14 +89,13 @@ class FollowupUpdate(BaseModel):
 async def startup_event():
     global _task_queue
     _task_queue = await get_task_queue()
-    if settings.task_queue_backend == "memory":
-        await _task_queue.start()
+    await _task_queue.start()
     logger.info("Agent PM service started")
 
 
 @app.on_event("shutdown")
 async def shutdown_event():
-    if settings.task_queue_backend == "memory" and _task_queue:
+    if _task_queue:
         await _task_queue.stop()
     logger.info("Agent PM service stopped")
 
@@ -473,6 +475,29 @@ async def list_tasks(status: str | None = None, limit: int = 50, _admin_key: Adm
     }
 
 
+@app.get("/tasks/dead-letter")
+async def list_dead_letter(limit: int = 50, _admin_key: AdminKeyDep = None) -> dict[str, Any]:
+    if not _task_queue:
+        raise HTTPException(status_code=503, detail="Task queue not initialized")
+    items = await _task_queue.list_dead_letters(limit)
+    return {"dead_letter": items, "total": len(items)}
+
+
+@app.delete("/tasks/dead-letter/{task_id}")
+async def delete_dead_letter(task_id: str, _admin_key: AdminKeyDep = None) -> dict[str, Any]:
+    if not _task_queue:
+        raise HTTPException(status_code=503, detail="Task queue not initialized")
+    await _task_queue.delete_dead_letter(task_id)
+    return {"task_id": task_id, "status": "deleted"}
+
+
+@app.get("/tasks/workers")
+async def worker_status(_admin_key: AdminKeyDep = None) -> dict[str, Any]:
+    if not _task_queue:
+        raise HTTPException(status_code=503, detail="Task queue not initialized")
+    return {"workers": await _task_queue.worker_heartbeats()}
+
+
 @app.post("/prd/{plan_id}/versions")
 async def create_prd_version(
     plan_id: str,