feat: migrate task queue to native redis client

Author: haasonsaas

agent_pm/storage/redis.py

@@ -1,76 +1,59 @@
-"""Redis-backed distributed task queue using ARQ."""
+"""Redis-backed task queue using redis.asyncio."""
 
 from __future__ import annotations
 
+import json
 import logging
+import uuid
 from typing import Any
 
-from arq import create_pool
-from arq.connections import ArqRedis, RedisSettings
+import redis.asyncio as redis
 
 from agent_pm.settings import settings
 
 logger = logging.getLogger(__name__)
 
+QUEUE_KEY = "agent_pm:tasks"
+RESULT_KEY = "agent_pm:results"
 
-async def get_redis_pool() -> ArqRedis:
-    """Get or create Redis connection pool for ARQ."""
-    # Parse Redis URL
-    url = settings.redis_url
-    if url.startswith("redis://"):
-        host_port = url.replace("redis://", "").split("/")[0]
-        host, port = host_port.split(":") if ":" in host_port else (host_port, 6379)
-    else:
-        host, port = "localhost", 6379
 
-    redis_settings = RedisSettings(host=host, port=int(port))
-    return await create_pool(redis_settings)
+def _queue_key() -> str:
+    return QUEUE_KEY
 
 
-async def enqueue_task(
-    pool: ArqRedis,
-    function_name: str,
-    *args: Any,
-    **kwargs: Any,
-) -> str:
-    """Enqueue a task to Redis queue.
+def _result_key() -> str:
+    return RESULT_KEY
 
-    Returns:
-        Job ID
-    """
-    job = await pool.enqueue_job(function_name, *args, **kwargs)
-    logger.info("Enqueued task: %s (job_id=%s)", function_name, job.job_id)
-    return job.job_id
 
+async def get_redis_client() -> redis.Redis:
+    return redis.from_url(settings.redis_url, decode_responses=True)
 
-async def get_job_status(pool: ArqRedis, job_id: str) -> dict[str, Any] | None:
-    """Get status of a job by ID."""
-    job_result = await pool.job_status(job_id)
-    if not job_result:
-        return None
 
-    return {
-        "job_id": job_id,
-        "status": job_result.job_status.value if job_result.job_status else "unknown",
-        "result": job_result.result,
-        "start_time": (job_result.start_time.isoformat() if job_result.start_time else None),
-        "finish_time": (job_result.finish_time.isoformat() if job_result.finish_time else None),
-    }
+async def enqueue_task(
+    client: redis.Redis,
+    name: str,
+    payload: dict[str, Any],
+) -> str:
+    task_id = payload.setdefault("task_id", uuid.uuid4().hex)
+    payload.setdefault("name", name)
+    await client.rpush(_queue_key(), json.dumps(payload))
+    logger.info("Queued redis task %s (%s)", name, task_id)
+    return task_id
+
 
+async def pop_task(client: redis.Redis) -> dict[str, Any] | None:
+    item = await client.lpop(_queue_key())
+    if not item:
+        return None
+    return json.loads(item)
 
-# ARQ worker functions (to be imported by worker process)
-async def example_background_task(ctx: dict, plan_id: str, user: str) -> dict:
-    """Example background task function."""
-    logger.info("Processing background task for plan_id=%s user=%s", plan_id, user)
-    # Perform async work here
-    return {"plan_id": plan_id, "status": "processed"}
 
+async def set_task_result(client: redis.Redis, task_id: str, result: dict[str, Any]) -> None:
+    await client.hset(_result_key(), task_id, json.dumps(result))
 
-# Worker class for ARQ
-class WorkerSettings:
-    """ARQ worker configuration."""
 
-    functions = [example_background_task]
-    redis_settings = RedisSettings(host="localhost", port=6379)
-    max_jobs = 10
-    job_timeout = 300  # 5 minutes
+async def get_task_result(client: redis.Redis, task_id: str) -> dict[str, Any] | None:
+    item = await client.hget(_result_key(), task_id)
+    if not item:
+        return None
+    return json.loads(item)

agent_pm/storage/tasks.py

@@ -14,7 +14,7 @@
 
 from ..settings import settings
 from ..utils.datetime import utc_now
-from .redis import enqueue_task as redis_enqueue_task, get_redis_pool
+from .redis import enqueue_task as redis_enqueue_task, get_redis_client
 
 logger = logging.getLogger(__name__)
 
@@ -187,7 +187,7 @@ async def get_task_queue() -> TaskQueue:
     if _task_queue is None:
         backend = settings.task_queue_backend
         if backend == "redis":
-            pool = await get_redis_pool()
+            client = await get_redis_client()
 
             class RedisTaskQueue(TaskQueue):
                 async def enqueue(  # type: ignore[override]
@@ -198,7 +198,14 @@ async def enqueue(  # type: ignore[override]
                     max_retries: int = 3,
                     **kwargs: Any,
                 ) -> str:
-                    return await redis_enqueue_task(pool, name, *args, **kwargs)
+                    payload = {
+                        "task_id": uuid.uuid4().hex,
+                        "name": name,
+                        "args": args,
+                        "kwargs": kwargs,
+                        "max_retries": max_retries,
+                    }
+                    return await redis_enqueue_task(client, name, payload)
 
             _task_queue = RedisTaskQueue(max_workers=settings.task_queue_workers)
         else:

pyproject.toml

@@ -21,8 +21,8 @@ dependencies = [
     "prometheus-client",
     "sqlalchemy>=2.0",
     "asyncpg",
-    "redis>=5.0",
-    "arq",
+    "redis[async]>=5.0",
+    "rq",
     "opentelemetry-api",
     "opentelemetry-sdk",
     "opentelemetry-instrumentation-fastapi",

tests/storage/test_redis.py

@@ -1,53 +1,54 @@
-import asyncio
+import json
 
 import pytest
 
 from agent_pm.storage import redis
 
 
-class FakeJobResult:
-    def __init__(self, job_id: str, status: str):
-        self.job_status = type("JobStatus", (), {"value": status})()
-        self.result = {"ok": True}
-        self.start_time = None
-        self.finish_time = None
-
-
-class FakeRedisPool:
+class DummyRedis:
     def __init__(self):
-        self.jobs: dict[str, tuple[str, tuple, dict]] = {}
+        self.items: list[str] = []
+        self.results: dict[str, str] = {}
 
-    async def enqueue_job(self, function_name, *args, **kwargs):
-        job_id = f"job-{len(self.jobs) + 1}"
-        self.jobs[job_id] = (function_name, args, kwargs)
-        return type("Job", (), {"job_id": job_id})()
+    async def rpush(self, key: str, value: str) -> None:
+        self.items.append(value)
 
-    async def job_status(self, job_id):
-        if job_id not in self.jobs:
+    async def lpop(self, key: str):
+        if not self.items:
             return None
-        return FakeJobResult(job_id, "complete")
+        return self.items.pop(0)
+
+    async def hset(self, key: str, field: str, value: str) -> None:
+        self.results[field] = value
+
+    async def hget(self, key: str, field: str):
+        return self.results.get(field)
 
 
 @pytest.mark.asyncio
-async def test_enqueue_task_records_job(monkeypatch):
-    pool = FakeRedisPool()
+async def test_enqueue_task_pushes_payload():
+    client = DummyRedis()
+    payload = {"args": [1], "kwargs": {"foo": "bar"}}
 
-    job_id = await redis.enqueue_task(pool, "do_work", 1, foo="bar")
+    task_id = await redis.enqueue_task(client, "do_work", payload)
 
-    assert job_id in pool.jobs
-    function_name, args, kwargs = pool.jobs[job_id]
-    assert function_name == "do_work"
-    assert args == (1,)
-    assert kwargs == {"foo": "bar"}
+    assert isinstance(task_id, str)
+    stored = json.loads(client.items[0])
+    assert stored["name"] == "do_work"
+    assert stored["args"] == [1]
+    assert stored["kwargs"] == {"foo": "bar"}
 
 
 @pytest.mark.asyncio
-async def test_get_job_status_returns_dict():
-    pool = FakeRedisPool()
-    job_id = await redis.enqueue_task(pool, "do_work")
-
-    status = await redis.get_job_status(pool, job_id)
-    assert status is not None
-    assert status["job_id"] == job_id
-    assert status["status"] == "complete"
-    assert status["result"] == {"ok": True}
+async def test_pop_and_result_roundtrip():
+    client = DummyRedis()
+    payload = {"args": [], "kwargs": {}}
+    task_id = await redis.enqueue_task(client, "noop", payload)
+
+    popped = await redis.pop_task(client)
+    assert popped is not None
+    assert popped["task_id"] == task_id
+
+    await redis.set_task_result(client, task_id, {"status": "done"})
+    result = await redis.get_task_result(client, task_id)
+    assert result == {"status": "done"}