matchmaking-v2/app/adk/worker.py

"""Redis Stream worker for matchmaking-service — the established agent-mesh pattern.

Consumes tasks from the ``tasks:matchmaking-service`` Redis Stream (consumer group),
executes each via MatchmakingAgentSession, and streams response messages to the
per-task pub/sub channel ``responses:{task_id}``, ending with a ``__task_complete__``
sentinel. Mirrors the sibling services (../matchmaking, interview-service, pathways …).
"""

from __future__ import annotations

import asyncio
import json
import logging

from app.adk.emitter import EventEmitter
from app.agent.session import MatchmakingAgentSession

logger = logging.getLogger(__name__)

AGENT_NAME = "matchmaking-service"


class RedisStreamWorker:
    """Consumes tasks from a Redis Stream, executes via the session handler."""

    def __init__(self, redis_url: str | None):
        self.redis_url = redis_url
        self.stream = f"tasks:{AGENT_NAME}"
        self.group = f"{AGENT_NAME}_workers"
        self.consumer = f"{AGENT_NAME}-1"
        self._redis = None
        self._task: asyncio.Task | None = None

    async def start(self):
        if not self.redis_url:
            logger.info("RedisStreamWorker: no REDIS_URL — idle (HTTP /a2a/tasks still serves).")
            return
        try:
            import redis.asyncio as aioredis  # lazy import
            self._redis = aioredis.from_url(self.redis_url, decode_responses=True)
            await self._redis.ping()
        except Exception as e:  # noqa: BLE001 — boot standalone if Redis is down
            logger.warning("RedisStreamWorker: Redis unavailable (%s) — idle, HTTP /a2a/tasks still serves.", e)
            self._redis = None
            return
        try:
            await self._redis.xgroup_create(self.stream, self.group, id="$", mkstream=True)
            logger.info("Created consumer group %s on %s", self.group, self.stream)
        except Exception:  # noqa: BLE001 — group already exists
            pass
        self._task = asyncio.create_task(self._consume_loop())
        logger.info("RedisStreamWorker started: stream=%s", self.stream)

    async def _consume_loop(self):
        while True:
            try:
                messages = await self._redis.xreadgroup(
                    self.group, self.consumer, {self.stream: ">"}, count=10, block=2000,
                )
                if not messages:
                    continue
                for _stream_name, entries in messages:
                    for entry_id, data in entries:
                        asyncio.create_task(self._process_task(entry_id, data))
            except asyncio.CancelledError:
                logger.info("%s worker shutting down", AGENT_NAME)
                break
            except Exception as e:  # noqa: BLE001
                if "NOGROUP" in str(e):
                    logger.warning("Consumer group lost, recreating: %s", self.group)
                    try:
                        await self._redis.xgroup_create(self.stream, self.group, id="$", mkstream=True)
                    except Exception:  # noqa: BLE001
                        pass
                else:
                    logger.error("Worker consume error: %s", e, exc_info=True)
                await asyncio.sleep(1)

    async def _process_task(self, entry_id: str, data: dict):
        task_id = data.get("task_id", "")
        response_channel = f"responses:{task_id}"
        try:
            task_params = json.loads(data.get("payload", "{}"))
            action = task_params.get("action", "")
            params = {**task_params.get("params", {})}
            user_id = task_params.get("user_id", "")
            user_context = task_params.get("user_context")
            session_start = task_params.get("session_start", False)
            if user_context:
                params["user_context"] = user_context

            emitter = EventEmitter()
            session = MatchmakingAgentSession(websocket=emitter, user_id=user_id)

            async def run_session():
                try:
                    if session_start:
                        await session.on_session_start(params)
                    elif action:
                        await session.on_user_action(action, params)
                finally:
                    emitter.mark_done()

            session_task = asyncio.create_task(run_session())

            while True:
                msg = await emitter.queue.get()
                if msg is None:
                    break
                await self._redis.publish(response_channel, json.dumps(msg))

            await session_task

        except Exception as e:  # noqa: BLE001
            logger.error("Task %s failed: %s", task_id, e, exc_info=True)
            await self._redis.publish(
                response_channel, json.dumps({"type": "agent_error", "message": str(e)}))
        finally:
            await self._redis.publish(
                response_channel, json.dumps({"type": "__task_complete__", "task_id": task_id}))
            await self._redis.xack(self.stream, self.group, entry_id)

    async def stop(self):
        if self._task:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
        if self._redis:
            await self._redis.aclose()
            self._redis = None
        logger.info("%s worker stopped", AGENT_NAME)