diff --git a/backend/package/yuxi/channel/message/__init__.py b/backend/package/yuxi/channel/message/__init__.py new file mode 100644 index 00000000..c69a0fa8 --- /dev/null +++ b/backend/package/yuxi/channel/message/__init__.py @@ -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", +] \ No newline at end of file diff --git a/backend/package/yuxi/channel/message/block_reply_pipeline.py b/backend/package/yuxi/channel/message/block_reply_pipeline.py new file mode 100644 index 00000000..a26be657 --- /dev/null +++ b/backend/package/yuxi/channel/message/block_reply_pipeline.py @@ -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, + ) diff --git a/backend/package/yuxi/channel/message/bridge.py b/backend/package/yuxi/channel/message/bridge.py new file mode 100644 index 00000000..3e6adff4 --- /dev/null +++ b/backend/package/yuxi/channel/message/bridge.py @@ -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 diff --git a/backend/package/yuxi/channel/message/circuit_breaker.py b/backend/package/yuxi/channel/message/circuit_breaker.py new file mode 100644 index 00000000..3451c3be --- /dev/null +++ b/backend/package/yuxi/channel/message/circuit_breaker.py @@ -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 \ No newline at end of file diff --git a/backend/package/yuxi/channel/message/conversation_fence.py b/backend/package/yuxi/channel/message/conversation_fence.py new file mode 100644 index 00000000..d2592d6a --- /dev/null +++ b/backend/package/yuxi/channel/message/conversation_fence.py @@ -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) \ No newline at end of file diff --git a/backend/package/yuxi/channel/message/dispatch.py b/backend/package/yuxi/channel/message/dispatch.py new file mode 100644 index 00000000..3ac8926d --- /dev/null +++ b/backend/package/yuxi/channel/message/dispatch.py @@ -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 diff --git a/backend/package/yuxi/channel/message/durable_send.py b/backend/package/yuxi/channel/message/durable_send.py new file mode 100644 index 00000000..d1e25f4a --- /dev/null +++ b/backend/package/yuxi/channel/message/durable_send.py @@ -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 diff --git a/backend/package/yuxi/channel/message/handler.py b/backend/package/yuxi/channel/message/handler.py new file mode 100644 index 00000000..3d74ed66 --- /dev/null +++ b/backend/package/yuxi/channel/message/handler.py @@ -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 diff --git a/backend/package/yuxi/channel/message/idempotency.py b/backend/package/yuxi/channel/message/idempotency.py new file mode 100644 index 00000000..b6ab9e13 --- /dev/null +++ b/backend/package/yuxi/channel/message/idempotency.py @@ -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 \ No newline at end of file diff --git a/backend/package/yuxi/channel/message/langfuse_trace.py b/backend/package/yuxi/channel/message/langfuse_trace.py new file mode 100644 index 00000000..8076bdc5 --- /dev/null +++ b/backend/package/yuxi/channel/message/langfuse_trace.py @@ -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 diff --git a/backend/package/yuxi/channel/message/media.py b/backend/package/yuxi/channel/message/media.py new file mode 100644 index 00000000..bc3be79f --- /dev/null +++ b/backend/package/yuxi/channel/message/media.py @@ -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" 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), + ) diff --git a/backend/package/yuxi/channel/message/media_cleaner.py b/backend/package/yuxi/channel/message/media_cleaner.py new file mode 100644 index 00000000..17f8e8c9 --- /dev/null +++ b/backend/package/yuxi/channel/message/media_cleaner.py @@ -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") diff --git a/backend/package/yuxi/channel/message/media_store.py b/backend/package/yuxi/channel/message/media_store.py new file mode 100644 index 00000000..d9fc5587 --- /dev/null +++ b/backend/package/yuxi/channel/message/media_store.py @@ -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) diff --git a/backend/package/yuxi/channel/message/metrics.py b/backend/package/yuxi/channel/message/metrics.py new file mode 100644 index 00000000..01afa3c0 --- /dev/null +++ b/backend/package/yuxi/channel/message/metrics.py @@ -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) \ No newline at end of file diff --git a/backend/package/yuxi/channel/message/models.py b/backend/package/yuxi/channel/message/models.py new file mode 100644 index 00000000..76f02a68 --- /dev/null +++ b/backend/package/yuxi/channel/message/models.py @@ -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) diff --git a/backend/package/yuxi/channel/message/processor.py b/backend/package/yuxi/channel/message/processor.py new file mode 100644 index 00000000..c327889d --- /dev/null +++ b/backend/package/yuxi/channel/message/processor.py @@ -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 diff --git a/backend/package/yuxi/channel/message/rate_limiter.py b/backend/package/yuxi/channel/message/rate_limiter.py new file mode 100644 index 00000000..c494dbfb --- /dev/null +++ b/backend/package/yuxi/channel/message/rate_limiter.py @@ -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() diff --git a/backend/package/yuxi/channel/message/reply_dispatcher.py b/backend/package/yuxi/channel/message/reply_dispatcher.py new file mode 100644 index 00000000..a6879569 --- /dev/null +++ b/backend/package/yuxi/channel/message/reply_dispatcher.py @@ -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