feat(channel/message): 新增消息处理核心模块及相关工具类

本次提交新增了渠道消息处理的完整核心模块,包含以下核心功能:
1.  新增会话围栏类,实现会话并发控制与过期清理
2.  新增媒体清理器,实现过期媒体文件自动清理
3.  新增熔断器组件,实现服务降级与故障隔离
4.  新增消息处理器,完成渠道消息的完整流转处理
5.  新增限流器组件,实现渠道级和账户级流量控制
6.  新增链路追踪模块,集成Langfuse实现调用链路监控
7.  新增指标统计模块,实现消息处理全链路指标采集
8.  新增统一消息模型,封装全渠道消息格式
9.  新增块回复流水线,实现流式回复的合并与去重
10. 新增本地媒体存储模块,实现媒体文件的本地管理
11. 新增回复分发器,实现回复内容的有序发送与延迟处理
12. 完善__init__.py导出所有核心模块与工具类
This commit is contained in:
Kris 2026-05-21 10:27:16 +08:00
parent a81279cf61
commit 0b19470c70
18 changed files with 4766 additions and 0 deletions

View File

@ -0,0 +1,114 @@
from yuxi.channel.message.block_reply_pipeline import create_block_reply_pipeline
from yuxi.channel.message.bridge import AgentBridge
from yuxi.channel.message.circuit_breaker import (
CircuitBreaker,
CircuitBreakerConfig,
CircuitBreakerOpenError,
CircuitState,
)
from yuxi.channel.message.dispatch import MessageDispatcher
from yuxi.channel.message.idempotency import (
ClaimStatus,
IdempotencyBackend,
InMemoryBackend,
RedisBackend,
build_idempotency_key,
check_and_set,
claim,
claim_with_status,
clear_inflight,
commit,
consume,
get_backend,
parse_idempotency_key,
release,
release_and_forget,
reset,
reset_sync,
)
from yuxi.channel.message.langfuse_trace import ChannelTrace, create_channel_trace
from yuxi.channel.message.media import shutdown_media
from yuxi.channel.message.metrics import (
MetricsTimer,
channel_agent_duration_ms,
channel_dispatch_duration_ms,
channel_messages_inflight,
channel_messages_total,
channel_rate_limit_rejects_total,
record_agent_duration_ms,
record_dispatch_duration_ms,
record_message,
record_rate_limit_reject,
registry,
set_inflight,
)
from yuxi.channel.message.models import (
ChunkType,
DispatchResult,
GroupContext,
MentionSource,
MessageReceipt,
MessageType,
PeerInfo,
ReplyPayload,
ReplyStage,
StreamingChunk,
UnifiedMessage,
)
from yuxi.channel.message.rate_limiter import ChannelRateLimiter as RateLimiter
from yuxi.channel.message.rate_limiter import rate_limit_manager
__all__ = [
"AgentBridge",
"ChannelTrace",
"ChunkType",
"CircuitBreaker",
"CircuitBreakerConfig",
"CircuitBreakerOpenError",
"CircuitState",
"ClaimStatus",
"DispatchResult",
"GroupContext",
"IdempotencyBackend",
"InMemoryBackend",
"MentionSource",
"MessageDispatcher",
"MessageReceipt",
"MessageType",
"MetricsTimer",
"PeerInfo",
"RateLimiter",
"rate_limit_manager",
"RedisBackend",
"ReplyPayload",
"ReplyStage",
"StreamingChunk",
"UnifiedMessage",
"build_idempotency_key",
"channel_agent_duration_ms",
"channel_dispatch_duration_ms",
"channel_messages_inflight",
"channel_messages_total",
"channel_rate_limit_rejects_total",
"check_and_set",
"claim",
"claim_with_status",
"clear_inflight",
"commit",
"consume",
"create_block_reply_pipeline",
"create_channel_trace",
"get_backend",
"parse_idempotency_key",
"record_agent_duration_ms",
"record_dispatch_duration_ms",
"record_message",
"record_rate_limit_reject",
"registry",
"release",
"release_and_forget",
"reset",
"reset_sync",
"set_inflight",
"shutdown_media",
]

View File

@ -0,0 +1,240 @@
import asyncio
import logging
from collections.abc import Callable, Coroutine
from typing import Any
from yuxi.channel.message.models import ChunkType, ReplyPayload, ReplyStage, StreamingChunk, UnifiedMessage
from yuxi.channel.message.reply_dispatcher import ReplyDispatcher
from yuxi.channel.streaming.block_chunker import BlockReplyCoalescer
from yuxi.channel.streaming.models import BlockReplyCoalescing
logger = logging.getLogger(__name__)
_DEFAULT_COALESCE_MIN_CHARS = 80
_DEFAULT_COALESCE_MAX_CHARS = 2000
_DEFAULT_COALESCE_IDLE_MS = 150.0
_DEFAULT_TIMEOUT_MS = 10_000.0
def _handle_pipeline_task_exception(task: asyncio.Task) -> None:
try:
task.result()
except Exception:
logger.exception("Unhandled exception in pipeline background task")
class BlockReplyPipeline:
def __init__(
self,
dispatcher: ReplyDispatcher,
*,
min_chars: int = _DEFAULT_COALESCE_MIN_CHARS,
max_chars: int = _DEFAULT_COALESCE_MAX_CHARS,
debounce_ms: float = _DEFAULT_COALESCE_IDLE_MS,
timeout_ms: float = _DEFAULT_TIMEOUT_MS,
):
self._dispatcher = dispatcher
self._timeout_ms = timeout_ms
self._sent_payload_keys: set[str] = set()
self._sent_content_keys: set[str] = set()
self._sent_media_urls: set[str] = set()
self._pending_keys: set[str] = set()
self._payload_seq: int = 0
self._streamed_text_fragments: list[str] = []
self._aborted = False
self._did_stream = False
self._did_log_timeout = False
self._send_chain = asyncio.Lock()
coalesce_config = BlockReplyCoalescing(
min_chars=max(1, min_chars),
max_chars=max(max(1, min_chars), max_chars),
idle_ms=int(debounce_ms),
joiner="",
)
self._coalescer = BlockReplyCoalescer(
config=coalesce_config,
on_flush=self._on_coalesced_flush,
is_stopped=lambda: self._aborted,
)
async def enqueue(self, chunk: StreamingChunk) -> None:
if self._aborted:
return
if chunk.stage == ReplyStage.TOOL:
self._coalescer.flush()
await self._enqueue_tool(chunk)
return
if chunk.stage == ReplyStage.BLOCK and chunk.chunk_type == ChunkType.TEXT_DELTA:
if chunk.content:
self._coalescer.enqueue(chunk.content)
return
if chunk.stage == ReplyStage.FINAL:
self._coalescer.flush()
await self._dispatcher.enqueue(chunk)
return
self._coalescer.flush()
await self._dispatcher.enqueue(chunk)
def _on_coalesced_flush(self, text: str) -> None:
payload = ReplyPayload(target_id="", content=text)
task = asyncio.create_task(self._send_payload(payload))
task.add_done_callback(_handle_pipeline_task_exception)
async def _enqueue_tool(self, chunk: StreamingChunk) -> None:
payload = ReplyPayload(
target_id="",
content=chunk.content or "",
)
await self._send_payload(payload)
async def _send_payload(self, payload: ReplyPayload) -> None:
if self._aborted:
return
self._payload_seq += 1
dedup_key = f"{self._payload_seq}|{payload.payload_key}"
content_key = payload.content_key
if dedup_key in self._sent_payload_keys or dedup_key in self._pending_keys:
return
self._pending_keys.add(dedup_key)
try:
async with self._send_chain:
if self._aborted:
return
chunk = StreamingChunk(
stage=ReplyStage.BLOCK,
chunk_type=ChunkType.TEXT_DELTA,
content=payload.content,
)
try:
await asyncio.wait_for(
self._dispatcher.enqueue(chunk),
timeout=self._timeout_ms / 1000,
)
except TimeoutError:
self._aborted = True
if not self._did_log_timeout:
self._did_log_timeout = True
logger.warning(
"Block reply delivery timed out after %.0fms; "
"aborting remaining replies to preserve ordering",
self._timeout_ms,
)
return
self._sent_payload_keys.add(dedup_key)
self._sent_content_keys.add(content_key)
for url in payload.media_urls:
self._sent_media_urls.add(url)
if not payload.media_urls and payload.content.strip():
self._streamed_text_fragments.append(payload.content.strip())
self._did_stream = True
finally:
self._pending_keys.discard(dedup_key)
async def flush(self, *, force: bool = False) -> None:
if force or self._coalescer.has_buffered():
self._coalescer.flush()
async with self._send_chain:
pass
def stop(self) -> None:
self._coalescer.stop()
def has_buffered(self) -> bool:
return self._coalescer.has_buffered()
def did_stream(self) -> bool:
return self._did_stream
def is_aborted(self) -> bool:
return self._aborted
def get_sent_media_urls(self) -> list[str]:
return list(self._sent_media_urls)
def has_sent_payload(self, content: str) -> bool:
stripped = content.strip()
if stripped in self._sent_content_keys:
return True
if not self._did_stream or not self._streamed_text_fragments:
return False
def _normalize(s: str) -> str:
return "".join(s.split())
return _normalize("".join(self._streamed_text_fragments)) == _normalize(stripped)
def track_media_sent(self, media_url: str) -> None:
self._sent_media_urls.add(media_url)
def has_sent_media(self, media_url: str) -> bool:
return media_url in self._sent_media_urls
@property
def streamed_text(self) -> str:
return "".join(self._streamed_text_fragments)
@property
def sent_payload_count(self) -> int:
return len(self._sent_payload_keys)
@property
def sent_content_count(self) -> int:
return len(self._sent_content_keys)
@property
def sent_media_count(self) -> int:
return len(self._sent_media_urls)
async def abort(self) -> None:
self._aborted = True
await self._dispatcher.abort()
async def wait_idle(self, timeout_ms: float | None = None) -> None:
await self._dispatcher.wait_idle(timeout_ms)
def create_block_reply_pipeline(
send_fn: Callable[..., Coroutine[Any, Any, str | None]],
msg: UnifiedMessage,
*,
response_prefix: str = "",
human_delay: tuple[float, float] | None = None,
min_chars: int = _DEFAULT_COALESCE_MIN_CHARS,
max_chars: int = _DEFAULT_COALESCE_MAX_CHARS,
debounce_ms: float = _DEFAULT_COALESCE_IDLE_MS,
timeout_ms: float = _DEFAULT_TIMEOUT_MS,
on_idle: Callable[[], None] | None = None,
on_error: Callable[[Exception], None] | None = None,
before_deliver: Callable[[str], Coroutine[Any, Any, str | None]] | None = None,
) -> BlockReplyPipeline:
dispatcher = ReplyDispatcher(
send_fn=send_fn,
msg=msg,
response_prefix=response_prefix,
human_delay=human_delay,
timeout_ms=timeout_ms,
on_idle=on_idle,
on_error=on_error,
before_deliver=before_deliver,
)
return BlockReplyPipeline(
dispatcher,
min_chars=min_chars,
max_chars=max_chars,
debounce_ms=debounce_ms,
timeout_ms=timeout_ms,
)

View File

@ -0,0 +1,372 @@
import asyncio
import contextlib
import json
import logging
import time
from collections.abc import AsyncIterator, Callable
from yuxi.channel.message.models import (
ChunkType,
DispatchResult,
ReplyStage,
StreamCancellationToken,
StreamCancelledError,
StreamingChunk,
StreamResult,
StreamStatus,
StreamTimeoutError,
UnifiedMessage,
)
from yuxi.channel.protocols import AgentPromptContext, AgentPromptProtocol, AgentToolProtocol
from yuxi.channel.plugins.registry import ChannelPluginRegistry
from yuxi.channel.security.external_content import sanitize_external_content
logger = logging.getLogger(__name__)
_STREAMING_TEXT_STATUSES = frozenset({"init", "loading", "agent_state", "streaming"})
_STREAMING_REASONING_STATUSES = frozenset({"thinking"})
_TOOL_STATUSES = frozenset({"tool_call", "tool_result"})
_TERMINAL_STATUSES = frozenset({"finished", "error", "interrupted", "warning"})
_INTERACTIVE_STATUSES = frozenset({"ask_user_question_required"})
_DEFAULT_TIMEOUT_SECONDS = 300.0
class AgentBridge:
def __init__(
self,
stream_fn: Callable[..., AsyncIterator[bytes]],
channel_manager=None,
):
self._stream_fn = stream_fn
self._channel_manager = channel_manager
def collect_channel_tools(self, channel_type: str) -> list[dict]:
plugin = ChannelPluginRegistry.get(channel_type)
if plugin is None or not isinstance(plugin, AgentToolProtocol):
return []
tools = plugin.get_agent_tools()
return [t.to_openai_schema() for t in tools]
async def execute_channel_tool(
self,
channel_type: str,
tool_name: str,
params: dict,
context: dict,
) -> dict:
plugin = ChannelPluginRegistry.get(channel_type)
if plugin is None or not isinstance(plugin, AgentToolProtocol):
return {"success": False, "error": f"Channel {channel_type} does not support agent tools"}
try:
return await plugin.execute_agent_tool(tool_name, params, context)
except Exception as e:
logger.exception("Channel tool execution failed: %s.%s", channel_type, tool_name)
return {"success": False, "error": str(e)}
# ── 构建 meta ──────────────────────────────────────────
def _build_meta(self, msg: UnifiedMessage) -> tuple[dict, str]:
meta = {
"source": "channel",
"channel_type": msg.channel_type,
"account_id": msg.account_id,
"peer_id": msg.sender.id,
}
plugin = ChannelPluginRegistry.get(msg.channel_type)
safe_content, _detection = sanitize_external_content(msg.content)
query = safe_content
if isinstance(plugin, AgentPromptProtocol):
context = AgentPromptContext(
channel_type=msg.channel_type,
account_id=msg.account_id,
peer_id=msg.sender.id,
peer_name=msg.sender.display_name,
group_id=msg.group.id if msg.group else None,
group_name=msg.group.name if msg.group else None,
thread_id=msg.group.thread_id if msg.group else None,
)
system_prompt = plugin.build_system_prompt(context)
if system_prompt:
meta["channel_system_prompt"] = system_prompt
context_note = plugin.build_context_note(context)
if context_note:
query = f"{context_note}\n{safe_content}"
channel_tools = self.collect_channel_tools(msg.channel_type)
if channel_tools:
meta["channel_tools"] = channel_tools
return meta, query
# ── invoke (原始字节流) ─────────────────────────────────
async def invoke(
self,
msg: UnifiedMessage,
dispatch_result: DispatchResult,
agent_config_id: int,
current_user,
db,
) -> AsyncIterator[bytes]:
if self._channel_manager:
self._channel_manager.inc_active_runs(msg.channel_type, msg.account_id)
try:
meta, query = self._build_meta(msg)
async for chunk in self._stream_fn(
query=query,
agent_config_id=agent_config_id,
thread_id=dispatch_result.thread_id,
meta=meta,
image_content=msg.image_base64,
current_user=current_user,
db=db,
):
yield chunk
finally:
if self._channel_manager:
self._channel_manager.dec_active_runs(msg.channel_type, msg.account_id)
# ── invoke_stream (结构化流) ────────────────────────────
async def invoke_stream(
self,
msg: UnifiedMessage,
dispatch_result: DispatchResult,
agent_config_id: int,
current_user,
db,
*,
cancel_token: StreamCancellationToken | None = None,
on_chunk: Callable[[StreamingChunk], None] | None = None,
timeout: float = _DEFAULT_TIMEOUT_SECONDS,
) -> AsyncIterator[StreamingChunk]:
if self._channel_manager:
self._channel_manager.inc_active_runs(msg.channel_type, msg.account_id)
try:
meta, query = self._build_meta(msg)
async def _stream_all() -> AsyncIterator[bytes]:
async for raw in self._stream_fn(
query=query,
agent_config_id=agent_config_id,
thread_id=dispatch_result.thread_id,
meta=meta,
image_content=msg.image_base64,
current_user=current_user,
db=db,
):
yield raw
stream_iter = _stream_all()
async with contextlib.aclosing(stream_iter):
while True:
if cancel_token is not None and cancel_token.is_cancelled:
raise StreamCancelledError("stream cancelled by cancel token")
try:
raw_chunk = await asyncio.wait_for(
stream_iter.__anext__(),
timeout=timeout,
)
except TimeoutError:
raise StreamTimeoutError(
f"No chunk received within {timeout}s"
)
except StopAsyncIteration:
break
parsed = self._parse_chunk(raw_chunk)
if parsed is None:
continue
if on_chunk is not None:
on_chunk(parsed)
yield parsed
yield StreamingChunk(
stage=ReplyStage.FINAL,
chunk_type=ChunkType.STATUS,
content="finished",
)
finally:
if self._channel_manager:
self._channel_manager.dec_active_runs(msg.channel_type, msg.account_id)
# ── invoke_with_result (结构化结果) ─────────────────────
async def invoke_with_result(
self,
msg: UnifiedMessage,
dispatch_result: DispatchResult,
agent_config_id: int,
current_user,
db,
*,
cancel_token: StreamCancellationToken | None = None,
timeout: float = _DEFAULT_TIMEOUT_SECONDS,
) -> StreamResult:
result = StreamResult(start_time=time.monotonic())
accumulated: list[str] = []
try:
async for chunk in self.invoke_stream(
msg,
dispatch_result,
agent_config_id,
current_user,
db,
cancel_token=cancel_token,
timeout=timeout,
):
result.chunk_count += 1
if chunk.chunk_type == ChunkType.TEXT_DELTA:
result.text_chunk_count += 1
if chunk.content:
accumulated.append(chunk.content)
elif chunk.chunk_type == ChunkType.REASONING:
result.reasoning_chunk_count += 1
elif chunk.chunk_type == ChunkType.TOOL_CALL:
result.tool_chunk_count += 1
elif chunk.chunk_type == ChunkType.TOOL_RESULT:
result.tool_chunk_count += 1
elif chunk.chunk_type == ChunkType.HEARTBEAT:
result.heartbeat_count += 1
elif chunk.chunk_type == ChunkType.CITATION:
result.citation_count += 1
elif chunk.chunk_type == ChunkType.ERROR:
result.error_count += 1
if chunk.content and not result.error_message:
result.error_message = chunk.content
if chunk.metadata and "token_usage" in chunk.metadata:
result.token_usage = chunk.metadata["token_usage"]
result.accumulated_text = "".join(accumulated)
result.status = StreamStatus.COMPLETED
except StreamCancelledError:
result.status = StreamStatus.CANCELLED
result.accumulated_text = "".join(accumulated)
result.error_message = "stream cancelled"
logger.debug("Agent stream cancelled for channel=%s peer=%s", msg.channel_type, msg.sender.id)
except StreamTimeoutError:
result.status = StreamStatus.TIMEOUT
result.accumulated_text = "".join(accumulated)
result.error_message = "stream timeout"
logger.warning("Agent stream timeout for channel=%s peer=%s", msg.channel_type, msg.sender.id)
except Exception as e:
result.status = StreamStatus.ERROR
result.accumulated_text = "".join(accumulated)
result.error_message = f"{type(e).__name__}: {e}"
logger.exception("Agent stream error for channel=%s peer=%s", msg.channel_type, msg.sender.id)
finally:
result.end_time = time.monotonic()
return result
# ── invoke_and_collect (兼容方法) ────────────────────────
async def invoke_and_collect(
self,
msg: UnifiedMessage,
dispatch_result: DispatchResult,
agent_config_id: int,
current_user,
db,
) -> str | None:
result = await self.invoke_with_result(msg, dispatch_result, agent_config_id, current_user, db)
return result.accumulated_text or None
# ── chunk 解析 ─────────────────────────────────────────
def _parse_chunk(self, raw: bytes) -> StreamingChunk | None:
try:
data = json.loads(raw.decode("utf-8"))
except (json.JSONDecodeError, UnicodeDecodeError):
logger.debug("Failed to parse stream chunk (len=%d): %s", len(raw), raw[:200])
return None
status = data.get("status", "")
content = data.get("response", "")
# 流式文本块
if status in _STREAMING_TEXT_STATUSES:
return StreamingChunk(
stage=ReplyStage.BLOCK,
chunk_type=ChunkType.TEXT_DELTA,
content=content,
)
# 推理/思考块
if status in _STREAMING_REASONING_STATUSES:
return StreamingChunk(
stage=ReplyStage.TOOL,
chunk_type=ChunkType.REASONING,
content=content,
)
# 工具调用
if status == "tool_call":
tool_name = data.get("tool_name", "")
tool_input = data.get("tool_input", {})
return StreamingChunk(
stage=ReplyStage.TOOL,
chunk_type=ChunkType.TOOL_CALL,
content=content,
tool_name=tool_name,
tool_input=tool_input if isinstance(tool_input, dict) else {},
)
# 工具结果
if status == "tool_result":
tool_name = data.get("tool_name", "")
tool_output = data.get("tool_output", content)
return StreamingChunk(
stage=ReplyStage.TOOL,
chunk_type=ChunkType.TOOL_RESULT,
content=content,
tool_name=tool_name,
tool_output=tool_output,
)
# 心跳
if status == "heartbeat":
return StreamingChunk(
stage=ReplyStage.BLOCK,
chunk_type=ChunkType.HEARTBEAT,
content=content,
)
# 引用
if status == "citation":
return StreamingChunk(
stage=ReplyStage.BLOCK,
chunk_type=ChunkType.CITATION,
content=content,
)
# 错误
if status == "error":
return StreamingChunk(
stage=ReplyStage.FINAL,
chunk_type=ChunkType.ERROR,
content=content,
)
# 终止态 / 交互态 — 不产出 chunk
if status in _TERMINAL_STATUSES | _INTERACTIVE_STATUSES:
return None
logger.debug("Unknown chunk status: %s", status)
return None

View File

@ -0,0 +1,112 @@
import asyncio
import logging
import time
from dataclasses import dataclass, field
from enum import StrEnum
logger = logging.getLogger(__name__)
class CircuitState(StrEnum):
CLOSED = "closed"
OPEN = "open"
HALF_OPEN = "half_open"
@dataclass(slots=True)
class CircuitBreakerConfig:
failure_threshold: int = 5
recovery_timeout_sec: float = 30.0
half_open_max_requests: int = 1
consecutive_successes_to_close: int = 2
@dataclass
class CircuitBreaker:
"""简单熔断器。
状态机
- CLOSED: 正常通行累计失败计数
- OPEN: 拒绝请求持续 recovery_timeout_sec 后转 HALF_OPEN
- HALF_OPEN: 允许少量探测请求连续成功则转 CLOSED任何失败则回 OPEN
"""
name: str
config: CircuitBreakerConfig = field(default_factory=CircuitBreakerConfig)
_state: CircuitState = CircuitState.CLOSED
_failure_count: int = 0
_success_count: int = 0
_last_failure_time: float = 0.0
_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
async def call(self, coro) -> object:
if not await self._allow_request():
raise CircuitBreakerOpenError(
f"Circuit breaker '{self.name}' is OPEN"
)
try:
result = await coro
except Exception:
await self._record_failure()
raise
else:
await self._record_success()
return result
async def _allow_request(self) -> bool:
async with self._lock:
if self._state == CircuitState.CLOSED:
return True
if self._state == CircuitState.OPEN:
if time.monotonic() - self._last_failure_time >= self.config.recovery_timeout_sec:
self._state = CircuitState.HALF_OPEN
self._success_count = 0
logger.info("Circuit '%s' transitioning OPEN -> HALF_OPEN", self.name)
return True
return False
if self._state == CircuitState.HALF_OPEN:
return True
return True
async def _record_success(self) -> None:
async with self._lock:
if self._state == CircuitState.HALF_OPEN:
self._success_count += 1
if self._success_count >= self.config.consecutive_successes_to_close:
self._state = CircuitState.CLOSED
self._failure_count = 0
self._success_count = 0
logger.info("Circuit '%s' transitioning HALF_OPEN -> CLOSED", self.name)
async def _record_failure(self) -> None:
async with self._lock:
self._failure_count += 1
self._last_failure_time = time.monotonic()
if self._state == CircuitState.HALF_OPEN:
self._state = CircuitState.OPEN
logger.warning(
"Circuit '%s' transitioning HALF_OPEN -> OPEN (probe failed)",
self.name,
)
elif self._state == CircuitState.CLOSED and self._failure_count >= self.config.failure_threshold:
self._state = CircuitState.OPEN
logger.warning(
"Circuit '%s' transitioning CLOSED -> OPEN after %d failures",
self.name,
self._failure_count,
)
@property
def state(self) -> CircuitState:
return self._state
async def reset(self) -> None:
async with self._lock:
self._state = CircuitState.CLOSED
self._failure_count = 0
self._success_count = 0
class CircuitBreakerOpenError(Exception):
pass

View File

@ -0,0 +1,53 @@
import asyncio
import logging
import time
from yuxi.channel.message.models import UnifiedMessage
logger = logging.getLogger(__name__)
_DEFAULT_FENCE_TTL = 300.0
class ConversationFence:
def __init__(self, ttl: float = _DEFAULT_FENCE_TTL):
self._ttl = ttl
self._versions: dict[str, int] = {}
self._version_ts: dict[str, float] = {}
self._locks: dict[str, asyncio.Lock] = {}
self._abort_events: dict[str, asyncio.Event] = {}
@staticmethod
def key_for(msg: UnifiedMessage) -> str:
if msg.group and msg.group.id:
return f"{msg.channel_type}:{msg.account_id}:group:{msg.group.id}"
return f"{msg.channel_type}:{msg.account_id}:dm:{msg.sender.id}"
def enter(self, key: str) -> tuple[int, asyncio.Event]:
now = time.monotonic()
self._gc(now)
self._versions[key] = self._versions.get(key, 0) + 1
self._version_ts[key] = now
old_event = self._abort_events.get(key)
if old_event is not None and not old_event.is_set():
old_event.set()
logger.debug("Foreground fence: aborting previous run for %s", key)
new_event = asyncio.Event()
self._abort_events[key] = new_event
return self._versions[key], new_event
def lock_for(self, key: str) -> asyncio.Lock:
return self._locks.setdefault(key, asyncio.Lock())
def current_version(self, key: str) -> int:
return self._versions.get(key, 0)
def _gc(self, now: float) -> None:
expired = [k for k, ts in self._version_ts.items() if now - ts > self._ttl]
for k in expired:
self._versions.pop(k, None)
self._version_ts.pop(k, None)
self._locks.pop(k, None)
self._abort_events.pop(k, None)

View File

@ -0,0 +1,407 @@
import logging
import time
from collections.abc import Awaitable, Callable
from typing import TYPE_CHECKING
from yuxi.channel.message.models import DispatchResult, PeerKind, UnifiedMessage
from yuxi.channel.protocols import (
ApprovalAction,
ApprovalDecision,
ConfigProtocol,
MessagingProtocol,
SessionResolution,
)
from yuxi.channel.plugins.registry import ChannelPluginRegistry
from yuxi.channel.routing.matcher import MessageContext, RouteMatcher, VerificationContext
from yuxi.channel.routing.models import RouteBinding
from yuxi.channel.routing.session_key import SessionKeyBuilder, session_key_to_thread_id
from yuxi.channel.security.allowlist import AllowlistChecker
from yuxi.channel.security.pairing import PairingManager
from yuxi.repositories.channel_thread_mapping_repo import ChannelThreadMappingRepository
from yuxi.repositories.channel_user_mapping_repo import ChannelUserMappingRepository
if TYPE_CHECKING:
from yuxi.channel.message.langfuse_trace import ChannelTrace
logger = logging.getLogger(__name__)
BeforeDispatchHook = Callable[[UnifiedMessage], Awaitable[bool | None]]
AfterDispatchHook = Callable[[UnifiedMessage, DispatchResult], Awaitable[None]]
class MessageDispatcher:
def __init__(
self,
bindings: list[RouteBinding],
allowlist: AllowlistChecker,
pairing: PairingManager,
*,
matcher: RouteMatcher | None = None,
session_key_builder: SessionKeyBuilder | None = None,
approval_engine=None,
user_mapping_repo: ChannelUserMappingRepository | None = None,
thread_mapping_repo: ChannelThreadMappingRepository | None = None,
on_before_dispatch: BeforeDispatchHook | None = None,
on_after_dispatch: AfterDispatchHook | None = None,
):
self._bindings = bindings
self._matcher = matcher
self._allowlist = allowlist
self._pairing = pairing
self._session_key_builder = session_key_builder or SessionKeyBuilder()
self._approval_engine = approval_engine
self._user_mapping_repo = user_mapping_repo
self._thread_mapping_repo = thread_mapping_repo
self._on_before_dispatch = on_before_dispatch
self._on_after_dispatch = on_after_dispatch
async def dispatch(self, msg: UnifiedMessage, *, trace: "ChannelTrace | None" = None) -> DispatchResult:
command_result = await self._try_handle_command(msg)
if command_result is not None:
return command_result
try:
return await self._dispatch_inner(msg, trace=trace)
except Exception:
logger.exception("Dispatch failed for message: %s", msg.msg_id)
return DispatchResult(success=False, error="dispatch_error")
@staticmethod
def _plugin_has_commands(plugin) -> bool:
return callable(getattr(plugin, "get_commands", None)) and callable(getattr(plugin, "handle_command", None))
async def _try_handle_command(self, msg: UnifiedMessage) -> DispatchResult | None:
text = (msg.content or "").strip()
if not text.startswith("/"):
return None
parts = text[1:].split()
if not parts:
return None
command_name = parts[0].lower()
args = parts[1:]
plugin = ChannelPluginRegistry.get(msg.channel_type)
if plugin is None or not self._plugin_has_commands(plugin):
return None
config = {}
if isinstance(plugin, ConfigProtocol):
try:
account_cfg = await plugin.resolve_account(msg.account_id)
config = account_cfg if account_cfg else {}
except Exception:
logger.debug("Failed to resolve account config for command: %s", command_name)
try:
response = await plugin.handle_command(
config,
command_name,
args,
msg,
None,
)
except Exception:
logger.exception("Command handling failed: %s", command_name)
return DispatchResult(
success=True,
handled_by="command",
command_response="命令执行失败,请稍后重试。",
)
return DispatchResult(
success=True,
handled_by="command",
command_response=response,
)
@staticmethod
def _plugin_has_approval(plugin) -> bool:
return (
callable(getattr(plugin, "check_approval_required", None))
and callable(getattr(plugin, "create_approval_request", None))
and callable(getattr(plugin, "check_approval_status", None))
)
async def _should_check_approval(self, msg: UnifiedMessage) -> bool:
if self._approval_engine is None:
return False
plugin = ChannelPluginRegistry.get(msg.channel_type)
return plugin is not None and self._plugin_has_approval(plugin)
async def _dispatch_inner(self, msg: UnifiedMessage, *, trace: "ChannelTrace | None" = None) -> DispatchResult:
plugin = ChannelPluginRegistry.get(msg.channel_type)
sec_span_id = None
if trace:
sec_span_id = await trace.add_span("security_admission", metadata={"peer_kind": msg.sender.kind.value})
if msg.sender.kind == PeerKind.DIRECT:
try:
dm_result = self._allowlist.check_dm(msg.sender.id)
if not dm_result.allowed:
if dm_result.pairing_required:
req = await self._pairing.upsert_code(msg.channel_type, msg.sender.id, msg.account_id)
if req is not None:
if sec_span_id:
await trace.end_span(sec_span_id, level="ERROR", status_message="pairing_required")
return DispatchResult(success=False, error=f"pairing_required:{req.code}")
if sec_span_id:
await trace.end_span(sec_span_id, level="ERROR", status_message="pairing_check_error")
return DispatchResult(success=False, error="pairing_check_error")
logger.info(
"Sender not in DM allowlist: %s/%s",
msg.channel_type,
msg.sender.id,
)
if sec_span_id:
await trace.end_span(sec_span_id, level="ERROR", status_message="not_in_allowlist")
return DispatchResult(success=False, error="not_in_allowlist")
except Exception:
logger.exception("DM allowlist check failed, blocking message")
if sec_span_id:
await trace.end_span(sec_span_id, level="ERROR", status_message="allowlist_check_error")
return DispatchResult(success=False, error="allowlist_check_error")
auth_scope = "dm"
auth_verified_by = "AllowlistChecker.dm"
else:
try:
if msg.group and msg.group.id:
group_result = self._allowlist.check_group(msg.group.id)
if not group_result.allowed:
logger.info(
"Group not in allowlist: %s/%s",
msg.channel_type,
msg.group.id,
)
if sec_span_id:
await trace.end_span(sec_span_id, level="ERROR", status_message="not_in_allowlist")
return DispatchResult(success=False, error="not_in_allowlist")
from_result = self._allowlist.check_group_allow_from(msg.sender.id)
if not from_result.allowed:
logger.info(
"Sender not in group allow_from: %s/%s",
msg.channel_type,
msg.sender.id,
)
if sec_span_id:
await trace.end_span(sec_span_id, level="ERROR", status_message="not_in_allowlist")
return DispatchResult(success=False, error="not_in_allowlist")
except Exception:
logger.exception("Group allowlist check failed, blocking message")
if sec_span_id:
await trace.end_span(sec_span_id, level="ERROR", status_message="allowlist_check_error")
return DispatchResult(success=False, error="allowlist_check_error")
auth_scope = "group"
auth_verified_by = "AllowlistChecker.group"
if sec_span_id:
await trace.end_span(sec_span_id, status_message="ok")
if self._matcher is None:
self._matcher = RouteMatcher()
await self._matcher.set_bindings(self._bindings)
explicit_target = None
if isinstance(plugin, MessagingProtocol):
session = plugin.resolve_session(msg)
explicit_target = plugin.parse_explicit_target(msg.content)
else:
if msg.sender.kind == PeerKind.DIRECT:
session = SessionResolution(kind="direct", conversation_id=msg.sender.id)
else:
gid = msg.group.id if msg.group else "unknown"
session = SessionResolution(kind="group", conversation_id=gid)
try:
peer_kind = PeerKind(session.kind)
except ValueError:
peer_kind = msg.sender.kind
ctx = MessageContext(
channel=msg.channel_type,
account_id=msg.account_id,
peer_kind=peer_kind,
peer_id=msg.sender.id,
channel_config_id=msg.channel_config_id,
member_role_ids=msg.member_role_ids,
)
ctx._verification_ctx = VerificationContext(
verified_by=auth_verified_by,
scope=auth_scope,
timestamp=time.time(),
)
if msg.group:
ctx.guild_id = msg.group.guild_id
ctx.team_id = msg.group.team_id
ctx.roles = msg.group.roles
if msg.group.thread_id or msg.message_thread_id:
if msg.group.route_peer_kind:
try:
ctx.parent_peer_kind = PeerKind(msg.group.route_peer_kind)
except ValueError:
pass
ctx.parent_peer_id = msg.group.route_peer_id or msg.group.id
route_span_id = None
if trace:
route_span_id = await trace.add_span("route_resolution")
try:
route = await self._matcher.resolve(ctx)
except Exception:
logger.exception("Route resolution failed")
if route_span_id:
await trace.end_span(route_span_id, level="ERROR", status_message="route_error")
return DispatchResult(success=False, error="route_error")
if route.agent_config_id == 0:
if route_span_id:
await trace.end_span(route_span_id, status_message="no_route")
return DispatchResult(success=False, error="no_route")
if route_span_id:
await trace.end_span(route_span_id, metadata={"agent_config_id": str(route.agent_config_id)})
if self._on_before_dispatch is not None:
try:
handled = await self._on_before_dispatch(msg)
if handled:
return DispatchResult(
success=True,
handled_by="hook:before_dispatch",
thread_id=None,
agent_config_id=None,
)
except Exception:
logger.exception("before_dispatch hook failed, continuing")
if await self._should_check_approval(msg):
try:
request = await self._approval_engine.request_approval(
msg.channel_type,
msg.account_id,
{},
ApprovalAction.APPROVE_EXEC,
msg.sender.id,
description=f"Agent 执行: {msg.content[:100]}",
context={"full_text": msg.content},
)
if request is not None:
result = await self._approval_engine.wait_for_decision(
msg.channel_type,
{},
request.id,
)
if result.decision != ApprovalDecision.APPROVED:
return DispatchResult(
success=False,
error=f"approval_denied:{result.reason or '审批未通过'}",
)
except Exception:
logger.exception("Approval check failed, blocking message")
return DispatchResult(
success=False,
error="approval_check_error",
)
agent_id_str = str(route.agent_config_id)
sk_span_id = None
if trace:
sk_span_id = await trace.add_span("session_key_build")
internal_user_id, thread_id = await self._resolve_thread_id(
msg,
agent_id_str,
route.dm_scope,
)
if sk_span_id:
await trace.end_span(sk_span_id, metadata={
"thread_id": thread_id,
"has_internal_user": internal_user_id is not None,
})
reply_sent = False
if explicit_target:
try:
target_plugin = ChannelPluginRegistry.get(explicit_target)
if target_plugin is not None and isinstance(target_plugin, MessagingProtocol):
logger.info(
"Message %s forwarded to explicit target: %s",
msg.msg_id,
explicit_target,
)
except Exception:
logger.exception("Explicit target forwarding failed: %s", explicit_target)
result = DispatchResult(
success=True,
thread_id=thread_id,
agent_config_id=route.agent_config_id,
internal_user_id=internal_user_id,
reply_sent=reply_sent,
)
if self._on_after_dispatch is not None:
try:
await self._on_after_dispatch(msg, result)
except Exception:
logger.exception("after_dispatch hook failed")
return result
async def _resolve_thread_id(
self,
msg: UnifiedMessage,
agent_id_str: str,
dm_scope: str,
) -> tuple[str | None, str]:
internal_user_id = None
if self._user_mapping_repo is not None:
try:
user_mapping = await self._user_mapping_repo.resolve_user(
msg.channel_type,
msg.sender.id,
)
internal_user_id = user_mapping.internal_user_id
except Exception:
logger.exception("Failed to resolve user mapping")
user_id = internal_user_id or msg.sender.id
if msg.sender.kind == PeerKind.DIRECT:
session_key = self._session_key_builder.build_dm_session(
agent_id_str,
msg.channel_type,
msg.account_id,
user_id,
dm_scope=dm_scope,
)
else:
session_key = self._session_key_builder.build_group_session(
agent_id_str,
msg.channel_type,
msg.sender.kind.value,
user_id,
)
thread_id = session_key_to_thread_id(session_key)
if internal_user_id and self._thread_mapping_repo is not None:
try:
channel_chat_id = msg.group.id if msg.group and msg.group.id else msg.sender.id
await self._thread_mapping_repo.resolve_thread(
msg.channel_type,
channel_chat_id,
internal_user_id,
agent_id=agent_id_str,
thread_id=thread_id,
)
except Exception:
logger.exception("Failed to persist thread mapping")
return internal_user_id, thread_id

View File

@ -0,0 +1,346 @@
import asyncio
import logging
import random
import time
import uuid
from collections.abc import Callable, Coroutine
from dataclasses import dataclass, field
from enum import StrEnum
from typing import Any
logger = logging.getLogger(__name__)
class DurableStrategy(StrEnum):
REQUIRED = "required"
BEST_EFFORT = "best_effort"
DISABLED = "disabled"
class MessageSendState(StrEnum):
IDLE = "idle"
RENDERING = "rendering"
PREVIEWING = "previewing"
SENDING = "sending"
SENT = "sent"
SUPPRESSED = "suppressed"
PARTIAL_FAILED = "partial_failed"
FAILED = "failed"
UNKNOWN_AFTER_SEND = "unknown_after_send"
FINALIZING = "finalizing"
EDITING = "editing"
EDITED = "edited"
DELETING = "deleting"
DELETED = "deleted"
CANCELLED = "cancelled"
class MessageReceiptPartKind(StrEnum):
TEXT = "text"
MEDIA = "media"
VOICE = "voice"
CARD = "card"
PREVIEW = "preview"
UNKNOWN = "unknown"
@dataclass
class MessageReceiptPart:
platform_message_id: str
kind: MessageReceiptPartKind = MessageReceiptPartKind.UNKNOWN
index: int = 0
thread_id: str | None = None
reply_to_id: str | None = None
@dataclass
class DurableMessageReceipt:
primary_platform_message_id: str = ""
platform_message_ids: list[str] = field(default_factory=list)
parts: list[MessageReceiptPart] = field(default_factory=list)
thread_id: str | None = None
reply_to_id: str | None = None
edit_token: str | None = None
delete_token: str | None = None
sent_at: float = 0.0
metadata: dict[str, Any] = field(default_factory=dict)
def __post_init__(self):
if self.sent_at == 0.0:
self.sent_at = time.time()
if not self.platform_message_ids and self.primary_platform_message_id:
self.platform_message_ids = [self.primary_platform_message_id]
@dataclass
class MessageSendContext:
target_id: str
content: str
id: str = ""
channel: str = ""
account_id: str | None = None
reply_to_id: str | None = None
thread_id: str | None = None
strategy: DurableStrategy = DurableStrategy.BEST_EFFORT
receipt: DurableMessageReceipt | None = None
previous_receipt: DurableMessageReceipt | None = None
state: MessageSendState = MessageSendState.IDLE
error: str | None = None
attempt: int = 1
retry_count: int = 0
max_retries: int = 3
min_delay_ms: int = 300
max_delay_ms: int = 30_000
jitter: float = 0.0
metadata: dict[str, Any] = field(default_factory=dict)
parts: list["MessageSendContext"] = field(default_factory=list)
_on_commit: Callable[..., Coroutine[Any, Any, None]] | None = field(default=None, repr=False)
_on_fail: Callable[..., Coroutine[Any, Any, None]] | None = field(default=None, repr=False)
def __post_init__(self):
if not self.id:
self.id = f"{self.channel or 'msg'}:{self.target_id}:{uuid.uuid4().hex[:8]}"
@property
def is_terminal(self) -> bool:
return self.state in (
MessageSendState.SENT,
MessageSendState.SUPPRESSED,
MessageSendState.PARTIAL_FAILED,
MessageSendState.FAILED,
MessageSendState.CANCELLED,
)
async def render(self) -> str:
self.state = MessageSendState.RENDERING
return self.content
def _backoff_delay(self, attempt: int) -> int:
base = self.min_delay_ms * (2 ** (attempt - 1))
delay = min(base, self.max_delay_ms)
if self.jitter > 0:
offset = (random.random() * 2 - 1) * self.jitter
delay = int(delay * (1 + offset))
return max(0, delay)
async def send(self, send_fn: Callable[..., Coroutine[Any, Any, str | None]]) -> DurableMessageReceipt | None:
self.state = MessageSendState.SENDING
last_error = None
total_attempts = self.max_retries + 1
for attempt_idx in range(total_attempts):
self.attempt = attempt_idx + 1
try:
message_id = await send_fn(self.content)
if message_id:
self.state = MessageSendState.SENT
self.receipt = DurableMessageReceipt(primary_platform_message_id=message_id)
return self.receipt
self.state = MessageSendState.SENT
return None
except Exception as e:
last_error = str(e)
self.retry_count = attempt_idx + 1
logger.warning(
"Message send attempt %d/%d failed: %s",
attempt_idx + 1,
total_attempts,
e,
)
if attempt_idx < total_attempts - 1:
delay = self._backoff_delay(attempt_idx + 1)
if delay > 0:
await asyncio.sleep(delay / 1000)
self.state = MessageSendState.FAILED
self.error = last_error
return None
async def send_batch(
self,
contents: list[str],
send_fn: Callable[..., Coroutine[Any, Any, str | None]],
) -> list[DurableMessageReceipt | None]:
if not contents:
return []
results: list[DurableMessageReceipt | None] = []
self.parts.clear()
failed_count = 0
for i, content in enumerate(contents):
part = MessageSendContext(
target_id=self.target_id,
content=content,
id=f"{self.id}#{i}",
channel=self.channel,
account_id=self.account_id,
reply_to_id=self.reply_to_id,
thread_id=self.thread_id,
strategy=self.strategy,
max_retries=self.max_retries,
min_delay_ms=self.min_delay_ms,
max_delay_ms=self.max_delay_ms,
jitter=self.jitter,
)
receipt = await part.send(send_fn)
results.append(receipt)
self.parts.append(part)
if receipt is None and part.state == MessageSendState.FAILED:
failed_count += 1
total = len(contents)
if failed_count == total:
self.state = MessageSendState.FAILED
self.error = f"All {total} parts failed"
elif failed_count > 0:
self.state = MessageSendState.PARTIAL_FAILED
self.error = f"{failed_count}/{total} parts failed"
else:
self.state = MessageSendState.SENT
return results
def mark_suppressed(self, reason: str = "") -> None:
self.state = MessageSendState.SUPPRESSED
self.error = reason
async def edit(
self, edit_fn: Callable[..., Coroutine[Any, Any, str | None]], new_content: str
) -> DurableMessageReceipt | None:
if self.receipt is None:
logger.warning("Cannot edit message without receipt")
return None
self.state = MessageSendState.EDITING
try:
new_id = await edit_fn(self.receipt.primary_platform_message_id, new_content)
if new_id:
self.receipt.primary_platform_message_id = new_id
if new_id not in self.receipt.platform_message_ids:
self.receipt.platform_message_ids.append(new_id)
self.state = MessageSendState.EDITED
return self.receipt
except Exception as e:
self.state = MessageSendState.FAILED
self.error = str(e)
return None
async def delete(self, delete_fn: Callable[..., Coroutine[Any, Any, None]]) -> bool:
if self.receipt is None:
return False
self.state = MessageSendState.DELETING
try:
await delete_fn(self.receipt.primary_platform_message_id)
self.state = MessageSendState.DELETED
return True
except Exception as e:
self.state = MessageSendState.FAILED
self.error = str(e)
return False
def mark_cancelled(self) -> None:
self.state = MessageSendState.CANCELLED
def mark_unknown_after_send(self) -> None:
self.state = MessageSendState.UNKNOWN_AFTER_SEND
async def commit(self) -> None:
if self._on_commit:
await self._on_commit(self.receipt)
async def fail(self, error: Exception | None = None) -> None:
if self._on_fail:
if error is None:
error = Exception(self.error or "send failed")
await self._on_fail(error)
class OutboundBridge:
def __init__(
self,
send_text_fn: Callable[..., Coroutine[Any, Any, str | None]] | None = None,
send_media_fn: Callable[..., Coroutine[Any, Any, str | None]] | None = None,
send_payload_fn: Callable[..., Coroutine[Any, Any, str | None]] | None = None,
):
self._send_text = send_text_fn
self._send_media = send_media_fn
self._send_payload = send_payload_fn
async def text(
self,
target_id: str,
content: str,
*,
reply_to_id: str | None = None,
thread_id: str | None = None,
) -> str | None:
if not self._send_text:
raise RuntimeError("OutboundBridge: send_text not configured")
return await self._send_text(target_id, content, reply_to_id=reply_to_id, thread_id=thread_id)
async def media(
self,
target_id: str,
media_url: str,
text: str = "",
*,
reply_to_id: str | None = None,
thread_id: str | None = None,
audio_as_voice: bool = False,
) -> str | None:
if not self._send_media:
raise RuntimeError("OutboundBridge: send_media not configured")
return await self._send_media(
target_id, media_url, text, reply_to_id=reply_to_id, thread_id=thread_id, audio_as_voice=audio_as_voice
)
async def payload(
self,
target_id: str,
payload: Any,
*,
reply_to_id: str | None = None,
thread_id: str | None = None,
) -> str | None:
if not self._send_payload:
raise RuntimeError("OutboundBridge: send_payload not configured")
return await self._send_payload(target_id, payload, reply_to_id=reply_to_id, thread_id=thread_id)
class DurableSendContextManager:
def __init__(
self,
ctx: "MessageSendContext",
*,
on_commit: Callable[..., Coroutine[Any, Any, None]] | None = None,
on_fail: Callable[..., Coroutine[Any, Any, None]] | None = None,
):
self.ctx = ctx
self.ctx._on_commit = on_commit
self.ctx._on_fail = on_fail
async def __aenter__(self) -> "MessageSendContext":
return self.ctx
async def __aexit__(self, exc_type, exc_val, exc_tb) -> bool:
if exc_type is not None:
await self.ctx.fail(exc_val)
return False
async def send_durable_message_batch(
ctx: "MessageSendContext",
send_fn: Callable[..., Coroutine[Any, Any, str | None]],
*,
on_commit: Callable[..., Coroutine[Any, Any, None]] | None = None,
on_fail: Callable[..., Coroutine[Any, Any, None]] | None = None,
) -> "DurableMessageReceipt | None":
if on_commit:
ctx._on_commit = on_commit
if on_fail:
ctx._on_fail = on_fail
await ctx.render()
result = await ctx.send(send_fn)
if result is not None and ctx.state in (MessageSendState.SENT, MessageSendState.SUPPRESSED):
await ctx.commit()
else:
await ctx.fail()
return result

View File

@ -0,0 +1,183 @@
import logging
from sqlalchemy import select
from yuxi.channel.message.bridge import AgentBridge
from yuxi.channel.message.block_reply_pipeline import create_block_reply_pipeline
from yuxi.channel.message.models import ChunkType, DispatchResult, ReplyStage, StreamingChunk, UnifiedMessage
from yuxi.channel.protocols import OutboundProtocol
from yuxi.channel.plugins.registry import ChannelPluginRegistry
from yuxi.channel.session import SessionRecorder
from yuxi.repositories.conversation_repository import ConversationRepository
from yuxi.services.chat_service import stream_agent_chat
from yuxi.storage.postgres.manager import pg_manager
from yuxi.storage.postgres.models_business import User
logger = logging.getLogger(__name__)
async def channel_message_handler(
msg: UnifiedMessage,
result: DispatchResult,
channel_manager=None,
) -> None:
channel_type = msg.channel_type
account_id = msg.account_id
logger.info(
"channel_turn start: msg_id=%s channel=%s agent_config_id=%s",
msg.msg_id,
channel_type,
result.agent_config_id,
)
user = await _resolve_user(result, msg.msg_id, channel_type)
if user is None:
logger.error(
"channel_turn drop: msg_id=%s channel=%s reason=no_user",
msg.msg_id,
channel_type,
)
return
if channel_manager:
channel_manager.inc_active_runs(channel_type, account_id)
channel_manager.update_last_message_at(channel_type, account_id)
try:
bridge = AgentBridge(stream_fn=stream_agent_chat, channel_manager=channel_manager)
async with pg_manager.get_async_session_context() as db:
await _record_session(msg, result, user, db)
pipeline = create_block_reply_pipeline(
send_fn=_make_send_fn(msg, result, channel_manager),
msg=msg,
group_msg_id=msg.group.id if msg.group else None,
)
async for chunk in bridge.invoke_stream(
msg,
result,
result.agent_config_id,
user,
db,
):
await pipeline.enqueue(chunk)
await pipeline.wait_idle()
if result.thread_id:
await _touch_session(msg, result, db)
except Exception:
logger.exception(
"channel_turn error: msg_id=%s channel=%s",
msg.msg_id,
channel_type,
)
finally:
if channel_manager:
channel_manager.dec_active_runs(channel_type, account_id)
def _make_send_fn(msg: UnifiedMessage, result: DispatchResult, channel_manager):
channel_type = msg.channel_type
plugin = ChannelPluginRegistry.get(channel_type)
if plugin is None or not isinstance(plugin, OutboundProtocol):
return _noop_send_fn
target_id = msg.group.id if msg.group and msg.group.id else msg.sender.id
async def send_fn(chunk: StreamingChunk) -> None:
try:
if chunk.stage == ReplyStage.BLOCK and chunk.chunk_type == ChunkType.TEXT_DELTA:
await plugin.send_text(
target_id,
chunk.content,
reply_to_id=result.thread_id,
)
except Exception:
logger.exception(
"channel_turn deliver_reply error: msg_id=%s channel=%s",
msg.msg_id,
channel_type,
)
return send_fn
async def _noop_send_fn(chunk: StreamingChunk) -> None:
pass
async def _record_session(
msg: UnifiedMessage,
result: DispatchResult,
user: User,
db,
) -> None:
session_key = result.session_key or f"{msg.channel_type}:{msg.account_id}:{msg.sender.id}"
try:
conv_repo = ConversationRepository(db)
recorder = SessionRecorder(conv_repo)
await recorder.record_or_update(
msg,
result,
session_key,
user_id=str(user.id),
title=msg.sender.display_name or str(msg.sender.id),
)
except Exception:
logger.exception(
"channel_turn session record failed: msg_id=%s channel=%s",
msg.msg_id,
msg.channel_type,
)
async def _touch_session(
msg: UnifiedMessage,
result: DispatchResult,
db,
) -> None:
try:
conv_repo = ConversationRepository(db)
recorder = SessionRecorder(conv_repo)
await recorder.touch(msg, result.thread_id)
except Exception:
logger.exception(
"channel_turn session touch failed: msg_id=%s thread=%s",
msg.msg_id,
result.thread_id,
)
async def _resolve_user(
result: DispatchResult,
msg_id: str,
channel_type: str,
) -> User | None:
async with pg_manager.get_async_session_context() as db:
if result.internal_user_id:
try:
uid = int(result.internal_user_id)
except (ValueError, TypeError):
uid = None
if uid is not None:
user = (await db.execute(select(User).where(User.id == uid))).scalar_one_or_none()
if user is not None:
logger.info(
"channel_turn resolve_user: msg_id=%s channel=%s resolved_by=internal_user_id",
msg_id,
channel_type,
)
return user
logger.warning(
"channel_turn resolve_user: msg_id=%s channel=%s internal_user_id=%s no match",
msg_id,
channel_type,
result.internal_user_id,
)
return None

View File

@ -0,0 +1,340 @@
import asyncio
import logging
import os
import time
from abc import ABC, abstractmethod
from enum import StrEnum
from cachetools import LRUCache
logger = logging.getLogger(__name__)
_DEFAULT_TTL_SECONDS = 300
_IDEMPOTENCY_KEY_SEPARATOR = ":"
_IDEMPOTENCY_NONCE_MARKER = ":nonce:"
class ClaimStatus(StrEnum):
INVALID = "invalid"
DUPLICATE = "duplicate"
INFLIGHT = "inflight"
CLAIMED = "claimed"
class IdempotencyBackend(ABC):
"""幂等性后端抽象基类"""
@abstractmethod
async def claim(self, msg_id: str, ttl: int = _DEFAULT_TTL_SECONDS) -> bool: ...
@abstractmethod
async def claim_with_status(self, msg_id: str) -> ClaimStatus: ...
@abstractmethod
async def check_and_set(self, msg_id: str, ttl: int = _DEFAULT_TTL_SECONDS) -> bool: ...
@abstractmethod
async def commit(self, msg_id: str) -> None: ...
@abstractmethod
async def release(self, msg_id: str) -> None: ...
@abstractmethod
async def release_and_forget(self, msg_id: str) -> None: ...
@abstractmethod
async def clear_inflight(self, msg_id: str | None = None) -> None: ...
@abstractmethod
async def reset(self) -> None: ...
class InMemoryBackend(IdempotencyBackend):
def __init__(self) -> None:
self._cache: LRUCache = LRUCache(maxsize=10_000)
self._inflight: set[str] = set()
self._lock = asyncio.Lock()
def _prune_expired(self, now: float) -> None:
expired = [mid for mid, (ts, _ttl) in self._cache.items() if now - ts >= _ttl]
for mid in expired:
del self._cache[mid]
self._inflight.discard(mid)
async def check_and_set(self, msg_id: str, ttl: int = _DEFAULT_TTL_SECONDS) -> bool:
if not msg_id:
return False
async with self._lock:
now = time.monotonic()
if msg_id in self._cache:
timestamp, entry_ttl = self._cache[msg_id]
if now - timestamp < entry_ttl:
return False
self._prune_expired(now)
self._cache[msg_id] = (now, ttl)
return True
async def claim(self, msg_id: str, ttl: int = _DEFAULT_TTL_SECONDS) -> bool:
if not msg_id:
return False
async with self._lock:
self._prune_expired(time.monotonic())
if msg_id in self._inflight:
return False
now = time.monotonic()
if msg_id in self._cache:
timestamp, entry_ttl = self._cache[msg_id]
if now - timestamp < entry_ttl:
return False
self._cache[msg_id] = (now, ttl)
self._inflight.add(msg_id)
return True
async def claim_with_status(self, msg_id: str) -> ClaimStatus:
if not msg_id:
return ClaimStatus.INVALID
async with self._lock:
self._prune_expired(time.monotonic())
if msg_id in self._inflight:
return ClaimStatus.INFLIGHT
now = time.monotonic()
if msg_id in self._cache:
timestamp, entry_ttl = self._cache[msg_id]
if now - timestamp < entry_ttl:
return ClaimStatus.DUPLICATE
self._cache[msg_id] = (now, _DEFAULT_TTL_SECONDS)
self._inflight.add(msg_id)
return ClaimStatus.CLAIMED
async def commit(self, msg_id: str) -> None:
async with self._lock:
self._inflight.discard(msg_id)
async def release(self, msg_id: str) -> None:
async with self._lock:
self._inflight.discard(msg_id)
async def release_and_forget(self, msg_id: str) -> None:
async with self._lock:
self._inflight.discard(msg_id)
self._cache.pop(msg_id, None)
async def clear_inflight(self, msg_id: str | None = None) -> None:
async with self._lock:
if msg_id is not None:
self._inflight.discard(msg_id)
else:
self._inflight.clear()
async def reset(self) -> None:
async with self._lock:
self._cache.clear()
self._inflight.clear()
class RedisBackend(IdempotencyBackend):
"""基于 Redis 的分布式幂等性后端
使用 SET NX EX 实现原子性检查和设置结合 SADD/SREM 管理 inflight 状态
"""
_INFLIGHT_SUFFIX = ":inflight"
def __init__(self, redis_client, key_prefix: str = "yuxi:idempotency") -> None:
self._redis = redis_client
self._key_prefix = key_prefix
def _cache_key(self, msg_id: str) -> str:
return f"{self._key_prefix}:msg:{msg_id}"
def _inflight_key(self, msg_id: str) -> str:
return f"{self._key_prefix}{self._INFLIGHT_SUFFIX}"
async def check_and_set(self, msg_id: str, ttl: int = _DEFAULT_TTL_SECONDS) -> bool:
if not msg_id:
return False
key = self._cache_key(msg_id)
return await self._redis.set(key, "1", nx=True, ex=ttl) or False
async def claim(self, msg_id: str, ttl: int = _DEFAULT_TTL_SECONDS) -> bool:
if not msg_id:
return False
inflight_key = self._inflight_key(msg_id)
added = await self._redis.sadd(inflight_key, msg_id)
if added == 0:
return False
key = self._cache_key(msg_id)
set_ok = await self._redis.set(key, "1", nx=True, ex=ttl)
if not set_ok:
await self._redis.srem(inflight_key, msg_id)
return False
return True
async def claim_with_status(self, msg_id: str) -> ClaimStatus:
if not msg_id:
return ClaimStatus.INVALID
inflight_key = self._inflight_key(msg_id)
if await self._redis.sismember(inflight_key, msg_id):
return ClaimStatus.INFLIGHT
key = self._cache_key(msg_id)
if await self._redis.exists(key):
return ClaimStatus.DUPLICATE
added = await self._redis.sadd(inflight_key, msg_id)
if added == 0:
return ClaimStatus.INFLIGHT
await self._redis.set(key, "1", ex=_DEFAULT_TTL_SECONDS)
return ClaimStatus.CLAIMED
async def commit(self, msg_id: str) -> None:
inflight_key = self._inflight_key(msg_id)
await self._redis.srem(inflight_key, msg_id)
async def release(self, msg_id: str) -> None:
inflight_key = self._inflight_key(msg_id)
await self._redis.srem(inflight_key, msg_id)
async def release_and_forget(self, msg_id: str) -> None:
inflight_key = self._inflight_key(msg_id)
cache_key = self._cache_key(msg_id)
await self._redis.srem(inflight_key, msg_id)
await self._redis.delete(cache_key)
async def clear_inflight(self, msg_id: str | None = None) -> None:
if msg_id is not None:
inflight_key = self._inflight_key(msg_id)
await self._redis.srem(inflight_key, msg_id)
else:
keys = await self._redis.keys(f"{self._key_prefix}{self._INFLIGHT_SUFFIX}:*")
for key in keys:
await self._redis.delete(key)
async def reset(self) -> None:
keys = await self._redis.keys(f"{self._key_prefix}:*")
for key in keys:
await self._redis.delete(key)
_backend: IdempotencyBackend | None = None
_backend_lock = asyncio.Lock()
def _get_redis_client():
redis_url = os.getenv("REDIS_URL", os.getenv("YUXI_REDIS_URL", ""))
if not redis_url:
return None
try:
from redis.asyncio import Redis
return Redis.from_url(redis_url, decode_responses=True)
except Exception:
logger.warning("Failed to create Redis client for idempotency, falling back to in-memory")
return None
async def get_backend() -> IdempotencyBackend:
global _backend
if _backend is not None:
return _backend
async with _backend_lock:
if _backend is not None:
return _backend
redis_client = _get_redis_client()
if redis_client is not None:
_backend = RedisBackend(redis_client)
logger.info("Idempotency backend: Redis")
else:
_backend = InMemoryBackend()
logger.info("Idempotency backend: InMemory (single-process)")
return _backend
async def _ensure_backend() -> IdempotencyBackend:
if _backend is not None:
return _backend
return await get_backend()
async def check_and_set(msg_id: str, ttl: int = _DEFAULT_TTL_SECONDS) -> bool:
backend = await _ensure_backend()
return await backend.check_and_set(msg_id, ttl)
async def claim(msg_id: str, ttl: int = _DEFAULT_TTL_SECONDS) -> bool:
backend = await _ensure_backend()
return await backend.claim(msg_id, ttl)
async def claim_with_status(msg_id: str) -> ClaimStatus:
backend = await _ensure_backend()
return await backend.claim_with_status(msg_id)
async def commit(msg_id: str) -> None:
backend = await _ensure_backend()
await backend.commit(msg_id)
async def release(msg_id: str) -> None:
backend = await _ensure_backend()
await backend.release(msg_id)
async def release_and_forget(msg_id: str) -> None:
backend = await _ensure_backend()
await backend.release_and_forget(msg_id)
async def consume(msg_id: str) -> bool:
if not await claim(msg_id):
return False
await release(msg_id)
return True
def build_idempotency_key(prefix: str, key: str, nonce: str | None = None) -> str:
base = f"{prefix}{_IDEMPOTENCY_KEY_SEPARATOR}{key}"
if nonce:
return f"{base}{_IDEMPOTENCY_NONCE_MARKER}{nonce}"
return base
def parse_idempotency_key(idempotency_key: str) -> tuple[str, str] | None:
if not idempotency_key:
return None
sep_idx = idempotency_key.find(_IDEMPOTENCY_KEY_SEPARATOR)
if sep_idx < 0:
return None
prefix = idempotency_key[:sep_idx]
body = idempotency_key[sep_idx + 1 :]
nonce_marker = body.rfind(_IDEMPOTENCY_NONCE_MARKER)
if nonce_marker >= 0:
return prefix, body[:nonce_marker]
return prefix, body
async def clear_inflight(msg_id: str | None = None) -> None:
backend = await _ensure_backend()
await backend.clear_inflight(msg_id)
async def reset() -> None:
backend = await _ensure_backend()
await backend.reset()
def reset_sync() -> None:
global _backend
if _backend is not None:
if isinstance(_backend, InMemoryBackend):
_backend._cache.clear()
_backend._inflight.clear()
else:
import asyncio
asyncio.get_event_loop().run_until_complete(reset())
else:
_backend = InMemoryBackend()
reset_for_tests = reset_sync

View File

@ -0,0 +1,202 @@
import logging
import time
from typing import Any
from yuxi.channel.message.models import DispatchResult, UnifiedMessage
from yuxi.services.langfuse_service import (
get_langfuse_client,
is_langfuse_enabled,
)
logger = logging.getLogger(__name__)
class ChannelTrace:
"""渠道消息处理链路的 Langfuse trace 包装器。
复用现有 langfuse_service Langfuse 客户端创建实际 Langfuse
span 以追踪消息处理每个阶段分发校验规则匹配等的耗时和状态
"""
def __init__(self, msg: UnifiedMessage, dispatch_result: DispatchResult | None = None):
self._msg = msg
self._dispatch_result = dispatch_result
self._client = get_langfuse_client()
self._enabled = is_langfuse_enabled() and self._client is not None
self._trace_id: str | None = None
self._trace_obj: Any = None
self._trace_start_time: float | None = None
self._live_spans: dict[str, Any] = {}
@property
def trace_name(self) -> str:
peer_kind = self._msg.sender.kind.value
return f"channel:{self._msg.channel_type}:{peer_kind}"
@property
def trace_id(self) -> str | None:
return self._trace_id
@property
def span_count(self) -> int:
return len(self._live_spans)
async def start(self) -> "ChannelTrace":
if not self._enabled:
return self
try:
self._trace_start_time = time.monotonic()
self._trace_id = self._client.create_trace_id()
self._trace_obj = self._client.trace(
id=self._trace_id,
name=self.trace_name,
metadata=self._build_base_metadata(),
tags=self._build_base_tags(),
)
except Exception:
logger.exception("Failed to create channel trace: %s", self.trace_name)
self._enabled = False
return self
async def add_span(
self,
name: str,
metadata: dict[str, Any] | None = None,
input_data: Any = None,
output_data: Any = None,
level: str = "DEFAULT",
status_message: str | None = None,
parent_span_id: str | None = None,
) -> str | None:
"""创建 Langfuse span 并追踪其生命周期。
返回 span_id 供后续 end_span() 结束使用
调用者必须在完成处理后调用 end_span() 或依赖 finish() 统一结束
"""
if not self._enabled or self._trace_obj is None:
return None
try:
parent = self._live_spans.get(parent_span_id) if parent_span_id else self._trace_obj
span = parent.span(
name=name,
input=input_data,
output=output_data,
metadata=metadata or {},
level=level,
status_message=status_message,
)
span_id: str = span.id
self._live_spans[span_id] = span
return span_id
except Exception:
logger.exception("Failed to create span '%s' for trace: %s", name, self.trace_name)
return None
async def end_span(
self,
span_id: str,
output_data: Any = None,
status_message: str | None = None,
level: str | None = None,
metadata: dict[str, Any] | None = None,
) -> None:
"""结束指定 span记录输出与最终状态。"""
span = self._live_spans.pop(span_id, None)
if span is None:
return
try:
update_kwargs: dict[str, Any] = {}
if output_data is not None:
update_kwargs["output"] = output_data
if level is not None:
update_kwargs["level"] = level
if status_message is not None:
update_kwargs["status_message"] = status_message
if metadata is not None:
update_kwargs["metadata"] = metadata
if update_kwargs:
span.update(**update_kwargs)
span.end()
except Exception:
logger.exception("Failed to end span '%s' for trace: %s", span_id, self.trace_name)
async def finish(self, error: str | None = None) -> None:
if not self._enabled or self._trace_obj is None:
return
try:
pending = dict(self._live_spans)
self._live_spans.clear()
for span_id, span in pending.items():
try:
span.end()
except Exception:
logger.exception("Failed to end pending span '%s'", span_id)
duration_ms = None
if self._trace_start_time is not None:
duration_ms = (time.monotonic() - self._trace_start_time) * 1000
metadata: dict[str, Any] = {
**self._build_base_metadata(),
"span_count": len(pending),
"dispatch_success": error is None,
**(self._build_dispatch_metadata() if self._dispatch_result else {}),
}
if duration_ms is not None:
metadata["duration_ms"] = round(duration_ms, 2)
self._trace_obj.update(
output=error or "success",
metadata=metadata,
)
except Exception:
logger.exception("Failed to finish channel trace: %s", self.trace_name)
def _build_base_metadata(self) -> dict[str, Any]:
return {
"source": "channel",
"channel_type": self._msg.channel_type,
"account_id": self._msg.account_id,
"peer_kind": self._msg.sender.kind.value,
"peer_id": self._msg.sender.id,
"msg_id": self._msg.msg_id,
"message_type": self._msg.message_type.value,
"feature": "channel_dispatch",
}
def _build_dispatch_metadata(self) -> dict[str, Any]:
if self._dispatch_result is None:
return {}
return {
"agent_config_id": str(self._dispatch_result.agent_config_id)
if self._dispatch_result.agent_config_id
else None,
"session_key": getattr(self._dispatch_result, "session_key", None),
"thread_id": self._dispatch_result.thread_id,
"matched_by": getattr(self._dispatch_result, "matched_by", None),
}
def _build_base_tags(self) -> list[str]:
tags = [
"yuxi",
"channel",
f"channel:{self._msg.channel_type}",
f"peer_kind:{self._msg.sender.kind.value}",
]
if self._dispatch_result and self._dispatch_result.agent_config_id:
tags.append(f"agent_config:{self._dispatch_result.agent_config_id}")
return tags
async def create_channel_trace(
msg: UnifiedMessage,
dispatch_result: DispatchResult | None = None,
) -> ChannelTrace:
trace = ChannelTrace(msg, dispatch_result)
await trace.start()
return trace

View File

@ -0,0 +1,427 @@
from __future__ import annotations
import asyncio
import base64
import io
import logging
from dataclasses import dataclass
from email.message import EmailMessage
from pathlib import Path, PurePosixPath
from urllib.parse import urlparse
import httpx
from PIL import Image
from yuxi.channel.message.media_cleaner import MediaCleaner
from yuxi.channel.message.media_store import MediaStore
from yuxi.config import config
logger = logging.getLogger(__name__)
_media_store: MediaStore | None = None
_media_store_initialized: bool = False
_media_cleaner: MediaCleaner | None = None
_background_tasks: set[asyncio.Task[object]] = set()
MAX_IMAGE_BYTES = 6 * 1024 * 1024 # 6MB
_MAX_STORE_CONCURRENCY = 5
_store_semaphore = asyncio.Semaphore(_MAX_STORE_CONCURRENCY)
_SVG_MIME_TYPES = frozenset({"image/svg+xml", "image/svg"})
_IMAGE_MAGIC_BYTES: dict[str, tuple[bytes, ...]] = {
"image/jpeg": (b"\xff\xd8\xff",),
"image/png": (b"\x89PNG\r\n\x1a\n",),
"image/gif": (b"GIF87a", b"GIF89a"),
"image/webp": (b"RIFF",),
"image/bmp": (b"BM",),
"image/tiff": (b"II*\x00", b"MM\x00*"),
"image/x-icon": (b"\x00\x00\x01\x00",),
}
_SVG_SCRIPT_MARKER = b"<script"
_media_client: httpx.AsyncClient | None = None
def _get_media_client() -> httpx.AsyncClient:
global _media_client
if _media_client is None:
_media_client = httpx.AsyncClient(limits=httpx.Limits(max_connections=10, max_keepalive_connections=5))
return _media_client
async def _cleanup_media_client() -> None:
global _media_client
if _media_client is not None:
await _media_client.aclose()
_media_client = None
_MAX_RETRIES = 2
_RETRY_DELAY = 1.0
_RETRYABLE_STATUSES = frozenset({408, 429, 500, 502, 503, 504})
@dataclass
class MediaDownloadResult:
data: bytes
content_type: str | None
filename: str | None
def _parse_content_disposition_filename(header: str | None) -> str | None:
if not header:
return None
msg = EmailMessage()
msg["Content-Disposition"] = header
filename = msg.get_filename()
if filename:
return filename
return None
def _is_transient_error(exc: Exception) -> bool:
if isinstance(exc, httpx.HTTPStatusError):
return exc.response.status_code in _RETRYABLE_STATUSES
if isinstance(exc, httpx.TimeoutException):
return True
if isinstance(exc, httpx.NetworkError):
return True
return False
def _resolve_filename(final_url: str, response: httpx.Response, file_path_hint: str | None) -> str | None:
header_filename = _parse_content_disposition_filename(response.headers.get("content-disposition"))
if header_filename:
return header_filename
if file_path_hint:
return PurePosixPath(file_path_hint).name or None
try:
path = urlparse(final_url).path
name = PurePosixPath(path).name
return name or None
except Exception:
return None
async def _read_with_limit(response: httpx.Response, max_bytes: int) -> bytes:
chunks: list[bytes] = []
total = 0
async for chunk in response.aiter_bytes(chunk_size=65536):
total += len(chunk)
if total > max_bytes:
raise ValueError(f"Payload exceeds maxBytes {max_bytes}")
chunks.append(chunk)
return b"".join(chunks)
def _validate_image_magic_bytes(data: bytes, content_type: str | None) -> bool:
ct = (content_type or "").split(";")[0].strip().lower()
if ct in _SVG_MIME_TYPES:
if data.lower().find(_SVG_SCRIPT_MARKER) != -1:
logger.warning("SVG contains script tag, rejecting for XSS safety")
return False
return True
expected_magics = _IMAGE_MAGIC_BYTES.get(ct)
if expected_magics is None:
return True
return any(data.startswith(magic) for magic in expected_magics)
async def _fetch_media(
url: str,
*,
max_bytes: int,
timeout: float,
file_path_hint: str | None = None,
) -> MediaDownloadResult | None:
parsed = urlparse(url)
if parsed.scheme not in ("http", "https"):
logger.warning("Unsupported media URL scheme: %s", parsed.scheme)
return None
for attempt in range(_MAX_RETRIES + 1):
try:
client = _get_media_client()
response = await client.get(url, follow_redirects=True, timeout=timeout)
content_length_raw = response.headers.get("content-length")
if content_length_raw:
try:
content_length = int(content_length_raw)
if content_length > max_bytes:
logger.error(
"Content-Length %d exceeds maxBytes %d for url=%s",
content_length,
max_bytes,
url,
)
return None
except ValueError:
pass
response.raise_for_status()
content_type = response.headers.get("content-type")
data = await _read_with_limit(response, max_bytes)
if not _validate_image_magic_bytes(data, content_type):
logger.warning(
"Media magic bytes mismatch: content-type=%s, url=%s",
content_type,
url,
)
return None
filename = _resolve_filename(str(response.url), response, file_path_hint)
return MediaDownloadResult(
data=data,
content_type=content_type,
filename=filename,
)
except (httpx.HTTPStatusError, httpx.TimeoutException, httpx.NetworkError) as e:
if attempt < _MAX_RETRIES and _is_transient_error(e):
delay = _RETRY_DELAY * (2**attempt)
logger.warning(
"Transient error fetching media (attempt %d/%d), retrying in %.1fs: url=%s, err=%s",
attempt + 1,
_MAX_RETRIES + 1,
delay,
url,
e,
)
await asyncio.sleep(delay)
continue
if isinstance(e, httpx.HTTPStatusError):
logger.error(
"HTTP error downloading media: url=%s, status=%d",
url,
e.response.status_code,
)
elif isinstance(e, httpx.TimeoutException):
logger.error("Timeout downloading media: url=%s", url)
else:
logger.error("Network error downloading media: url=%s, err=%s", url, e)
return None
except Exception:
logger.exception("Failed to download media: url=%s", url)
return None
logger.error("All retry attempts exhausted for media: url=%s", url)
return None
async def download_image_to_base64(
url: str,
timeout: float = 30.0,
) -> str | None:
result = await _fetch_media(url, max_bytes=MAX_IMAGE_BYTES, timeout=timeout)
if result is None:
return None
if not (result.content_type or "").startswith("image/"):
logger.warning(
"Non-image response: content-type=%s, url=%s",
result.content_type,
url,
)
return None
_maybe_store(result.data, result.content_type, result.filename, url)
return base64.b64encode(result.data).decode("utf-8")
async def resolve_image_base64(media_urls: list[str]) -> str | None:
if not media_urls:
return None
for url in media_urls:
result = await download_image_to_base64(url)
if result:
return result
return None
def get_media_store() -> MediaStore | None:
global _media_store, _media_store_initialized
if _media_store_initialized:
return _media_store
_media_store_initialized = True
if not config.media_store_enabled:
return None
media_dir = Path(config.save_dir) / "data" / "media"
_media_store = MediaStore(
base_dir=media_dir,
max_total_bytes=config.media_store_max_total_bytes,
max_file_bytes=config.media_store_max_file_bytes,
ttl_seconds=config.media_store_ttl_seconds,
)
return _media_store
def get_media_cleaner() -> MediaCleaner | None:
global _media_cleaner
store = get_media_store()
if store is None:
return None
if _media_cleaner is None:
_media_cleaner = MediaCleaner(store, config.media_cleanup_interval_seconds)
return _media_cleaner
async def start_media_cleaner() -> None:
cleaner = get_media_cleaner()
if cleaner is not None:
await cleaner.start()
async def stop_media_cleaner() -> None:
cleaner = get_media_cleaner()
if cleaner is not None:
await cleaner.stop()
await _cleanup_media_client()
async def shutdown_media() -> None:
await stop_media_cleaner()
await _cancel_background_tasks()
await _cleanup_media_client()
async def _cancel_background_tasks() -> None:
_cleanup_completed_tasks()
pending = list(_background_tasks)
if not pending:
return
for t in pending:
t.cancel()
results = await asyncio.gather(*pending, return_exceptions=True)
for r in results:
if isinstance(r, Exception) and not isinstance(r, asyncio.CancelledError):
logger.warning("Background store task error during shutdown: %s", r)
_background_tasks.clear()
def _cleanup_completed_tasks() -> None:
done = {t for t in _background_tasks if t.done()}
_background_tasks.difference_update(done)
if done:
logger.debug("Cleaned up %d completed background store tasks, %d remaining", len(done), len(_background_tasks))
async def _store_with_semaphore(store: MediaStore, data: bytes, content_type: str | None, filename: str | None, source_url: str) -> None:
async with _store_semaphore:
await store.store(data, content_type, filename, source_url)
def _maybe_store(data: bytes, content_type: str | None, filename: str | None, source_url: str) -> None:
store = get_media_store()
if store is None:
return
_cleanup_completed_tasks()
if len(_background_tasks) >= _MAX_STORE_CONCURRENCY * 2:
logger.warning(
"Background store tasks piling up: %d pending, concurrency limit=%d",
len(_background_tasks),
_MAX_STORE_CONCURRENCY,
)
task = asyncio.create_task(_store_with_semaphore(store, data, content_type, filename, source_url))
task.add_done_callback(_handle_store_task_done)
_background_tasks.add(task)
def _handle_store_task_done(task: asyncio.Task[object]) -> None:
try:
task.result()
except asyncio.CancelledError:
pass
except Exception:
logger.exception("Failed to store media in background task")
finally:
_background_tasks.discard(task)
@dataclass
class ImageProcessResult:
original_file_id: str | None
thumbnail_file_id: str | None
width: int
height: int
format: str
mime_type: str
size_bytes: int
async def download_and_process_image(
url: str,
*,
timeout: float = 30.0,
thumbnail_size: tuple[int, int] | None = None,
convert_to: str | None = None,
) -> ImageProcessResult | None:
result = await _fetch_media(url, max_bytes=MAX_IMAGE_BYTES, timeout=timeout)
if result is None:
return None
if not (result.content_type or "").startswith("image/"):
logger.warning(
"Non-image response for download_and_process_image: content-type=%s, url=%s",
result.content_type,
url,
)
return None
from yuxi.utils.image_processor import image_processor as img_proc
image_data = result.data
if convert_to and convert_to in img_proc.CONVERTIBLE_FORMATS:
try:
image_data = img_proc.convert_format(result.data, convert_to)
except Exception:
logger.exception("Failed to convert image format to %s for url=%s", convert_to, url)
store = get_media_store()
original_file_id: str | None = None
thumbnail_file_id: str | None = None
if store is not None:
original_file_id = await store.store(
image_data,
content_type=result.content_type,
filename=result.filename,
source_url=url,
)
try:
thumbnail_data = img_proc.generate_thumbnail(image_data, size=thumbnail_size)
except Exception:
logger.exception("Failed to generate thumbnail for url=%s", url)
thumbnail_data = None
if thumbnail_data and store is not None:
thumbnail_file_id = await store.store(
thumbnail_data,
content_type="image/jpeg",
filename=None,
source_url=url,
)
try:
with io.BytesIO(image_data) as buf:
with Image.open(buf) as img:
width, height = img.size
fmt = img.format or "JPEG"
except Exception:
width, height = 0, 0
fmt = "JPEG"
return ImageProcessResult(
original_file_id=original_file_id,
thumbnail_file_id=thumbnail_file_id,
width=width,
height=height,
format=fmt,
mime_type=f"image/{fmt.lower()}",
size_bytes=len(image_data),
)

View File

@ -0,0 +1,63 @@
from __future__ import annotations
import asyncio
import logging
from yuxi.channel.message.media_store import MediaStore
logger = logging.getLogger(__name__)
class MediaCleaner:
def __init__(self, media_store: MediaStore, interval_seconds: int):
self._media_store = media_store
self._interval = interval_seconds
self._task: asyncio.Task | None = None
self._stop_event = asyncio.Event()
@property
def is_running(self) -> bool:
return self._task is not None and not self._task.done()
async def start(self) -> None:
if self.is_running:
return
self._stop_event.clear()
self._task = asyncio.create_task(self._loop(), name="media-cleaner")
logger.info(
"MediaCleaner started (interval=%ds, ttl=%ds)",
self._interval,
self._media_store.ttl_seconds,
)
async def stop(self) -> None:
if not self.is_running:
return
self._stop_event.set()
try:
await asyncio.wait_for(self._task, timeout=10)
except TimeoutError:
logger.warning("MediaCleaner stop timed out, cancelling task")
self._task.cancel()
except asyncio.CancelledError:
pass
self._task = None
logger.info("MediaCleaner stopped")
async def cleanup_once(self) -> int:
return await self._media_store.cleanup_expired()
async def _loop(self) -> None:
while not self._stop_event.is_set():
try:
await asyncio.wait_for(self._stop_event.wait(), timeout=self._interval)
break
except TimeoutError:
pass
try:
deleted = await self._media_store.cleanup_expired()
if deleted:
logger.info("MediaCleaner cycle: deleted %d expired files", deleted)
except Exception:
logger.exception("MediaCleaner cycle error")

View File

@ -0,0 +1,242 @@
from __future__ import annotations
import asyncio
import json
import logging
import os
import uuid
from datetime import UTC, datetime
from pathlib import Path
logger = logging.getLogger(__name__)
_META_SUFFIX = ".meta.json"
def _now_iso() -> str:
return datetime.now(UTC).isoformat()
def _date_dir_name() -> str:
return datetime.now(UTC).strftime("%Y-%m-%d")
class MediaStore:
def __init__(
self,
base_dir: Path,
max_total_bytes: int,
max_file_bytes: int,
ttl_seconds: int,
):
self.base_dir = base_dir
self.max_total_bytes = max_total_bytes
self.max_file_bytes = max_file_bytes
self.ttl_seconds = ttl_seconds
self._size_initialized = False
self._running_total: int = 0
self._size_lock = asyncio.Lock()
def _ensure_dir(self, dir_path: Path) -> None:
dir_path.mkdir(parents=True, exist_ok=True)
def _make_file_id(self, date_dir: str, uid: str, ext: str) -> str:
return f"{date_dir}/{uid}{ext}"
def _get_file_path(self, file_id: str) -> Path:
return self.base_dir / file_id
def _get_meta_path(self, file_id: str) -> Path:
return self.base_dir / f"{file_id}{_META_SUFFIX}"
def _dir_size(self, dir_path: Path) -> int:
total = 0
try:
for entry in dir_path.rglob("*"):
if entry.is_file():
try:
total += entry.stat().st_size
except OSError:
logger.debug("Failed to stat file during size calculation: %s", entry)
except OSError:
logger.debug("Failed to calculate directory size: %s", dir_path)
return total
async def _check_quota(self, incoming_bytes: int) -> bool:
if incoming_bytes > self.max_file_bytes:
logger.warning(
"File size %d exceeds per-file limit %d",
incoming_bytes,
self.max_file_bytes,
)
return False
current_total = await self._get_running_total()
if current_total + incoming_bytes > self.max_total_bytes:
logger.warning(
"Storage quota exceeded: current=%d, incoming=%d, max=%d",
current_total,
incoming_bytes,
self.max_total_bytes,
)
return False
return True
async def _get_running_total(self) -> int:
async with self._size_lock:
if not self._size_initialized:
self._running_total = self._dir_size(self.base_dir)
self._size_initialized = True
return self._running_total
def _infer_ext(self, content_type: str | None, filename: str | None) -> str:
if filename:
suffix = Path(filename).suffix
if suffix:
return suffix.lower()
if content_type:
ct = content_type.split(";")[0].strip().lower()
mapping = {
"image/jpeg": ".jpg",
"image/png": ".png",
"image/gif": ".gif",
"image/webp": ".webp",
"image/bmp": ".bmp",
"image/svg+xml": ".svg",
"audio/mpeg": ".mp3",
"audio/wav": ".wav",
"audio/ogg": ".ogg",
"video/mp4": ".mp4",
"application/pdf": ".pdf",
}
if ct in mapping:
return mapping[ct]
return ".bin"
async def store(
self,
data: bytes,
content_type: str | None = None,
filename: str | None = None,
source_url: str | None = None,
) -> str | None:
if not await self._check_quota(len(data)):
return None
date_dir = _date_dir_name()
uid = uuid.uuid4().hex
ext = self._infer_ext(content_type, filename)
file_id = self._make_file_id(date_dir, uid, ext)
file_path = self._get_file_path(file_id)
meta_path = self._get_meta_path(file_id)
self._ensure_dir(file_path.parent)
try:
file_path.write_bytes(data)
except OSError:
logger.exception("Failed to write media file: %s", file_path)
return None
meta = {
"original_filename": filename,
"content_type": content_type,
"source_url": source_url,
"stored_at": _now_iso(),
"ttl_seconds": self.ttl_seconds,
"size": len(data),
}
try:
meta_path.write_text(json.dumps(meta, ensure_ascii=False), encoding="utf-8")
except OSError:
logger.exception("Failed to write media metadata: %s", meta_path)
try:
file_path.unlink(missing_ok=True)
except OSError:
logger.debug("Failed to clean up orphaned media file: %s", file_path)
return None
logger.info("Stored media: file_id=%s, size=%d", file_id, len(data))
async with self._size_lock:
self._running_total += len(data)
return file_id
def resolve_path(self, file_id: str) -> Path | None:
file_path = self._get_file_path(file_id)
if file_path.exists():
return file_path
return None
def read_meta(self, file_id: str) -> dict | None:
meta_path = self._get_meta_path(file_id)
try:
return json.loads(meta_path.read_text(encoding="utf-8"))
except (OSError, json.JSONDecodeError):
return None
async def delete(self, file_id: str) -> bool:
file_path = self._get_file_path(file_id)
meta_path = self._get_meta_path(file_id)
deleted = False
for p in (file_path, meta_path):
if not p.exists():
continue
try:
p.unlink()
deleted = True
except OSError:
logger.debug("Failed to delete file during media cleanup: %s", p)
if deleted:
logger.info("Deleted media: file_id=%s", file_id)
return deleted
def get_total_size(self) -> int:
return self._dir_size(self.base_dir)
async def cleanup_expired(self) -> int:
deleted_count = 0
now = datetime.now(UTC)
try:
for entry in sorted(self.base_dir.rglob(f"*{_META_SUFFIX}")):
try:
meta = json.loads(entry.read_text(encoding="utf-8"))
except (OSError, json.JSONDecodeError):
continue
stored_at_str = meta.get("stored_at")
ttl_s = meta.get("ttl_seconds", self.ttl_seconds)
if not stored_at_str:
continue
try:
stored_at = datetime.fromisoformat(stored_at_str)
except ValueError:
continue
if (now - stored_at).total_seconds() <= ttl_s:
continue
file_id = str(entry.relative_to(self.base_dir)).removesuffix(_META_SUFFIX)
if await self.delete(file_id):
deleted_count += 1
self._remove_empty_dirs()
except OSError:
logger.exception("Error during media cleanup scan")
if deleted_count:
logger.info("Media cleanup: deleted %d expired files", deleted_count)
return deleted_count
def _remove_empty_dirs(self) -> None:
try:
for dirpath, dirnames, filenames in os.walk(self.base_dir, topdown=False):
if dirpath == str(self.base_dir):
continue
if not dirnames and not filenames:
try:
os.rmdir(dirpath)
except OSError:
logger.debug("Failed to remove empty directory: %s", dirpath)
except OSError:
logger.debug("Failed to walk directory for cleanup: %s", self.base_dir)

View File

@ -0,0 +1,211 @@
import time
from collections import defaultdict
from collections.abc import Callable
from dataclasses import dataclass, field
from enum import StrEnum
class MetricType(StrEnum):
COUNTER = "counter"
GAUGE = "gauge"
HISTOGRAM = "histogram"
@dataclass(slots=True)
class _Counter:
value: int = 0
def inc(self, amount: int = 1) -> None:
self.value += amount
@dataclass(slots=True)
class _Gauge:
value: float = 0.0
def set(self, value: float) -> None:
self.value = value
def inc(self, amount: float = 1.0) -> None:
self.value += amount
def dec(self, amount: float = 1.0) -> None:
self.value -= amount
@dataclass(slots=True)
class _Histogram:
buckets: list[float]
values: list[int] = field(default_factory=list)
_sum: float = 0.0
_count: int = 0
def observe(self, value: float) -> None:
self._sum += value
self._count += 1
while len(self.values) < len(self.buckets):
self.values.append(0)
for i, bound in enumerate(self.buckets):
if value <= bound:
self.values[i] += 1
return
def quantile(self, q: float) -> float | None:
if self._count == 0:
return None
if not self.values:
return None
target_rank = q * self._count
cumulative = 0
for i, (bound, count) in enumerate(zip(self.buckets, self.values)):
cumulative += count
if cumulative >= target_rank:
if i == 0:
lower_bound = 0.0
else:
lower_bound = self.buckets[i - 1]
upper_bound = bound
prev_cumulative = cumulative - count
fraction = (target_rank - prev_cumulative) / max(count, 1)
return lower_bound + fraction * (upper_bound - lower_bound)
return float(self.buckets[-1]) if self.buckets else None
class MetricsRegistry:
def __init__(self) -> None:
self._label_values: dict[str, dict[tuple[str, ...], object]] = defaultdict(dict)
def counter(self, name: str, label_keys: tuple[str, ...] = ()) -> Callable[..., None]:
def inc(labels: dict[str, str] | None = None, amount: int = 1) -> None:
key = self._resolve_key(labels, label_keys)
counter = self._label_values[name].get(key)
if counter is None:
counter = _Counter()
self._label_values[name][key] = counter
counter.inc(amount)
return inc
def gauge(self, name: str, label_keys: tuple[str, ...] = ()) -> Callable[..., None]:
def set(value: float, labels: dict[str, str] | None = None) -> None:
key = self._resolve_key(labels, label_keys)
gauge = self._label_values[name].get(key)
if gauge is None:
gauge = _Gauge()
self._label_values[name][key] = gauge
gauge.set(value)
return set
def histogram(self, name: str, buckets: list[float], label_keys: tuple[str, ...] = ()) -> Callable[..., None]:
def observe(value: float, labels: dict[str, str] | None = None) -> None:
key = self._resolve_key(labels, label_keys)
hist = self._label_values[name].get(key)
if hist is None:
hist = _Histogram(buckets=buckets)
self._label_values[name][key] = hist
hist.observe(value)
return observe
@staticmethod
def _resolve_key(labels: dict[str, str] | None, label_keys: tuple[str, ...]) -> tuple[str, ...]:
if not labels:
return tuple("" for _ in label_keys)
return tuple(labels.get(k, "") for k in label_keys)
def snapshot(self) -> dict:
result: dict = {}
for metric_name, entries in self._label_values.items():
result[metric_name] = {}
for label_tuple, metric in entries.items():
if isinstance(metric, _Counter):
result[metric_name][str(label_tuple)] = {"type": "counter", "value": metric.value}
elif isinstance(metric, _Gauge):
result[metric_name][str(label_tuple)] = {"type": "gauge", "value": metric.value}
elif isinstance(metric, _Histogram):
result[metric_name][str(label_tuple)] = {
"type": "histogram",
"buckets": metric.buckets,
"values": metric.values,
"sum": metric._sum,
"count": metric._count,
"p50": metric.quantile(0.50),
"p95": metric.quantile(0.95),
"p99": metric.quantile(0.99),
}
return result
registry = MetricsRegistry()
channel_messages_total = registry.counter(
"channel_messages_total",
("channel_type", "status"),
)
channel_dispatch_duration_ms = registry.histogram(
"channel_dispatch_duration_ms",
buckets=[5, 25, 50, 100, 250, 500, 1000, 2500, 5000],
label_keys=("channel_type",),
)
channel_agent_duration_ms = registry.histogram(
"channel_agent_duration_ms",
buckets=[100, 500, 1000, 2500, 5000, 10000, 30000, 60000],
label_keys=("channel_type",),
)
channel_rate_limit_rejects_total = registry.counter(
"channel_rate_limit_rejects_total",
("channel_type",),
)
channel_messages_inflight = registry.gauge(
"channel_messages_inflight",
)
def record_message(channel_type: str, status: str) -> None:
channel_messages_total(labels={"channel_type": channel_type, "status": status})
def record_dispatch_duration_ms(channel_type: str, duration_ms: float) -> None:
channel_dispatch_duration_ms(duration_ms, labels={"channel_type": channel_type})
def record_agent_duration_ms(channel_type: str, duration_ms: float) -> None:
channel_agent_duration_ms(duration_ms, labels={"channel_type": channel_type})
def record_rate_limit_reject(channel_type: str) -> None:
channel_rate_limit_rejects_total(labels={"channel_type": channel_type})
def set_inflight(count: int) -> None:
channel_messages_inflight(count)
class MetricsTimer:
def __init__(self, on_finish: Callable[[float], None]) -> None:
self._on_finish = on_finish
self._start = 0.0
def __enter__(self) -> "MetricsTimer":
self._start = time.monotonic()
return self
def __exit__(self, *args) -> None:
elapsed = (time.monotonic() - self._start) * 1000
self._on_finish(elapsed)
async def __aenter__(self) -> "MetricsTimer":
self._start = time.monotonic()
return self
async def __aexit__(self, *args) -> None:
elapsed = (time.monotonic() - self._start) * 1000
self._on_finish(elapsed)

View File

@ -0,0 +1,283 @@
import hashlib
import json
from dataclasses import dataclass, field
from datetime import datetime
from enum import StrEnum
from typing import Any
from yuxi.channel.routing.models import PeerKind
class MessageType(StrEnum):
TEXT = "text"
IMAGE = "image"
VOICE = "voice"
FILE = "file"
EVENT = "event"
class ReplyStage(StrEnum):
TOOL = "tool"
BLOCK = "block"
FINAL = "final"
class ChunkType(StrEnum):
TEXT_DELTA = "text-delta"
REASONING = "reasoning"
TOOL_CALL = "tool-call"
TOOL_RESULT = "tool-result"
HEARTBEAT = "heartbeat"
CITATION = "citation"
STATUS = "status"
ERROR = "error"
class MentionSource(StrEnum):
EXPLICIT_BOT = "explicit_bot"
SUBTEAM = "subteam"
MENTION_PATTERN = "mention_pattern"
IMPLICIT_THREAD = "implicit_thread"
COMMAND_BYPASS = "command_bypass"
NONE = "none"
@dataclass
class PeerInfo:
kind: PeerKind
id: str
display_name: str | None = None
username: str | None = None
tag: str | None = None
is_bot: bool = False
is_self: bool = False
display_label: str | None = None
roles: list[str] = field(default_factory=list)
@dataclass
class GroupContext:
id: str | None = None
name: str | None = None
guild_id: str | None = None
team_id: str | None = None
thread_id: str | None = None
roles: list[str] = field(default_factory=list)
kind: str | None = None
space_id: str | None = None
parent_id: str | None = None
native_channel_id: str | None = None
route_peer_kind: str | None = None
route_peer_id: str | None = None
@dataclass
class UnifiedMessage:
msg_id: str
channel_type: str
account_id: str
content: str
sender: PeerInfo
channel_config_id: str | None = None
message_type: MessageType = MessageType.TEXT
media_urls: list[str] = field(default_factory=list)
image_base64: str | None = None
group: GroupContext | None = None
timestamp: datetime | None = None
raw_payload: dict[str, Any] = field(default_factory=dict)
metadata: dict[str, Any] = field(default_factory=dict)
body_for_agent: str | None = None
command_body: str | None = None
body_for_commands: str | None = None
inbound_history: list[dict[str, Any]] = field(default_factory=list)
reply_chain: list[dict[str, Any]] = field(default_factory=list)
reply_to_id: str | None = None
reply_to_id_full: str | None = None
root_message_id: str | None = None
message_thread_id: str | None = None
reply_to_tag: bool = False
reply_to_current: bool = False
forwarded_from: str | None = None
forwarded_from_type: str | None = None
forwarded_from_id: str | None = None
forwarded_from_username: str | None = None
forwarded_date: float | None = None
media_path: str | None = None
media_paths: list[str] = field(default_factory=list)
media_types: list[str] = field(default_factory=list)
conversation_label: str | None = None
group_channel: str | None = None
group_space: str | None = None
group_members: str | None = None
group_system_prompt: str | None = None
member_role_ids: list[str] = field(default_factory=list)
was_mentioned: bool = False
explicitly_mentioned_bot: bool = False
mentioned_user_ids: list[str] = field(default_factory=list)
mention_source: MentionSource | None = None
surface: str | None = None
originating_channel: str | None = None
originating_to: str | None = None
native_channel_id: str | None = None
native_direct_user_id: str | None = None
thread_parent_id: str | None = None
location_lat: float | None = None
location_lon: float | None = None
@dataclass
class DispatchResult:
success: bool
thread_id: str | None = None
agent_config_id: int | None = None
error: str | None = None
session_key: str | None = None
matched_by: str | None = None
handled_by: str | None = None
command_response: str | None = None
internal_user_id: str | None = None
reply_sent: bool = False
explicit_target: str | None = None
queued_final: bool = False
counts: dict[str, int] = field(default_factory=dict)
failed_counts: dict[str, int] = field(default_factory=dict)
source_reply_delivery_mode: str | None = None
before_agent_run_blocked: bool = False
@dataclass
class StreamingChunk:
"""Agent 推理流式块"""
stage: ReplyStage = ReplyStage.BLOCK
chunk_type: ChunkType = ChunkType.TEXT_DELTA
content: str | None = None
tool_name: str | None = None
tool_input: dict[str, Any] | None = None
tool_output: str | None = None
metadata: dict[str, Any] = field(default_factory=dict)
agent_run_id: str | None = None
is_final: bool = False
item_id: str | None = None
title: str | None = None
phase: str | None = None
steps: list[str] | None = None
@dataclass
class ReplyPayload:
"""出站回复载荷(经过去重和合并后)"""
target_id: str
content: str
reply_to_id: str | None = None
thread_id: str | None = None
media_urls: list[str] = field(default_factory=list)
presentation: dict[str, Any] | None = None
metadata: dict[str, Any] = field(default_factory=dict)
btw: dict[str, str] | None = None
audio_as_voice: bool = False
spoken_text: str | None = None
is_error: bool = False
is_reasoning: bool = False
is_compaction_notice: bool = False
channel_data: dict[str, Any] | None = None
trusted_local_media: bool = False
sensitive_media: bool = False
delivery: dict[str, Any] | None = None
reply_to_tag: bool = False
reply_to_current: bool = False
@property
def payload_key(self) -> str:
key = f"{self.content}|{','.join(sorted(self.media_urls))}|{self.reply_to_id or ''}"
if self.presentation:
stable_hash = hashlib.md5(
json.dumps(self.presentation, sort_keys=True, ensure_ascii=False).encode()
).hexdigest()
key += f"|{stable_hash}"
return key
@property
def content_key(self) -> str:
return f"{self.content}|{','.join(sorted(self.media_urls))}"
class StreamStatus(StrEnum):
"""流执行终止状态"""
COMPLETED = "completed"
CANCELLED = "cancelled"
ERROR = "error"
TIMEOUT = "timeout"
@dataclass
class StreamResult:
"""流执行完成后的结构化结果,对应 openclaw 的 MessageReceipt 概念"""
status: StreamStatus = StreamStatus.COMPLETED
accumulated_text: str = ""
chunk_count: int = 0
text_chunk_count: int = 0
tool_chunk_count: int = 0
reasoning_chunk_count: int = 0
heartbeat_count: int = 0
citation_count: int = 0
error_count: int = 0
start_time: float = 0.0
end_time: float = 0.0
token_usage: dict[str, int] | None = None
error_message: str | None = None
@property
def elapsed_ms(self) -> float:
return (self.end_time - self.start_time) * 1000 if self.start_time > 0 else 0.0
@property
def is_success(self) -> bool:
return self.status == StreamStatus.COMPLETED
class StreamCancellationToken:
"""流取消令牌,对应 openclaw 的 AbortSignal"""
def __init__(self) -> None:
self._cancelled = False
def cancel(self) -> None:
self._cancelled = True
@property
def is_cancelled(self) -> bool:
return self._cancelled
def throw_if_cancelled(self) -> None:
if self._cancelled:
raise StreamCancelledError()
class StreamCancelledError(Exception):
pass
class StreamTimeoutError(Exception):
pass
@dataclass
class MessageReceipt:
platform_message_id: str
thread_id: str | None = None
reply_to_id: str | None = None
sent_at: float | None = None
parts: list[dict[str, Any]] = field(default_factory=list)
raw: list[dict[str, Any]] = field(default_factory=list)

View File

@ -0,0 +1,685 @@
import asyncio
import logging
import time
from collections.abc import AsyncIterator, Callable
from sqlalchemy import select
from yuxi.channel.hooks.lifecycle import HookEvent, LifecycleHookRegistry
from yuxi.channel.message.block_reply_pipeline import create_block_reply_pipeline
from yuxi.channel.message.bridge import AgentBridge
from yuxi.channel.message.conversation_fence import ConversationFence
from yuxi.channel.message.dispatch import MessageDispatcher
from yuxi.channel.message.idempotency import claim, release
from yuxi.channel.message.langfuse_trace import ChannelTrace, create_channel_trace
from yuxi.channel.message.metrics import (
record_dispatch_duration_ms,
record_message,
record_rate_limit_reject,
set_inflight,
)
from yuxi.channel.message.models import (
PeerKind,
StreamCancelledError,
StreamStatus,
StreamTimeoutError,
UnifiedMessage,
)
from yuxi.channel.message.rate_limiter import rate_limit_manager
from yuxi.channel.protocols import (
InboundMessageHook,
MessageSendingHook,
OutboundProtocol,
SendLifecycleHook,
)
from yuxi.channel.plugins.registry import ChannelPluginRegistry
from yuxi.channel.routing.matcher import RouteMatcher
from yuxi.channel.routing.models import RouteBinding
from yuxi.channel.routing.session_key import SessionKeyBuilder
from yuxi.channel.security.allowlist import AllowlistChecker
from yuxi.channel.security.identity_link import IdentityLinkResolver
from yuxi.channel.security.pairing import InMemoryPairingStore, PairingManager
from yuxi.channel.message.circuit_breaker import CircuitBreaker
from yuxi.repositories.channel_binding_repo import ChannelBindingRepository, to_route_binding
from yuxi.repositories.channel_msg_record_repo import ChannelMsgRecordRepository
from yuxi.repositories.channel_thread_mapping_repo import ChannelThreadMappingRepository
from yuxi.repositories.channel_user_mapping_repo import ChannelUserMappingRepository
from yuxi.storage.postgres.manager import pg_manager
from yuxi.storage.postgres.models_business import User
logger = logging.getLogger(__name__)
_BINDINGS_CACHE_TTL = 60.0
_IDEMPOTENCY_TTL = 300
_FALLBACK_MESSAGES: dict[str, str] = {
"no_route": "暂不支持该渠道的消息处理。",
"route_error": "消息路由异常,请稍后重试。",
"dispatch_error": "消息处理异常,请稍后重试。",
"approval_denied": "操作未通过审批。",
"approval_check_error": "审批校验异常,请稍后重试。",
"rate_limited": "当前消息量较大,请稍后重试。",
}
class MessageProcessor:
def __init__(
self,
stream_fn: Callable[..., AsyncIterator[bytes]],
channel_manager=None,
*,
msg_repo=None,
binding_repo=None,
user_mapping_repo=None,
thread_mapping_repo=None,
allowlist=None,
pairing=None,
identity_link_resolver=None,
response_prefix: str = "",
human_delay: tuple[float, float] | None = None,
):
self._bridge = AgentBridge(stream_fn, channel_manager)
self._msg_repo = msg_repo or ChannelMsgRecordRepository()
self._binding_repo = binding_repo or ChannelBindingRepository()
self._user_mapping_repo = user_mapping_repo or ChannelUserMappingRepository()
self._thread_mapping_repo = thread_mapping_repo or ChannelThreadMappingRepository()
self._hooks = LifecycleHookRegistry()
self._allowlist = allowlist or AllowlistChecker()
if pairing:
self._pairing = pairing
else:
self._pairing = PairingManager(InMemoryPairingStore(), allowlist=self._allowlist)
if allowlist is None:
self._allowlist.load_dm_allowlist(["*"])
self._allowlist.load_group_allowlist(["*"])
self._response_prefix = response_prefix
self._human_delay = human_delay
self._identity_links = identity_link_resolver or IdentityLinkResolver()
self._matcher = RouteMatcher()
self._session_key_builder = SessionKeyBuilder()
self._active_message_count = 0
self._active_count_lock = asyncio.Lock()
self._cached_bindings: list[RouteBinding] = []
self._bindings_ts: float = 0.0
self._fence = ConversationFence()
self._self_user_ids: dict[str, set[str]] = {}
def set_self_user_ids(self, channel_type: str, account_id: str, user_ids: set[str]) -> None:
key = f"{channel_type}:{account_id}"
self._self_user_ids[key] = user_ids
@property
def hooks(self) -> LifecycleHookRegistry:
return self._hooks
async def _get_bindings(self) -> list[RouteBinding]:
now = time.monotonic()
if now - self._bindings_ts < _BINDINGS_CACHE_TTL and self._cached_bindings:
return self._cached_bindings
db_bindings = await self._binding_repo.list_all()
self._cached_bindings = [to_route_binding(b) for b in db_bindings]
self._bindings_ts = now
await self._matcher.set_bindings(self._cached_bindings)
return self._cached_bindings
async def _get_user_for_channel(self, internal_user_id: str | None) -> User | None:
if not internal_user_id:
return None
try:
uid = int(internal_user_id)
except (ValueError, TypeError):
return None
async with pg_manager.get_async_session_context() as db:
return (await db.execute(select(User).where(User.id == uid))).scalar_one_or_none()
def _is_echo(self, msg: UnifiedMessage) -> bool:
key = f"{msg.channel_type}:{msg.account_id}"
self_ids = self._self_user_ids.get(key)
if not self_ids:
return False
return msg.sender.id in self_ids
async def _ensure_security_loaded(self) -> None:
await self._identity_links.ensure_loaded()
self._session_key_builder = SessionKeyBuilder(
identity_links=self._identity_links.list_links()
)
async def process(self, msg: UnifiedMessage) -> None:
if self._is_echo(msg):
logger.debug("Skipping echo message: %s/%s", msg.channel_type, msg.msg_id)
return
limiter = await rate_limit_manager.get_limiter(msg.channel_type, msg.account_id)
acquire_result = await limiter.try_acquire()
if not acquire_result:
record_rate_limit_reject(msg.channel_type)
logger.warning(
"Rate limited: channel=%s account=%s retry_after=%.2fs",
msg.channel_type, msg.account_id, acquire_result.retry_after_sec,
)
await self._send_fallback(msg, _FALLBACK_MESSAGES["rate_limited"])
return
if not await claim(msg.msg_id, ttl=_IDEMPOTENCY_TTL):
logger.debug("Duplicate/inflight message ignored: %s", msg.msg_id)
return
try:
async with self._active_count_lock:
self._active_message_count += 1
set_inflight(self._active_message_count)
await self._ensure_security_loaded()
await self._process_impl(msg, limiter)
finally:
await release(msg.msg_id)
async with self._active_count_lock:
self._active_message_count = max(0, self._active_message_count - 1)
set_inflight(self._active_message_count)
async def _process_impl(self, msg: UnifiedMessage, limiter) -> None:
trace = await create_channel_trace(msg)
await trace.add_span("idempotency_check", metadata={"claimed": True})
conversation_key = self._fence.key_for(msg)
version, stop_event = self._fence.enter(conversation_key)
lock = self._fence.lock_for(conversation_key)
logger.debug(
"processor: phase=enter msg_id=%s channel=%s conv_key=%s version=%d",
msg.msg_id, msg.channel_type, conversation_key, version,
)
async with lock:
if version != self._fence.current_version(conversation_key):
logger.debug("Foreground fence: message superseded for %s", conversation_key)
record_message(msg.channel_type, "superseded")
await trace.finish(error="superseded")
return
record_message(msg.channel_type, "received")
plugin = ChannelPluginRegistry.get(msg.channel_type)
if isinstance(plugin, InboundMessageHook):
try:
result = await plugin.on_message_received(msg)
if result is None:
logger.info("Message dropped by inbound hook: %s", msg.msg_id)
await trace.finish(error="dropped_by_hook")
return
except Exception:
logger.exception("Inbound hook failed, proceeding: %s", msg.msg_id)
try:
await self._hooks.fire(HookEvent.MESSAGE_RECEIVED, msg)
except Exception:
logger.exception("MESSAGE_RECEIVED hook failed: %s", msg.msg_id)
chat_id = msg.group.id if msg.group and msg.group.id else msg.sender.id
chat_type = "group" if msg.sender.kind != PeerKind.DIRECT else "direct"
record_data = {
"channel_id": msg.account_id,
"channel_type": msg.channel_type,
"message_id": msg.msg_id,
"chat_id": chat_id,
"chat_type": chat_type,
"content_type": msg.message_type.value if msg.message_type else "text",
"sender_user_id": msg.sender.id,
"content_preview": (msg.content or "")[:500],
"reply_to_message_id": msg.metadata.get("reply_to_message_id") if msg.metadata else None,
"extra_data": msg.raw_payload,
}
record = None
start_time = time.monotonic()
try:
record = await self._msg_repo.create_record(record_data)
if stop_event.is_set():
logger.debug("Aborted before dispatch for %s", conversation_key)
elapsed_ms = int((time.monotonic() - start_time) * 1000)
await self._msg_repo.mark_error(record.id, "aborted_by_fence", elapsed_ms)
await trace.finish(error="aborted_by_fence")
return
bindings = await self._get_bindings()
dispatcher = MessageDispatcher(
bindings,
self._allowlist,
self._pairing,
matcher=self._matcher,
session_key_builder=self._session_key_builder,
user_mapping_repo=self._user_mapping_repo,
thread_mapping_repo=self._thread_mapping_repo,
)
try:
await self._hooks.fire(HookEvent.BEFORE_AGENT_RUN, msg)
except Exception:
logger.exception("BEFORE_AGENT_RUN hook failed: %s", msg.msg_id)
result = await dispatcher.dispatch(msg, trace=trace)
record_dispatch_duration_ms(
msg.channel_type,
(time.monotonic() - start_time) * 1000,
)
logger.debug(
"processor: phase=dispatched msg_id=%s handled_by=%s success=%s agent_config_id=%s",
msg.msg_id, result.handled_by, result.success, result.agent_config_id,
)
if result.handled_by == "command":
if result.command_response:
await self._send_reply(msg, result.command_response)
try:
await self._hooks.fire(HookEvent.AFTER_AGENT_RUN, msg, result, result.command_response)
except Exception:
logger.exception("AFTER_AGENT_RUN hook failed: %s", msg.msg_id)
elapsed_ms = int((time.monotonic() - start_time) * 1000)
await self._msg_repo.mark_success(
record.id,
"",
(result.command_response or "")[:500],
elapsed_ms,
)
await trace.finish()
return
if not result.success:
try:
await self._hooks.fire(HookEvent.AFTER_AGENT_RUN, msg, result, None)
except Exception:
logger.exception("AFTER_AGENT_RUN hook failed: %s", msg.msg_id)
error_msg = result.error or "dispatch_failed"
fallback = _FALLBACK_MESSAGES.get(error_msg)
if fallback is None and result.error and result.error.startswith("approval_denied:"):
fallback = _FALLBACK_MESSAGES["approval_denied"]
if fallback:
await self._send_fallback(msg, fallback)
elapsed_ms = int((time.monotonic() - start_time) * 1000)
await self._msg_repo.mark_error(record.id, error_msg, elapsed_ms)
await trace.finish(error=error_msg)
return
if stop_event.is_set():
logger.debug("Aborted after dispatch for %s", conversation_key)
try:
await self._hooks.fire(HookEvent.AFTER_AGENT_RUN, msg, result, None)
except Exception:
logger.exception("AFTER_AGENT_RUN hook failed: %s", msg.msg_id)
elapsed_ms = int((time.monotonic() - start_time) * 1000)
await self._msg_repo.mark_error(record.id, "aborted_by_fence", elapsed_ms)
await trace.finish(error="aborted_by_fence")
return
logger.debug(
"processor: phase=reply_start msg_id=%s channel=%s agent_config_id=%s",
msg.msg_id, msg.channel_type, result.agent_config_id,
)
await self._handle_agent_reply(msg, result, plugin, stop_event, record, start_time, conversation_key, trace, limiter)
except Exception as e:
logger.exception("Message processing failed: %s", msg.msg_id)
await trace.finish(error=f"{type(e).__name__}: {e}")
try:
await self._hooks.fire(HookEvent.AFTER_AGENT_RUN, msg, None, None)
except Exception:
logger.exception("AFTER_AGENT_RUN hook failed: %s", msg.msg_id)
elapsed_ms = int((time.monotonic() - start_time) * 1000)
error_text = f"{type(e).__name__}: {e}"
if record is not None:
await self._msg_repo.mark_error(record.id, error_text[:500], elapsed_ms)
record_message(msg.channel_type, "error")
async def _handle_agent_reply(
self,
msg: UnifiedMessage,
result,
plugin,
stop_event: asyncio.Event | None,
record,
start_time: float,
conversation_key: str,
trace: ChannelTrace,
limiter,
) -> None:
user = await self._get_user_for_channel(result.internal_user_id)
if user is None:
logger.error(
"channel_turn skip: msg_id=%s channel=%s reason=no_user_for_internal_id:%s",
msg.msg_id,
msg.channel_type,
result.internal_user_id,
)
elapsed_ms = int((time.monotonic() - start_time) * 1000)
await self._msg_repo.mark_error(record.id, "no_user_for_internal_id", elapsed_ms)
await trace.finish(error="no_user_for_internal_id")
return
agent_config_id = result.agent_config_id
chat_id = msg.group.id if msg.group and msg.group.id else msg.sender.id
try:
await self._hooks.fire(HookEvent.AGENT_BOOTSTRAP, msg, result)
except Exception:
logger.exception("AGENT_BOOTSTRAP hook failed: %s", msg.msg_id)
stream_status: StreamStatus | None = None
if plugin is not None and isinstance(plugin, OutboundProtocol):
try:
await self._hooks.fire(HookEvent.MESSAGE_SENDING, msg, "")
except Exception:
logger.exception("MESSAGE_SENDING hook failed: %s", msg.msg_id)
reply_text, stream_status = await self._process_with_pipeline(
msg,
result,
agent_config_id,
user,
plugin,
stop_event,
limiter=limiter,
trace=trace,
)
else:
async with pg_manager.get_async_session_context() as db:
stream_result = await self._bridge.invoke_with_result(
msg,
result,
agent_config_id,
user,
db,
)
reply_text = stream_result.accumulated_text or None
stream_status = stream_result.status
if reply_text:
if plugin is not None:
reply_text = await self._apply_sending_hooks(msg, plugin, chat_id, reply_text)
if reply_text is None:
elapsed_ms = int((time.monotonic() - start_time) * 1000)
await self._msg_repo.mark_error(record.id, "blocked_by_sending_hook", elapsed_ms)
await trace.finish(error="blocked_by_sending_hook")
return
success = await self._send_with_retry(plugin, chat_id, reply_text, reply_to_id=msg.msg_id)
if not success:
logger.exception("Failed to send reply via %s after retries", msg.channel_type)
else:
logger.warning("No outbound plugin for channel: %s", msg.channel_type)
elapsed_ms = int((time.monotonic() - start_time) * 1000)
await self._msg_repo.mark_error(record.id, "no_outbound_plugin", elapsed_ms)
await trace.finish(error="no_outbound_plugin")
return
if not reply_text and stream_status in (StreamStatus.TIMEOUT, StreamStatus.ERROR, StreamStatus.CANCELLED):
fallback_map = {
StreamStatus.TIMEOUT: "处理超时,请稍后重试。",
StreamStatus.ERROR: "处理异常,已通知管理员。",
StreamStatus.CANCELLED: "处理已取消。",
}
fallback = fallback_map.get(stream_status)
if fallback:
await self._send_fallback(msg, fallback)
elapsed_ms = int((time.monotonic() - start_time) * 1000)
await self._msg_repo.mark_error(record.id, stream_status.value, elapsed_ms)
await trace.finish(error=stream_status.value)
return
if isinstance(plugin, SendLifecycleHook):
try:
await plugin.after_send_success(
chat_id,
{"content": reply_text, "msg_id": msg.msg_id},
)
except Exception:
logger.exception("Send lifecycle hook failed: %s", msg.msg_id)
try:
await self._hooks.fire(HookEvent.MESSAGE_SENT, msg, reply_text)
except Exception:
logger.exception("MESSAGE_SENT hook failed: %s", msg.msg_id)
try:
await self._hooks.fire(HookEvent.AFTER_AGENT_RUN, msg, result, reply_text)
except Exception:
logger.exception("AFTER_AGENT_RUN hook failed: %s", msg.msg_id)
elapsed_ms = int((time.monotonic() - start_time) * 1000)
await self._msg_repo.mark_success(
record.id,
"",
(reply_text or "")[:500],
elapsed_ms,
)
record_message(msg.channel_type, "success")
await trace.finish()
async def _apply_sending_hooks(
self,
msg: UnifiedMessage,
plugin,
chat_id: str,
reply_text: str,
) -> str | None:
if isinstance(plugin, MessageSendingHook):
try:
modified = await plugin.on_message_sending(msg, reply_text)
if modified is None:
logger.info("Message blocked by MessageSendingHook: %s", msg.msg_id)
return None
reply_text = modified
except Exception:
logger.exception("MessageSendingHook failed: %s", msg.msg_id)
if isinstance(plugin, SendLifecycleHook):
try:
await plugin.before_send_attempt(chat_id, reply_text)
except Exception:
logger.exception("before_send_attempt failed: %s", msg.msg_id)
try:
result = await self._hooks.fire_sequential(
HookEvent.MESSAGE_SENDING,
initial=reply_text,
msg=msg,
)
reply_text = result if result is not None else reply_text
except Exception:
logger.exception("MESSAGE_SENDING hook failed: %s", msg.msg_id)
return reply_text
async def _process_with_pipeline(
self,
msg: UnifiedMessage,
result,
agent_config_id: int,
user,
plugin,
stop_event: asyncio.Event | None = None,
*,
limiter=None,
trace: "ChannelTrace | None" = None,
) -> tuple[str | None, StreamStatus | None]:
target_id = msg.group.id if msg.group and msg.group.id else msg.sender.id
async def send_fn(content: str) -> str | None:
if stop_event and stop_event.is_set():
return None
if not content:
return None
send_content = content
if isinstance(plugin, MessageSendingHook):
try:
modified = await plugin.on_message_sending(msg, send_content)
if modified is None:
logger.info("Pipeline send blocked by MessageSendingHook: %s", msg.msg_id)
return None
send_content = modified
except Exception:
logger.exception("Pipeline MessageSendingHook failed: %s", msg.msg_id)
if isinstance(plugin, SendLifecycleHook):
try:
await plugin.before_send_attempt(target_id, send_content)
except Exception:
logger.exception("Pipeline before_send_attempt failed: %s", msg.msg_id)
success = await self._send_with_retry(plugin, target_id, send_content, reply_to_id=msg.msg_id)
if not success:
logger.exception("Pipeline send failed via %s after retries", msg.channel_type)
return None
pipeline = create_block_reply_pipeline(
send_fn=send_fn,
msg=msg,
response_prefix=self._response_prefix,
human_delay=self._human_delay,
)
accumulated: list[str] = []
stream_error: StreamStatus | None = None
exec_span_id = None
if trace:
exec_span_id = await trace.add_span(
"agent_execution",
metadata={
"agent_config_id": agent_config_id,
"thread_id": result.thread_id,
},
)
async with pg_manager.get_async_session_context() as db:
try:
stream_iter = self._bridge.invoke_stream(
msg,
result,
agent_config_id,
user,
db,
)
if limiter is not None:
stream_iter = limiter.run_with_limit(stream_iter)
async for chunk in stream_iter:
if stop_event and stop_event.is_set():
await pipeline.abort()
logger.debug("Pipeline aborted by foreground fence for %s", self._fence.key_for(msg))
break
await pipeline.enqueue(chunk)
if chunk.content:
accumulated.append(chunk.content)
except StreamTimeoutError:
stream_error = StreamStatus.TIMEOUT
logger.warning("Agent stream timeout for channel=%s peer=%s", msg.channel_type, msg.sender.id)
except StreamCancelledError:
stream_error = StreamStatus.CANCELLED
logger.debug("Agent stream cancelled for channel=%s peer=%s", msg.channel_type, msg.sender.id)
await pipeline.wait_idle()
if exec_span_id:
status_msg = stream_error.value if stream_error else "ok"
await trace.end_span(
exec_span_id,
level="ERROR" if stream_error else "DEFAULT",
status_message=status_msg,
metadata={"stream_status": status_msg, "chunks": len(accumulated)},
)
if stop_event and stop_event.is_set():
return None, None
return ("".join(accumulated) if accumulated else None, stream_error)
async def _send_fallback(self, msg: UnifiedMessage, fallback_text: str) -> None:
plugin = ChannelPluginRegistry.get(msg.channel_type)
if plugin is None or not isinstance(plugin, OutboundProtocol):
return
try:
target_id = msg.group.id if msg.group and msg.group.id else msg.sender.id
await plugin.send_text(target_id, fallback_text, reply_to_id=msg.msg_id)
except Exception:
logger.exception("Failed to send fallback reply via %s", msg.channel_type)
async def _send_with_retry(
self,
plugin,
target_id: str,
content: str,
*,
reply_to_id: str | None = None,
max_retries: int = 3,
circuit_breaker: "CircuitBreaker | None" = None,
) -> bool:
last_error = None
total_attempts = max_retries + 1
for attempt in range(total_attempts):
if circuit_breaker is not None and not circuit_breaker.allow_request():
logger.warning(
"Circuit breaker OPEN for %s, skipping send",
plugin.id if hasattr(plugin, "id") else "unknown",
)
return False
try:
await plugin.send_text(target_id, content, reply_to_id=reply_to_id)
if circuit_breaker is not None:
circuit_breaker.record_success()
return True
except Exception as e:
last_error = e
if circuit_breaker is not None:
circuit_breaker.record_failure()
logger.warning(
"Send attempt %d/%d failed (target=%s): %s",
attempt + 1, total_attempts, target_id, e,
)
if attempt < max_retries:
delay = min(0.3 * (2 ** attempt), 30.0)
await asyncio.sleep(delay)
logger.error("All %d send attempts failed (target=%s): %s", total_attempts, target_id, last_error)
if isinstance(plugin, SendLifecycleHook):
try:
await plugin.after_send_failure(target_id, str(last_error))
except Exception:
pass
return False
async def _send_reply(self, msg: UnifiedMessage, content: str) -> str | None:
plugin = ChannelPluginRegistry.get(msg.channel_type)
if plugin is None or not isinstance(plugin, OutboundProtocol):
logger.warning("No outbound plugin for channel: %s", msg.channel_type)
return None
target_id = msg.group.id if msg.group and msg.group.id else msg.sender.id
reply_to_id = msg.msg_id
send_content = await self._apply_sending_hooks(msg, plugin, target_id, content)
if send_content is None:
return None
success = await self._send_with_retry(plugin, target_id, send_content, reply_to_id=reply_to_id)
if not success:
logger.exception("Failed to send reply via %s after retries", msg.channel_type)
return None

View File

@ -0,0 +1,217 @@
import asyncio
import logging
import time
from collections.abc import AsyncIterator, Callable
from dataclasses import dataclass
logger = logging.getLogger(__name__)
_global_semaphore = asyncio.Semaphore(50)
def get_global_semaphore() -> asyncio.Semaphore:
return _global_semaphore
DEFAULT_RATE = 10.0
DEFAULT_CAPACITY = 20
DEFAULT_PER_ACCOUNT_CONCURRENCY = 5
DEFAULT_QUEUE_DEPTH = 20
DEFAULT_MAX_LIMITERS = 10_000
@dataclass(slots=True)
class AcquireResult:
allowed: bool
retry_after_sec: float = 0.0
def __bool__(self) -> bool:
return self.allowed
class TokenBucket:
"""令牌桶算法 — 控制单账户消息速率。
参数
- rate: 每秒补充的令牌数默认 10
- capacity: 桶容量上限默认 20允许短时突发
- time_func: 可选时钟注入便于测试默认 time.monotonic
"""
def __init__(
self,
rate: float = DEFAULT_RATE,
capacity: int = DEFAULT_CAPACITY,
time_func: Callable[[], float] = time.monotonic,
):
self._rate = rate
self._capacity = capacity
self._tokens = float(capacity)
self._time = time_func
self._last_refill = self._time()
self._lock = asyncio.Lock()
async def acquire(self) -> AcquireResult:
async with self._lock:
now = self._time()
elapsed = now - self._last_refill
self._tokens = min(self._capacity, self._tokens + elapsed * self._rate)
self._last_refill = now
if self._tokens >= 1.0:
self._tokens -= 1.0
return AcquireResult(allowed=True)
if self._rate <= 0.0:
return AcquireResult(allowed=False, retry_after_sec=float("inf"))
retry_after = (1.0 - self._tokens) / self._rate
return AcquireResult(allowed=False, retry_after_sec=retry_after)
def reset(self) -> None:
self._tokens = float(self._capacity)
self._last_refill = self._time()
class ChannelRateLimiter:
"""单渠道账户的限流控制器,组合令牌桶 + 并发信号量。
线程安全所有方法在 asyncio 上下文中调用
"""
def __init__(
self,
rate: float = DEFAULT_RATE,
per_account_concurrency: int = DEFAULT_PER_ACCOUNT_CONCURRENCY,
queue_depth: int = DEFAULT_QUEUE_DEPTH,
):
self._rate = rate
self._per_account_concurrency = per_account_concurrency
self._bucket = TokenBucket(rate=rate, capacity=DEFAULT_CAPACITY)
self._semaphore = asyncio.Semaphore(per_account_concurrency)
self._queue_depth = queue_depth
self._waiting: int = 0
async def try_acquire(self) -> AcquireResult:
"""尝试获取消息处理许可。
返回 AcquireResult(allowed=..., retry_after_sec=...)
"""
result = await self._bucket.acquire()
if result.allowed:
return result
if self._waiting >= self._queue_depth:
logger.warning(
"Rate limit queue full (depth=%d), rejecting message",
self._queue_depth,
)
return AcquireResult(allowed=False, retry_after_sec=result.retry_after_sec)
self._waiting += 1
try:
await asyncio.sleep(result.retry_after_sec)
retry = await self._bucket.acquire()
if not retry.allowed:
return AcquireResult(allowed=False, retry_after_sec=retry.retry_after_sec)
return retry
finally:
self._waiting -= 1
async def run_with_limit(self, coro) -> AsyncIterator:
"""在并发限制内执行 Agent 调用。
先获取账户级信号量再获取全局信号量
"""
async with self._semaphore:
async with _global_semaphore:
async for chunk in coro:
yield chunk
def reset(self) -> None:
self._bucket.reset()
@property
def waiting_count(self) -> int:
return self._waiting
@property
def config(self) -> dict:
return {
"rate": self._rate,
"per_account_concurrency": self._per_account_concurrency,
"queue_depth": self._queue_depth,
}
class RateLimitManager:
"""限流管理器 — 按账户维度管理限流器实例。
每个 (channel_type, account_id) 组合共享同一个 ChannelRateLimiter
内置自动清理机制防止无限增长
"""
def __init__(self, max_limiters: int = DEFAULT_MAX_LIMITERS):
self._limiters: dict[str, tuple[float, ChannelRateLimiter]] = {}
self._max_limiters = max_limiters
self._last_cleanup = time.monotonic()
self._lock = asyncio.Lock()
def _key(self, channel_type: str, account_id: str) -> str:
return f"{channel_type}:{account_id}"
async def get_limiter(
self,
channel_type: str,
account_id: str,
rate: float = DEFAULT_RATE,
per_account_concurrency: int = DEFAULT_PER_ACCOUNT_CONCURRENCY,
) -> ChannelRateLimiter:
async with self._lock:
self._maybe_cleanup()
key = self._key(channel_type, account_id)
if key not in self._limiters:
self._limiters[key] = (
time.monotonic(),
ChannelRateLimiter(
rate=rate,
per_account_concurrency=per_account_concurrency,
),
)
else:
self._limiters[key] = (time.monotonic(), self._limiters[key][1])
return self._limiters[key][1]
async def remove_limiter(self, channel_type: str, account_id: str) -> None:
async with self._lock:
key = self._key(channel_type, account_id)
self._limiters.pop(key, None)
async def clear(self) -> None:
async with self._lock:
self._limiters.clear()
async def get_stats(self) -> dict[str, dict]:
async with self._lock:
return {
key: {
"waiting": limiter.waiting_count,
**limiter.config,
}
for key, (_, limiter) in self._limiters.items()
}
def _maybe_cleanup(self) -> None:
if len(self._limiters) <= self._max_limiters:
return
now = time.monotonic()
if now - self._last_cleanup < 300:
return
self._last_cleanup = now
expired = [key for key, (last_touch, _) in self._limiters.items() if now - last_touch > 3600]
for key in expired:
self._limiters.pop(key, None)
if expired:
logger.info("Cleaned up %d stale rate limiters", len(expired))
rate_limit_manager = RateLimitManager()

View File

@ -0,0 +1,269 @@
import asyncio
import logging
import random
import time
from collections.abc import Callable, Coroutine
from typing import Any
from yuxi.channel.message.models import ReplyPayload, ReplyStage, StreamingChunk, UnifiedMessage
logger = logging.getLogger(__name__)
_DEFAULT_MIN_CHARS = 80
_DEFAULT_DEBOUNCE_MS = 150.0
_DEFAULT_TIMEOUT_MS = 10_000.0
OnIdleCallback = Callable[[], None]
OnErrorCallback = Callable[[Exception], None]
BeforeDeliverCallback = Callable[[str], Coroutine[Any, Any, str | None]]
class ReplyDispatcher:
def __init__(
self,
send_fn: Callable[..., Coroutine[Any, Any, str | None]],
msg: UnifiedMessage,
*,
response_prefix: str = "",
human_delay: tuple[float, float] | None = None,
timeout_ms: float = _DEFAULT_TIMEOUT_MS,
on_idle: OnIdleCallback | None = None,
on_error: OnErrorCallback | None = None,
before_deliver: BeforeDeliverCallback | None = None,
):
self._send_fn = send_fn
self._msg = msg
self._response_prefix = response_prefix
self._human_delay = human_delay
self._timeout_ms = timeout_ms
self._on_idle = on_idle
self._on_error = on_error
self._before_deliver = before_deliver
self._send_chain: asyncio.Lock = asyncio.Lock()
self._flush_lock: asyncio.Lock = asyncio.Lock()
self._pending = 0
self._complete = asyncio.Event()
self._done = False
self._blocks: list[str] = []
self._block_deadline: float = 0.0
self._block_timer: asyncio.Task | None = None
self._sent_payload_keys: set[str] = set()
self._sent_content_keys: set[str] = set()
self._tool_count = 0
self._block_count = 0
self._final_count = 0
self._failed_count = 0
self._cancelled_count = 0
self._started_at: float | None = None
self._abort_event = asyncio.Event()
async def enqueue(self, chunk: StreamingChunk) -> bool:
if self._abort_event.is_set():
self._cancelled_count += 1
return False
if chunk.stage == ReplyStage.TOOL:
return await self._dispatch_tool(chunk)
elif chunk.stage == ReplyStage.BLOCK:
return await self._dispatch_block(chunk)
elif chunk.stage == ReplyStage.FINAL:
return await self._dispatch_final(chunk)
return False
async def _dispatch_tool(self, chunk: StreamingChunk) -> bool:
self._pending += 1
self._tool_count += 1
try:
prefix = self._response_prefix if self._response_prefix else ""
content = f"{prefix}{chunk.content}" if chunk.content else ""
if not content:
return False
await self._send_ordered(content)
self._block_count += 1
return True
except Exception as e:
logger.exception("Tool dispatch failed")
self._failed_count += 1
if self._on_error:
self._on_error(e)
raise
finally:
self._pending -= 1
self._maybe_fire_idle()
async def _dispatch_block(self, chunk: StreamingChunk) -> bool:
if not chunk.content:
return False
async with self._flush_lock:
self._blocks.append(chunk.content)
now = time.monotonic()
if self._block_deadline == 0.0:
self._block_deadline = now + _DEFAULT_DEBOUNCE_MS / 1000.0
self._block_timer = asyncio.create_task(self._flush_after_debounce())
if sum(len(b) for b in self._blocks) >= _DEFAULT_MIN_CHARS:
await self._flush_blocks()
return True
async def _flush_after_debounce(self) -> None:
try:
await asyncio.sleep(_DEFAULT_DEBOUNCE_MS / 1000.0)
except asyncio.CancelledError:
return
async with self._flush_lock:
await self._flush_blocks()
async def _flush_blocks(self) -> None:
if not self._blocks:
return
combined = "".join(self._blocks)
self._blocks.clear()
self._block_deadline = 0.0
if self._block_timer and not self._block_timer.done():
self._block_timer.cancel()
self._block_timer = None
payload = ReplyPayload(
target_id=self._target_id,
content=combined,
reply_to_id=self._msg.msg_id if self._block_count == 0 else None,
)
if payload.content_key in self._sent_content_keys:
return
self._sent_content_keys.add(payload.content_key)
self._pending += 1
self._block_count += 1
try:
await self._send_ordered(combined)
except Exception as e:
logger.exception("Block dispatch failed")
self._failed_count += 1
if self._on_error:
self._on_error(e)
raise
finally:
self._pending -= 1
self._maybe_fire_idle()
async def _dispatch_final(self, chunk: StreamingChunk) -> bool:
async with self._flush_lock:
await self._flush_blocks()
if chunk.content and chunk.content != "finished":
self._pending += 1
self._final_count += 1
try:
await self._send_ordered(chunk.content)
except Exception as e:
logger.exception("Final dispatch failed")
self._failed_count += 1
if self._on_error:
self._on_error(e)
raise
finally:
self._pending -= 1
self._maybe_fire_idle()
self._done = True
self._complete.set()
return True
self._done = True
self._complete.set()
self._maybe_fire_idle()
return True
async def _send_ordered(self, content: str) -> None:
async with self._send_chain:
if self._abort_event.is_set():
return
if self._human_delay and self._block_count > 0:
delay = self._human_delay[0] + (self._human_delay[1] - self._human_delay[0]) * random.random()
await asyncio.sleep(delay / 1000.0)
if self._response_prefix and self._block_count == 0:
prefix = self._response_prefix.format(
agent_name=self._msg.channel_type,
model="",
)
content = f"{prefix}{content}"
if self._before_deliver:
try:
content = await self._before_deliver(content)
except Exception as e:
logger.exception("before_deliver hook failed")
if self._on_error:
self._on_error(e)
return
if content is None:
self._cancelled_count += 1
return
await self._send_fn(content)
@property
def _target_id(self) -> str:
return self._msg.group.id if self._msg.group and self._msg.group.id else self._msg.sender.id
async def wait_idle(self, timeout_ms: float | None = None) -> None:
timeout = timeout_ms or self._timeout_ms
try:
await asyncio.wait_for(self._complete.wait(), timeout=timeout / 1000.0)
except TimeoutError:
self._abort_event.set()
self._done = True
self._complete.set()
logger.warning("ReplyDispatcher idle timeout after %.0fms, aborting", timeout)
def mark_complete(self) -> None:
self._done = True
self._complete.set()
self._maybe_fire_idle()
async def abort(self) -> None:
async with self._flush_lock:
if self._blocks:
await self._flush_blocks()
self._abort_event.set()
self._complete.set()
def _maybe_fire_idle(self) -> None:
if self._done and self._pending == 0 and self._on_idle:
self._on_idle()
@property
def tool_count(self) -> int:
return self._tool_count
@property
def block_count(self) -> int:
return self._block_count
@property
def final_count(self) -> int:
return self._final_count
@property
def failed_count(self) -> int:
return self._failed_count
@property
def cancelled_count(self) -> int:
return self._cancelled_count
@property
def queued_counts(self) -> dict[str, int]:
return {"tool": self._tool_count, "block": self._block_count, "final": self._final_count}
@property
def is_idle(self) -> bool:
return self._done and self._pending == 0