base_agent/uas/comm/adapters/redis_adapter.py

# -*- coding: utf-8 -*-
"""
Redis 适配器
基于 redis-py >= 4.x 的异步客户端
支持：Pub/Sub、Stream（持久化）、Key/Value 消息存储
"""

import asyncio
import json
import logging
from typing import Optional

import redis.asyncio as aioredis

from uas.comm.base import CommBaseAdapter, ConnectionConfig, MessageHandler
from uas.comm.message import Message

logger = logging.getLogger(__name__)


class RedisAdapterComm(CommBaseAdapter):
    """
    Redis 通信适配器
    URL 格式: redis://[:password@]host[:port][/db][?options]
    示例:     redis://:secret@127.0.0.1:6379/0?socket_timeout=5
    """

    def __init__(self, config: ConnectionConfig):
        super().__init__(config)
        self._client:  Optional[aioredis.Redis]  = None
        self._pubsub:  Optional[aioredis.client.PubSub] = None
        self._listen_task: Optional[asyncio.Task] = None

    # ── 生命周期 ──────────────────────────────────────────────────

    async def connect(self) -> None:
        opts = self.config.options
        db   = int(self.config.database) if self.config.database.isdigit() else 0

        self._client = aioredis.Redis(
            host            = self.config.host,
            port            = self.config.port,
            password        = self.config.password or None,
            db              = db,
            socket_timeout  = float(opts.get("socket_timeout", 5)),
            decode_responses = True,
        )
        # 连通性测试
        await self._client.ping()
        self._pubsub  = self._client.pubsub()
        self._running = True
        self.logger.info(f"✅ Redis 已连接: {self.config.host}:{self.config.port} db={db}")

    async def disconnect(self) -> None:
        self._running = False
        if self._listen_task:
            self._listen_task.cancel()
            try:
                await self._listen_task
            except asyncio.CancelledError:
                pass
        if self._pubsub:
            await self._pubsub.close()
        if self._client:
            await self._client.aclose()
        self.logger.info("🔌 Redis 连接已关闭")

    # ── 核心功能 ──────────────────────────────────────────────────

    async def publish(self, message: Message) -> bool:
        if not self._client:
            raise RuntimeError("Redis 未连接")
        try:
            payload = json.dumps(message.to_dict(), ensure_ascii=False)
            receivers = await self._client.publish(message.topic, payload)
            self.logger.debug(f"📤 发布 topic={message.topic!r} receivers={receivers}")
            return True
        except Exception as e:
            self.logger.error(f"发布失败: {e}")
            return False

    async def subscribe(self, topic: str, handler: MessageHandler) -> None:
        if not self._pubsub:
            raise RuntimeError("Redis 未连接")

        self._register_handler(topic, handler)

        # 支持通配符订阅
        if "*" in topic or "?" in topic:
            await self._pubsub.psubscribe(topic)
        else:
            await self._pubsub.subscribe(topic)

        # 启动监听循环（幂等）
        if self._listen_task is None or self._listen_task.done():
            self._listen_task = asyncio.create_task(self._listen_loop())

        self.logger.info(f"📥 订阅 Redis topic: {topic!r}")

    async def unsubscribe(self, topic: str) -> None:
        if self._pubsub:
            if "*" in topic or "?" in topic:
                await self._pubsub.punsubscribe(topic)
            else:
                await self._pubsub.unsubscribe(topic)
        self._handlers.pop(topic, None)
        self.logger.info(f"🚫 取消订阅: {topic!r}")

    # ── Stream 持久化发布（可选）──────────────────────────────────

    async def publish_stream(self, message: Message, maxlen: int = 1000) -> str:
        """使用 Redis Stream 发布（消息持久化，支持消费组）"""
        stream_id = await self._client.xadd(
            message.topic,
            {"data": json.dumps(message.to_dict(), ensure_ascii=False)},
            maxlen=maxlen,
        )
        self.logger.debug(f"📤 Stream 发布 id={stream_id}")
        return stream_id

    # ── 内部监听循环 ──────────────────────────────────────────────

    async def _listen_loop(self):
        self.logger.info("🔄 Redis 监听循环启动")
        try:
            async for raw in self._pubsub.listen():
                if raw["type"] not in ("message", "pmessage"):
                    continue
                try:
                    data    = json.loads(raw["data"])
                    message = Message.from_dict(data)
                    await self._dispatch(message)
                except Exception as e:
                    self.logger.error(f"消息解析异常: {e}", exc_info=True)
        except asyncio.CancelledError:
            self.logger.info("🛑 Redis 监听循环已停止")

    async def health_check(self) -> bool:
        try:
            return bool(await self._client.ping())
        except Exception:
            return False