253 lines
9.0 KiB
Python
253 lines
9.0 KiB
Python
"""HandoffTransport - Agent 间 Handoff 传输层抽象
|
||
|
||
提供统一的传输接口，支持：
|
||
- InProcessHandoffTransport: 进程内 asyncio.Queue 传输（Expert Team 模式）
|
||
- RedisHandoffTransport: Redis Pub/Sub 传输（分布式模式）
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import logging
|
||
import os
|
||
from typing import AsyncIterator, Awaitable, Callable
|
||
|
||
import redis.asyncio as aioredis
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class HandoffTransport:
|
||
"""Handoff 传输层协议
|
||
|
||
定义 Agent 间 Handoff 消息的传输接口，支持多种底层实现。
|
||
"""
|
||
|
||
async def send(self, channel: str, message: dict) -> None:
|
||
"""发送消息到指定频道"""
|
||
...
|
||
|
||
async def listen(self, channel: str) -> AsyncIterator[dict]:
|
||
"""监听指定频道的消息（异步生成器）"""
|
||
...
|
||
yield
|
||
|
||
def register_handler(self, channel: str, handler: Callable[[dict], Awaitable[None]]) -> None:
|
||
"""注册消息处理器"""
|
||
...
|
||
|
||
|
||
class InProcessHandoffTransport:
|
||
"""进程内 Handoff 传输
|
||
|
||
基于 asyncio.Queue 实现的进程内消息传输，适用于 Expert Team 模式下
|
||
同一进程内的 Agent 间 Handoff，无需 Redis 依赖。
|
||
|
||
支持广播模式：同一频道可有多个消费者，每条消息会投递到所有消费者。
|
||
"""
|
||
|
||
_CLOSED_SENTINEL: dict = {"_closed": True}
|
||
_MAX_QUEUE_SIZE: int = 1024
|
||
|
||
def __init__(self) -> None:
|
||
self._channels: dict[str, list[asyncio.Queue[dict]]] = {}
|
||
self._handlers: dict[str, list[Callable[[dict], Awaitable[None]]]] = {}
|
||
self._closed = False
|
||
|
||
async def send(self, channel: str, message: dict) -> None:
|
||
"""发送消息到指定频道
|
||
|
||
将消息放入该频道所有消费者的队列中，并调用所有已注册的处理器。
|
||
如果频道不存在（无消费者），消息将被丢弃。
|
||
"""
|
||
queues = self._channels.get(channel, [])
|
||
for queue in queues:
|
||
await queue.put(message)
|
||
|
||
handlers = self._handlers.get(channel, [])
|
||
for handler in handlers:
|
||
try:
|
||
await handler(message)
|
||
except Exception as e:
|
||
logger.error(f"InProcessHandoffTransport handler error on channel '{channel}': {e}")
|
||
|
||
async def listen(self, channel: str) -> AsyncIterator[dict]: # type: ignore[override]
|
||
"""监听指定频道的消息
|
||
|
||
为调用者创建一个独立的队列，通过异步生成器持续产出消息。
|
||
每个调用者获得独立的队列，实现广播语义。
|
||
"""
|
||
queue: asyncio.Queue[dict] = asyncio.Queue(maxsize=self._MAX_QUEUE_SIZE)
|
||
if channel not in self._channels:
|
||
self._channels[channel] = []
|
||
self._channels[channel].append(queue)
|
||
|
||
try:
|
||
while True:
|
||
message = await queue.get()
|
||
if message is self._CLOSED_SENTINEL:
|
||
break
|
||
yield message
|
||
finally:
|
||
# 清理：移除当前消费者的队列
|
||
if channel in self._channels:
|
||
try:
|
||
self._channels[channel].remove(queue)
|
||
except ValueError:
|
||
pass
|
||
if not self._channels[channel]:
|
||
del self._channels[channel]
|
||
|
||
def register_handler(self, channel: str, handler: Callable[[dict], Awaitable[None]]) -> None:
|
||
"""注册消息处理器
|
||
|
||
当频道收到消息时，处理器将被调用。
|
||
"""
|
||
if channel not in self._handlers:
|
||
self._handlers[channel] = []
|
||
self._handlers[channel].append(handler)
|
||
|
||
def close(self) -> None:
|
||
"""关闭传输，清理所有频道和处理器。
|
||
|
||
向所有活跃的监听器队列放入哨兵值，使其能够优雅退出。
|
||
"""
|
||
self._closed = True
|
||
# Put sentinel into every queue so active listeners unblock and exit
|
||
for channel_queues in self._channels.values():
|
||
for queue in channel_queues:
|
||
try:
|
||
queue.put_nowait(self._CLOSED_SENTINEL)
|
||
except asyncio.QueueFull:
|
||
pass
|
||
self._channels.clear()
|
||
self._handlers.clear()
|
||
|
||
|
||
class RedisHandoffTransport:
|
||
"""Redis Pub/Sub Handoff 传输
|
||
|
||
基于 Redis Pub/Sub 实现的分布式消息传输，适用于跨进程、跨节点的
|
||
Agent 间 Handoff。
|
||
|
||
使用懒连接：在首次使用时才建立 Redis 连接，而非初始化时。
|
||
"""
|
||
|
||
def __init__(self, redis_url: str) -> None:
|
||
self._redis_url = redis_url
|
||
self._redis: aioredis.Redis | None = None
|
||
self._active_pubsubs: list[aioredis.client.PubSub] = []
|
||
self._handlers: dict[str, list[Callable[[dict], Awaitable[None]]]] = {}
|
||
self._listen_tasks: dict[str, asyncio.Task] = {}
|
||
|
||
async def _ensure_connection(self) -> aioredis.Redis:
|
||
"""确保 Redis 连接已建立（懒初始化）"""
|
||
if self._redis is None:
|
||
try:
|
||
# DEPLOY-2: 显式传入 REDIS_PASSWORD env var（Helm 注入但 redis-py 不自动读取）。
|
||
self._redis = aioredis.from_url(
|
||
self._redis_url,
|
||
decode_responses=True,
|
||
password=os.environ.get("REDIS_PASSWORD") or None,
|
||
)
|
||
await self._redis.ping()
|
||
logger.info("RedisHandoffTransport connected to Redis")
|
||
except Exception:
|
||
# Reset on failure so future calls retry
|
||
self._redis = None
|
||
raise
|
||
return self._redis
|
||
|
||
async def send(self, channel: str, message: dict) -> None:
|
||
"""发送消息到指定频道
|
||
|
||
通过 Redis PUBLISH 命令将消息发布到频道。
|
||
"""
|
||
redis = await self._ensure_connection()
|
||
await redis.publish(channel, json.dumps(message))
|
||
logger.debug(f"RedisHandoffTransport sent message to channel '{channel}'")
|
||
|
||
async def listen(self, channel: str) -> AsyncIterator[dict]: # type: ignore[override]
|
||
"""监听指定频道的消息
|
||
|
||
订阅 Redis 频道，通过异步生成器持续产出消息。
|
||
"""
|
||
redis = await self._ensure_connection()
|
||
pubsub = redis.pubsub()
|
||
await pubsub.subscribe(channel)
|
||
self._active_pubsubs.append(pubsub)
|
||
|
||
try:
|
||
async for message in pubsub.listen():
|
||
if message["type"] == "message":
|
||
data = message["data"]
|
||
if isinstance(data, str):
|
||
yield json.loads(data)
|
||
else:
|
||
yield json.loads(data.decode())
|
||
except asyncio.CancelledError:
|
||
pass
|
||
finally:
|
||
await pubsub.unsubscribe(channel)
|
||
await pubsub.aclose()
|
||
if pubsub in self._active_pubsubs:
|
||
self._active_pubsubs.remove(pubsub)
|
||
|
||
def register_handler(self, channel: str, handler: Callable[[dict], Awaitable[None]]) -> None:
|
||
"""注册消息处理器并启动后台监听任务
|
||
|
||
注册处理器后，自动启动一个后台 asyncio.Task 来监听该频道的消息。
|
||
"""
|
||
if channel not in self._handlers:
|
||
self._handlers[channel] = []
|
||
self._handlers[channel].append(handler)
|
||
|
||
# 启动后台监听任务
|
||
if channel not in self._listen_tasks:
|
||
self._listen_tasks[channel] = asyncio.create_task(self._background_listen(channel))
|
||
|
||
async def _background_listen(self, channel: str) -> None:
|
||
"""后台监听频道消息并调用处理器"""
|
||
try:
|
||
async for message in self.listen(channel):
|
||
handlers = self._handlers.get(channel, [])
|
||
for handler in handlers:
|
||
try:
|
||
await handler(message)
|
||
except Exception as e:
|
||
logger.error(
|
||
f"RedisHandoffTransport handler error on channel '{channel}': {e}"
|
||
)
|
||
except asyncio.CancelledError:
|
||
pass
|
||
except Exception as e:
|
||
logger.error(f"RedisHandoffTransport listen error on channel '{channel}': {e}")
|
||
|
||
async def close(self) -> None:
|
||
"""关闭传输，取消所有监听任务并关闭 Redis 连接"""
|
||
# 取消所有后台监听任务
|
||
for task in self._listen_tasks.values():
|
||
if not task.done():
|
||
task.cancel()
|
||
try:
|
||
await task
|
||
except asyncio.CancelledError:
|
||
pass
|
||
self._listen_tasks.clear()
|
||
self._handlers.clear()
|
||
|
||
# 关闭所有活跃的 PubSub 连接
|
||
for pubsub in self._active_pubsubs:
|
||
try:
|
||
await pubsub.unsubscribe()
|
||
await pubsub.aclose()
|
||
except Exception:
|
||
pass
|
||
self._active_pubsubs.clear()
|
||
|
||
# 关闭 Redis 连接
|
||
if self._redis is not None:
|
||
await self._redis.close()
|
||
self._redis = None
|