chore: checkpoint current backend/runtime changes

backend/src/v1/agent/dependencies.py

@@ -1,57 +1,98 @@
 from __future__ import annotations
 
-from typing import Any, cast
+import asyncio
+from typing import Any
 from uuid import UUID
 
 from fastapi import Depends
+from redis.asyncio import Redis
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from core.agent.infrastructure.events.redis_stream import RedisStreamEventStore
-from core.agent.infrastructure.queue.tasks import run_command_task
+from core.agent.infrastructure.queue.tasks import (
+    run_command_task,
+    run_command_task_bulk,
+    run_command_task_critical,
+)
 from core.config.settings import config
 from core.db import get_db
-from services.base.redis import redis_service
+from services.base.redis import get_or_init_redis_client
 from v1.agent.repository import AgentRepository
 from v1.agent.service import AgentService
 
+DEDUP_WAIT_RETRIES = 20
+DEDUP_WAIT_SECONDS = 0.05
+DEDUP_LOCK_SECONDS = 300
+DEDUP_INFLIGHT_MARKER = "__inflight__"
 
-class CeleryQueueClient:
+
+class TaskiqQueueClient:
     def __init__(self) -> None:
-        self._redis = redis_service.get_client()
+        self._redis: Redis | None = None
+
+    async def _get_redis(self) -> Redis:
+        if self._redis is None:
+            self._redis = await get_or_init_redis_client()
+        return self._redis
+
+    @staticmethod
+    def _select_queue_task(command: dict[str, object]) -> Any:
+        queue = str(command.get("queue", "default")).strip().lower()
+        if queue == "critical":
+            return run_command_task_critical
+        if queue == "bulk":
+            return run_command_task_bulk
+        return run_command_task
 
     async def enqueue(
         self, *, command: dict[str, object], dedup_key: str | None
     ) -> str:
+        redis_client = await self._get_redis()
         redis_key = None
         if dedup_key:
             redis_key = f"agent:dedup:{dedup_key}"
-            locked = await self._redis.set(redis_key, "__inflight__", nx=True, ex=300)
+            locked = await redis_client.set(
+                redis_key,
+                DEDUP_INFLIGHT_MARKER,
+                nx=True,
+                ex=DEDUP_LOCK_SECONDS,
+            )
             if not locked:
-                existing = await self._redis.get(redis_key)
-                if existing and existing != "__inflight__":
-                    return existing
+                for _ in range(DEDUP_WAIT_RETRIES):
+                    existing = await redis_client.get(redis_key)
+                    if existing and existing != DEDUP_INFLIGHT_MARKER:
+                        return existing
+                    await asyncio.sleep(DEDUP_WAIT_SECONDS)
+                raise RuntimeError("duplicate request is still in progress")
 
         payload = dict(command)
-        if dedup_key:
-            payload["dedup_key"] = dedup_key
-        delay = getattr(run_command_task, "delay")
-        result = delay(payload)
-        task_id = str(result.id)
-        if redis_key is not None:
-            await self._redis.set(redis_key, task_id, ex=300)
-        return task_id
+        queue_task = self._select_queue_task(payload)
+        try:
+            result = await queue_task.kiq(payload)
+            task_id = str(result.task_id)
+            if redis_key is not None:
+                await redis_client.set(redis_key, task_id, ex=DEDUP_LOCK_SECONDS)
+            return task_id
+        except Exception:
+            if redis_key is not None:
+                await redis_client.delete(redis_key)
+            raise
 
 
 class RedisEventStream:
     def __init__(self) -> None:
-        settings = cast(Any, config)
-        client = redis_service.get_client()
-        self._store = RedisStreamEventStore(
-            client=client,
-            stream_prefix=settings.agent_runtime.redis_stream_prefix,
-            read_count=settings.agent_runtime.redis_stream_read_count,
-            block_ms=settings.agent_runtime.redis_stream_block_ms,
-        )
+        self._store: RedisStreamEventStore | None = None
+
+    async def _get_store(self) -> RedisStreamEventStore:
+        if self._store is None:
+            client = await get_or_init_redis_client()
+            self._store = RedisStreamEventStore(
+                client=client,
+                stream_prefix=config.agent_runtime.redis_stream_prefix,
+                read_count=config.agent_runtime.redis_stream_read_count,
+                block_ms=config.agent_runtime.redis_stream_block_ms,
+            )
+        return self._store
 
     async def read(
         self,
@@ -59,7 +100,8 @@ class RedisEventStream:
         session_id: str,
         last_event_id: str | None,
     ) -> list[dict[str, Any]]:
-        rows = await self._store.read_events(
+        store = await self._get_store()
+        rows = await store.read_events(
             session_id=UUID(session_id),
             last_event_id=last_event_id,
         )
@@ -69,6 +111,6 @@ class RedisEventStream:
 def get_agent_service(session: AsyncSession = Depends(get_db)) -> AgentService:
     return AgentService(
         repository=AgentRepository(session),
-        queue=CeleryQueueClient(),
+        queue=TaskiqQueueClient(),
         stream=RedisEventStream(),
     )