feat(queue): add audit trail and retry policies

Author: haasonsaas

agent_pm/storage/redis.py

@@ -17,6 +17,8 @@
 QUEUE_KEY = "agent_pm:tasks"
 RESULT_KEY = "agent_pm:results"
 DEAD_LETTER_KEY = "agent_pm:dead_letter"
+DEAD_LETTER_AUDIT_KEY = "agent_pm:dead_letter:audit"
+RETRY_POLICY_KEY = "agent_pm:retry_policy"
 HEARTBEAT_KEY = "agent_pm:worker_heartbeats"
 
 
@@ -32,6 +34,14 @@ def _dead_letter_key() -> str:
     return DEAD_LETTER_KEY
 
 
+def _dead_letter_audit_key() -> str:
+    return DEAD_LETTER_AUDIT_KEY
+
+
+def _retry_policy_key() -> str:
+    return RETRY_POLICY_KEY
+
+
 def _heartbeat_key() -> str:
     return HEARTBEAT_KEY
 
@@ -122,6 +132,57 @@ async def clear_dead_letter(client: redis.Redis, task_id: str) -> None:
     await client.hdel(_dead_letter_key(), task_id)
 
 
+async def append_dead_letter_audit(client: redis.Redis, entry: dict[str, Any], *, max_entries: int = 1000) -> None:
+    payload = entry.copy()
+    payload.setdefault("timestamp", datetime.now(UTC).isoformat())
+    await client.lpush(_dead_letter_audit_key(), json.dumps(payload))
+    await client.ltrim(_dead_letter_audit_key(), 0, max_entries - 1)
+
+
+async def fetch_dead_letter_audit(client: redis.Redis, limit: int = 100) -> list[dict[str, Any]]:
+    raw_entries = await client.lrange(_dead_letter_audit_key(), 0, max(0, limit - 1))
+    entries: list[dict[str, Any]] = []
+    for value in raw_entries:
+        try:
+            entries.append(json.loads(value))
+        except json.JSONDecodeError:
+            entries.append({"raw": value})
+    return entries
+
+
+async def set_retry_policy(client: redis.Redis, task_name: str, policy: dict[str, Any]) -> None:
+    if not policy:
+        await client.hdel(_retry_policy_key(), task_name)
+        return
+    await client.hset(_retry_policy_key(), task_name, json.dumps(policy))
+
+
+async def get_retry_policy(client: redis.Redis, task_name: str) -> dict[str, Any] | None:
+    value = await client.hget(_retry_policy_key(), task_name)
+    if not value:
+        return None
+    try:
+        data = json.loads(value)
+    except json.JSONDecodeError:
+        return None
+    return data
+
+
+async def delete_retry_policy(client: redis.Redis, task_name: str) -> None:
+    await client.hdel(_retry_policy_key(), task_name)
+
+
+async def list_retry_policies(client: redis.Redis) -> dict[str, dict[str, Any]]:
+    values = await client.hgetall(_retry_policy_key())
+    policies: dict[str, dict[str, Any]] = {}
+    for name, raw in values.items():
+        try:
+            policies[name] = json.loads(raw)
+        except json.JSONDecodeError:
+            continue
+    return policies
+
+
 async def purge_dead_letters(client: redis.Redis, *, older_than: datetime | None = None) -> int:
     if older_than is None:
         count = await client.hlen(_dead_letter_key())

agent_pm/storage/tasks.py

@@ -25,16 +25,22 @@
 from ..settings import settings
 from ..utils.datetime import utc_now
 from .redis import (
+    append_dead_letter_audit,
     clear_dead_letter,
     count_dead_letters,
+    delete_retry_policy,
     enqueue_task as redis_enqueue_task,
+    fetch_dead_letter_audit,
     fetch_dead_letters,
     get_dead_letter,
     get_redis_client,
+    get_retry_policy,
     list_heartbeats,
+    list_retry_policies,
     pop_task,
     purge_dead_letters,
     record_dead_letter,
+    set_retry_policy,
     set_task_result,
     write_heartbeat,
 )
@@ -217,6 +223,21 @@ async def list_dead_letters(
     ) -> tuple[list[dict[str, Any]], int]:
         return [], 0
 
+    async def list_dead_letter_audit(self, limit: int = 100) -> list[dict[str, Any]]:
+        return []
+
+    async def get_retry_policy(self, task_name: str) -> dict[str, Any] | None:
+        return None
+
+    async def set_retry_policy(self, task_name: str, policy: dict[str, Any]) -> None:
+        return None
+
+    async def delete_retry_policy(self, task_name: str) -> None:
+        return None
+
+    async def list_retry_policies(self) -> dict[str, dict[str, Any]]:
+        return {}
+
     async def delete_dead_letter(self, task_id: str) -> None:
         return None
 
@@ -314,23 +335,23 @@ async def _worker(self, worker_id: int):
                             await record_dead_letter(self._redis, payload)
                             dead_letter_recorded_total.labels(
                                 queue=self.queue_name,
-                                error_type="TimeoutError",
-                            ).inc()
-                            dead_letter_recorded_total.labels(
-                                queue=self.queue_name,
-                                error_type=payload.get("error_type", payload.get("last_error", "unknown")),
+                                error_type=payload.get("error_type", "MissingCallable"),
                             ).inc()
                             record_task_completion(self.queue_name, TaskStatus.FAILED.value)
                             continue
 
                         retries = payload.get("retry_count", 0)
-                        max_retries = payload.get("max_retries", 3)
+                        base_max_retries = payload.get("max_retries", 3)
+
+                        policy = await get_retry_policy(self._redis, name) or {}
+                        timeout = float(policy.get("timeout", self._task_timeout))
+                        max_retries = int(policy.get("max_retries", base_max_retries))
 
                         start = utc_now()
                         try:
                             result = await asyncio.wait_for(
                                 coro_fn(*payload.get("args", ()), **payload.get("kwargs", {})),
-                                timeout=self._task_timeout,
+                                timeout=timeout,
                             )
                         except TimeoutError:
                             payload.setdefault("metadata", {})
@@ -343,6 +364,10 @@ async def _worker(self, worker_id: int):
                             payload["stack_trace"] = None
                             payload["worker_id"] = worker_id
                             await record_dead_letter(self._redis, payload)
+                            dead_letter_recorded_total.labels(
+                                queue=self.queue_name,
+                                error_type="TimeoutError",
+                            ).inc()
                             record_task_completion(self.queue_name, TaskStatus.FAILED.value)
                             record_task_latency(self.queue_name, (utc_now() - start).total_seconds())
                             continue
@@ -365,7 +390,10 @@ async def _worker(self, worker_id: int):
                                 record_task_latency(self.queue_name, (utc_now() - start).total_seconds())
                                 continue
 
-                            backoff = min(self._backoff_base**retries, self._backoff_max)
+                            backoff = min(
+                                float(policy.get("backoff_base", self._backoff_base)) ** retries,
+                                float(policy.get("backoff_max", self._backoff_max)),
+                            )
                             await asyncio.sleep(backoff)
                             await redis_enqueue_task(self._redis, name, payload)
                             continue
@@ -404,6 +432,21 @@ async def list_dead_letters(
                     dead_letter_active_gauge.labels(queue=self.queue_name).set(total_filtered)
                     return window, total_filtered
 
+                async def list_dead_letter_audit(self, limit: int = 100) -> list[dict[str, Any]]:
+                    return await fetch_dead_letter_audit(self._redis, limit)
+
+                async def get_retry_policy(self, task_name: str) -> dict[str, Any] | None:
+                    return await get_retry_policy(self._redis, task_name)
+
+                async def set_retry_policy(self, task_name: str, policy: dict[str, Any]) -> None:
+                    await set_retry_policy(self._redis, task_name, policy)
+
+                async def delete_retry_policy(self, task_name: str) -> None:
+                    await delete_retry_policy(self._redis, task_name)
+
+                async def list_retry_policies(self) -> dict[str, dict[str, Any]]:
+                    return await list_retry_policies(self._redis)
+
                 async def delete_dead_letter(self, task_id: str) -> None:
                     await clear_dead_letter(self._redis, task_id)
 

tests/storage/test_redis.py

@@ -12,6 +12,8 @@ def __init__(self):
         self.results: dict[str, str] = {}
         self.dead: dict[str, str] = {}
         self.heartbeats: dict[str, str] = {}
+        self.audit: list[str] = []
+        self.retry_policies: dict[str, str] = {}
 
     async def rpush(self, key: str, value: str) -> None:
         self.items.append(value)
@@ -29,28 +31,38 @@ async def hset(self, key: str, field: str, value: str) -> None:
             self.dead[field] = value
         elif key.endswith("heartbeats"):
             self.heartbeats[field] = value
+        elif key.endswith("retry_policy"):
+            self.retry_policies[field] = value
         else:
             self.results[field] = value
 
     async def hget(self, key: str, field: str):
         if key.endswith("dead_letter"):
             return self.dead.get(field)
+        if key.endswith("retry_policy"):
+            return self.retry_policies.get(field)
         return self.results.get(field)
 
     async def hgetall(self, key: str):
         if key.endswith("dead_letter"):
             return self.dead
         if key.endswith("heartbeats"):
             return self.heartbeats
+        if key.endswith("retry_policy"):
+            return self.retry_policies
         return self.results
 
     async def hdel(self, key: str, field: str):
         if key.endswith("dead_letter"):
             self.dead.pop(field, None)
+        elif key.endswith("retry_policy"):
+            self.retry_policies.pop(field, None)
 
     async def hlen(self, key: str) -> int:
         if key.endswith("dead_letter"):
             return len(self.dead)
+        if key.endswith("retry_policy"):
+            return len(self.retry_policies)
         return len(self.results)
 
     async def expire(self, key: str, ttl: int):
@@ -59,6 +71,18 @@ async def expire(self, key: str, ttl: int):
     async def delete(self, key: str):
         if key.endswith("dead_letter"):
             self.dead.clear()
+    async def lpush(self, key: str, value: str):
+        if key.endswith("dead_letter:audit"):
+            self.audit.insert(0, value)
+
+    async def ltrim(self, key: str, start: int, stop: int):
+        if key.endswith("dead_letter:audit"):
+            self.audit = self.audit[start : stop + 1]
+
+    async def lrange(self, key: str, start: int, stop: int):
+        if key.endswith("dead_letter:audit"):
+            return self.audit[start : stop + 1]
+        return []
 
 
 @pytest.mark.asyncio
@@ -119,3 +143,15 @@ async def test_dead_letter_and_heartbeat_helpers():
     await redis.write_heartbeat(client, "worker:1", {"status": "ok"}, ttl=60)
     beats = await redis.list_heartbeats(client)
     assert beats["worker:1"]["status"] == "ok"
+
+    await redis.append_dead_letter_audit(client, {"event": "record", "task_id": "abc"}, max_entries=10)
+    audit_entries = await redis.fetch_dead_letter_audit(client, limit=1)
+    assert audit_entries[0]["task_id"] == "abc"
+
+    await redis.set_retry_policy(client, "job", {"timeout": 123, "backoff_base": 2.5})
+    policy = await redis.get_retry_policy(client, "job")
+    assert policy == {"timeout": 123, "backoff_base": 2.5}
+    policies = await redis.list_retry_policies(client)
+    assert "job" in policies
+    await redis.delete_retry_policy(client, "job")
+    assert await redis.get_retry_policy(client, "job") is None