feat(channel/message): 新增消息处理核心模块及相关工具类
本次提交新增了渠道消息处理的完整核心模块,包含以下核心功能: 1. 新增会话围栏类,实现会话并发控制与过期清理 2. 新增媒体清理器,实现过期媒体文件自动清理 3. 新增熔断器组件,实现服务降级与故障隔离 4. 新增消息处理器,完成渠道消息的完整流转处理 5. 新增限流器组件,实现渠道级和账户级流量控制 6. 新增链路追踪模块,集成Langfuse实现调用链路监控 7. 新增指标统计模块,实现消息处理全链路指标采集 8. 新增统一消息模型,封装全渠道消息格式 9. 新增块回复流水线,实现流式回复的合并与去重 10. 新增本地媒体存储模块,实现媒体文件的本地管理 11. 新增回复分发器,实现回复内容的有序发送与延迟处理 12. 完善__init__.py导出所有核心模块与工具类
This commit is contained in:
parent
a81279cf61
commit
0b19470c70
114
backend/package/yuxi/channel/message/__init__.py
Normal file
114
backend/package/yuxi/channel/message/__init__.py
Normal file
@ -0,0 +1,114 @@
|
||||
from yuxi.channel.message.block_reply_pipeline import create_block_reply_pipeline
|
||||
from yuxi.channel.message.bridge import AgentBridge
|
||||
from yuxi.channel.message.circuit_breaker import (
|
||||
CircuitBreaker,
|
||||
CircuitBreakerConfig,
|
||||
CircuitBreakerOpenError,
|
||||
CircuitState,
|
||||
)
|
||||
from yuxi.channel.message.dispatch import MessageDispatcher
|
||||
from yuxi.channel.message.idempotency import (
|
||||
ClaimStatus,
|
||||
IdempotencyBackend,
|
||||
InMemoryBackend,
|
||||
RedisBackend,
|
||||
build_idempotency_key,
|
||||
check_and_set,
|
||||
claim,
|
||||
claim_with_status,
|
||||
clear_inflight,
|
||||
commit,
|
||||
consume,
|
||||
get_backend,
|
||||
parse_idempotency_key,
|
||||
release,
|
||||
release_and_forget,
|
||||
reset,
|
||||
reset_sync,
|
||||
)
|
||||
from yuxi.channel.message.langfuse_trace import ChannelTrace, create_channel_trace
|
||||
from yuxi.channel.message.media import shutdown_media
|
||||
from yuxi.channel.message.metrics import (
|
||||
MetricsTimer,
|
||||
channel_agent_duration_ms,
|
||||
channel_dispatch_duration_ms,
|
||||
channel_messages_inflight,
|
||||
channel_messages_total,
|
||||
channel_rate_limit_rejects_total,
|
||||
record_agent_duration_ms,
|
||||
record_dispatch_duration_ms,
|
||||
record_message,
|
||||
record_rate_limit_reject,
|
||||
registry,
|
||||
set_inflight,
|
||||
)
|
||||
from yuxi.channel.message.models import (
|
||||
ChunkType,
|
||||
DispatchResult,
|
||||
GroupContext,
|
||||
MentionSource,
|
||||
MessageReceipt,
|
||||
MessageType,
|
||||
PeerInfo,
|
||||
ReplyPayload,
|
||||
ReplyStage,
|
||||
StreamingChunk,
|
||||
UnifiedMessage,
|
||||
)
|
||||
from yuxi.channel.message.rate_limiter import ChannelRateLimiter as RateLimiter
|
||||
from yuxi.channel.message.rate_limiter import rate_limit_manager
|
||||
|
||||
__all__ = [
|
||||
"AgentBridge",
|
||||
"ChannelTrace",
|
||||
"ChunkType",
|
||||
"CircuitBreaker",
|
||||
"CircuitBreakerConfig",
|
||||
"CircuitBreakerOpenError",
|
||||
"CircuitState",
|
||||
"ClaimStatus",
|
||||
"DispatchResult",
|
||||
"GroupContext",
|
||||
"IdempotencyBackend",
|
||||
"InMemoryBackend",
|
||||
"MentionSource",
|
||||
"MessageDispatcher",
|
||||
"MessageReceipt",
|
||||
"MessageType",
|
||||
"MetricsTimer",
|
||||
"PeerInfo",
|
||||
"RateLimiter",
|
||||
"rate_limit_manager",
|
||||
"RedisBackend",
|
||||
"ReplyPayload",
|
||||
"ReplyStage",
|
||||
"StreamingChunk",
|
||||
"UnifiedMessage",
|
||||
"build_idempotency_key",
|
||||
"channel_agent_duration_ms",
|
||||
"channel_dispatch_duration_ms",
|
||||
"channel_messages_inflight",
|
||||
"channel_messages_total",
|
||||
"channel_rate_limit_rejects_total",
|
||||
"check_and_set",
|
||||
"claim",
|
||||
"claim_with_status",
|
||||
"clear_inflight",
|
||||
"commit",
|
||||
"consume",
|
||||
"create_block_reply_pipeline",
|
||||
"create_channel_trace",
|
||||
"get_backend",
|
||||
"parse_idempotency_key",
|
||||
"record_agent_duration_ms",
|
||||
"record_dispatch_duration_ms",
|
||||
"record_message",
|
||||
"record_rate_limit_reject",
|
||||
"registry",
|
||||
"release",
|
||||
"release_and_forget",
|
||||
"reset",
|
||||
"reset_sync",
|
||||
"set_inflight",
|
||||
"shutdown_media",
|
||||
]
|
||||
240
backend/package/yuxi/channel/message/block_reply_pipeline.py
Normal file
240
backend/package/yuxi/channel/message/block_reply_pipeline.py
Normal file
@ -0,0 +1,240 @@
|
||||
import asyncio
|
||||
import logging
|
||||
from collections.abc import Callable, Coroutine
|
||||
from typing import Any
|
||||
|
||||
from yuxi.channel.message.models import ChunkType, ReplyPayload, ReplyStage, StreamingChunk, UnifiedMessage
|
||||
from yuxi.channel.message.reply_dispatcher import ReplyDispatcher
|
||||
from yuxi.channel.streaming.block_chunker import BlockReplyCoalescer
|
||||
from yuxi.channel.streaming.models import BlockReplyCoalescing
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_DEFAULT_COALESCE_MIN_CHARS = 80
|
||||
_DEFAULT_COALESCE_MAX_CHARS = 2000
|
||||
_DEFAULT_COALESCE_IDLE_MS = 150.0
|
||||
_DEFAULT_TIMEOUT_MS = 10_000.0
|
||||
|
||||
|
||||
def _handle_pipeline_task_exception(task: asyncio.Task) -> None:
|
||||
try:
|
||||
task.result()
|
||||
except Exception:
|
||||
logger.exception("Unhandled exception in pipeline background task")
|
||||
|
||||
|
||||
class BlockReplyPipeline:
|
||||
def __init__(
|
||||
self,
|
||||
dispatcher: ReplyDispatcher,
|
||||
*,
|
||||
min_chars: int = _DEFAULT_COALESCE_MIN_CHARS,
|
||||
max_chars: int = _DEFAULT_COALESCE_MAX_CHARS,
|
||||
debounce_ms: float = _DEFAULT_COALESCE_IDLE_MS,
|
||||
timeout_ms: float = _DEFAULT_TIMEOUT_MS,
|
||||
):
|
||||
self._dispatcher = dispatcher
|
||||
self._timeout_ms = timeout_ms
|
||||
|
||||
self._sent_payload_keys: set[str] = set()
|
||||
self._sent_content_keys: set[str] = set()
|
||||
self._sent_media_urls: set[str] = set()
|
||||
self._pending_keys: set[str] = set()
|
||||
self._payload_seq: int = 0
|
||||
|
||||
self._streamed_text_fragments: list[str] = []
|
||||
self._aborted = False
|
||||
self._did_stream = False
|
||||
self._did_log_timeout = False
|
||||
|
||||
self._send_chain = asyncio.Lock()
|
||||
|
||||
coalesce_config = BlockReplyCoalescing(
|
||||
min_chars=max(1, min_chars),
|
||||
max_chars=max(max(1, min_chars), max_chars),
|
||||
idle_ms=int(debounce_ms),
|
||||
joiner="",
|
||||
)
|
||||
self._coalescer = BlockReplyCoalescer(
|
||||
config=coalesce_config,
|
||||
on_flush=self._on_coalesced_flush,
|
||||
is_stopped=lambda: self._aborted,
|
||||
)
|
||||
|
||||
async def enqueue(self, chunk: StreamingChunk) -> None:
|
||||
if self._aborted:
|
||||
return
|
||||
|
||||
if chunk.stage == ReplyStage.TOOL:
|
||||
self._coalescer.flush()
|
||||
await self._enqueue_tool(chunk)
|
||||
return
|
||||
|
||||
if chunk.stage == ReplyStage.BLOCK and chunk.chunk_type == ChunkType.TEXT_DELTA:
|
||||
if chunk.content:
|
||||
self._coalescer.enqueue(chunk.content)
|
||||
return
|
||||
|
||||
if chunk.stage == ReplyStage.FINAL:
|
||||
self._coalescer.flush()
|
||||
await self._dispatcher.enqueue(chunk)
|
||||
return
|
||||
|
||||
self._coalescer.flush()
|
||||
await self._dispatcher.enqueue(chunk)
|
||||
|
||||
def _on_coalesced_flush(self, text: str) -> None:
|
||||
payload = ReplyPayload(target_id="", content=text)
|
||||
task = asyncio.create_task(self._send_payload(payload))
|
||||
task.add_done_callback(_handle_pipeline_task_exception)
|
||||
|
||||
async def _enqueue_tool(self, chunk: StreamingChunk) -> None:
|
||||
payload = ReplyPayload(
|
||||
target_id="",
|
||||
content=chunk.content or "",
|
||||
)
|
||||
await self._send_payload(payload)
|
||||
|
||||
async def _send_payload(self, payload: ReplyPayload) -> None:
|
||||
if self._aborted:
|
||||
return
|
||||
|
||||
self._payload_seq += 1
|
||||
dedup_key = f"{self._payload_seq}|{payload.payload_key}"
|
||||
content_key = payload.content_key
|
||||
|
||||
if dedup_key in self._sent_payload_keys or dedup_key in self._pending_keys:
|
||||
return
|
||||
|
||||
self._pending_keys.add(dedup_key)
|
||||
|
||||
try:
|
||||
async with self._send_chain:
|
||||
if self._aborted:
|
||||
return
|
||||
|
||||
chunk = StreamingChunk(
|
||||
stage=ReplyStage.BLOCK,
|
||||
chunk_type=ChunkType.TEXT_DELTA,
|
||||
content=payload.content,
|
||||
)
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
self._dispatcher.enqueue(chunk),
|
||||
timeout=self._timeout_ms / 1000,
|
||||
)
|
||||
except TimeoutError:
|
||||
self._aborted = True
|
||||
if not self._did_log_timeout:
|
||||
self._did_log_timeout = True
|
||||
logger.warning(
|
||||
"Block reply delivery timed out after %.0fms; "
|
||||
"aborting remaining replies to preserve ordering",
|
||||
self._timeout_ms,
|
||||
)
|
||||
return
|
||||
|
||||
self._sent_payload_keys.add(dedup_key)
|
||||
self._sent_content_keys.add(content_key)
|
||||
for url in payload.media_urls:
|
||||
self._sent_media_urls.add(url)
|
||||
if not payload.media_urls and payload.content.strip():
|
||||
self._streamed_text_fragments.append(payload.content.strip())
|
||||
self._did_stream = True
|
||||
finally:
|
||||
self._pending_keys.discard(dedup_key)
|
||||
|
||||
async def flush(self, *, force: bool = False) -> None:
|
||||
if force or self._coalescer.has_buffered():
|
||||
self._coalescer.flush()
|
||||
async with self._send_chain:
|
||||
pass
|
||||
|
||||
def stop(self) -> None:
|
||||
self._coalescer.stop()
|
||||
|
||||
def has_buffered(self) -> bool:
|
||||
return self._coalescer.has_buffered()
|
||||
|
||||
def did_stream(self) -> bool:
|
||||
return self._did_stream
|
||||
|
||||
def is_aborted(self) -> bool:
|
||||
return self._aborted
|
||||
|
||||
def get_sent_media_urls(self) -> list[str]:
|
||||
return list(self._sent_media_urls)
|
||||
|
||||
def has_sent_payload(self, content: str) -> bool:
|
||||
stripped = content.strip()
|
||||
if stripped in self._sent_content_keys:
|
||||
return True
|
||||
if not self._did_stream or not self._streamed_text_fragments:
|
||||
return False
|
||||
|
||||
def _normalize(s: str) -> str:
|
||||
return "".join(s.split())
|
||||
|
||||
return _normalize("".join(self._streamed_text_fragments)) == _normalize(stripped)
|
||||
|
||||
def track_media_sent(self, media_url: str) -> None:
|
||||
self._sent_media_urls.add(media_url)
|
||||
|
||||
def has_sent_media(self, media_url: str) -> bool:
|
||||
return media_url in self._sent_media_urls
|
||||
|
||||
@property
|
||||
def streamed_text(self) -> str:
|
||||
return "".join(self._streamed_text_fragments)
|
||||
|
||||
@property
|
||||
def sent_payload_count(self) -> int:
|
||||
return len(self._sent_payload_keys)
|
||||
|
||||
@property
|
||||
def sent_content_count(self) -> int:
|
||||
return len(self._sent_content_keys)
|
||||
|
||||
@property
|
||||
def sent_media_count(self) -> int:
|
||||
return len(self._sent_media_urls)
|
||||
|
||||
async def abort(self) -> None:
|
||||
self._aborted = True
|
||||
await self._dispatcher.abort()
|
||||
|
||||
async def wait_idle(self, timeout_ms: float | None = None) -> None:
|
||||
await self._dispatcher.wait_idle(timeout_ms)
|
||||
|
||||
|
||||
def create_block_reply_pipeline(
|
||||
send_fn: Callable[..., Coroutine[Any, Any, str | None]],
|
||||
msg: UnifiedMessage,
|
||||
*,
|
||||
response_prefix: str = "",
|
||||
human_delay: tuple[float, float] | None = None,
|
||||
min_chars: int = _DEFAULT_COALESCE_MIN_CHARS,
|
||||
max_chars: int = _DEFAULT_COALESCE_MAX_CHARS,
|
||||
debounce_ms: float = _DEFAULT_COALESCE_IDLE_MS,
|
||||
timeout_ms: float = _DEFAULT_TIMEOUT_MS,
|
||||
on_idle: Callable[[], None] | None = None,
|
||||
on_error: Callable[[Exception], None] | None = None,
|
||||
before_deliver: Callable[[str], Coroutine[Any, Any, str | None]] | None = None,
|
||||
) -> BlockReplyPipeline:
|
||||
dispatcher = ReplyDispatcher(
|
||||
send_fn=send_fn,
|
||||
msg=msg,
|
||||
response_prefix=response_prefix,
|
||||
human_delay=human_delay,
|
||||
timeout_ms=timeout_ms,
|
||||
on_idle=on_idle,
|
||||
on_error=on_error,
|
||||
before_deliver=before_deliver,
|
||||
)
|
||||
return BlockReplyPipeline(
|
||||
dispatcher,
|
||||
min_chars=min_chars,
|
||||
max_chars=max_chars,
|
||||
debounce_ms=debounce_ms,
|
||||
timeout_ms=timeout_ms,
|
||||
)
|
||||
372
backend/package/yuxi/channel/message/bridge.py
Normal file
372
backend/package/yuxi/channel/message/bridge.py
Normal file
@ -0,0 +1,372 @@
|
||||
import asyncio
|
||||
import contextlib
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import AsyncIterator, Callable
|
||||
|
||||
from yuxi.channel.message.models import (
|
||||
ChunkType,
|
||||
DispatchResult,
|
||||
ReplyStage,
|
||||
StreamCancellationToken,
|
||||
StreamCancelledError,
|
||||
StreamingChunk,
|
||||
StreamResult,
|
||||
StreamStatus,
|
||||
StreamTimeoutError,
|
||||
UnifiedMessage,
|
||||
)
|
||||
from yuxi.channel.protocols import AgentPromptContext, AgentPromptProtocol, AgentToolProtocol
|
||||
from yuxi.channel.plugins.registry import ChannelPluginRegistry
|
||||
from yuxi.channel.security.external_content import sanitize_external_content
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_STREAMING_TEXT_STATUSES = frozenset({"init", "loading", "agent_state", "streaming"})
|
||||
_STREAMING_REASONING_STATUSES = frozenset({"thinking"})
|
||||
_TOOL_STATUSES = frozenset({"tool_call", "tool_result"})
|
||||
_TERMINAL_STATUSES = frozenset({"finished", "error", "interrupted", "warning"})
|
||||
_INTERACTIVE_STATUSES = frozenset({"ask_user_question_required"})
|
||||
|
||||
_DEFAULT_TIMEOUT_SECONDS = 300.0
|
||||
|
||||
|
||||
class AgentBridge:
|
||||
def __init__(
|
||||
self,
|
||||
stream_fn: Callable[..., AsyncIterator[bytes]],
|
||||
channel_manager=None,
|
||||
):
|
||||
self._stream_fn = stream_fn
|
||||
self._channel_manager = channel_manager
|
||||
|
||||
def collect_channel_tools(self, channel_type: str) -> list[dict]:
|
||||
plugin = ChannelPluginRegistry.get(channel_type)
|
||||
if plugin is None or not isinstance(plugin, AgentToolProtocol):
|
||||
return []
|
||||
tools = plugin.get_agent_tools()
|
||||
return [t.to_openai_schema() for t in tools]
|
||||
|
||||
async def execute_channel_tool(
|
||||
self,
|
||||
channel_type: str,
|
||||
tool_name: str,
|
||||
params: dict,
|
||||
context: dict,
|
||||
) -> dict:
|
||||
plugin = ChannelPluginRegistry.get(channel_type)
|
||||
if plugin is None or not isinstance(plugin, AgentToolProtocol):
|
||||
return {"success": False, "error": f"Channel {channel_type} does not support agent tools"}
|
||||
|
||||
try:
|
||||
return await plugin.execute_agent_tool(tool_name, params, context)
|
||||
except Exception as e:
|
||||
logger.exception("Channel tool execution failed: %s.%s", channel_type, tool_name)
|
||||
return {"success": False, "error": str(e)}
|
||||
|
||||
# ── 构建 meta ──────────────────────────────────────────
|
||||
|
||||
def _build_meta(self, msg: UnifiedMessage) -> tuple[dict, str]:
|
||||
meta = {
|
||||
"source": "channel",
|
||||
"channel_type": msg.channel_type,
|
||||
"account_id": msg.account_id,
|
||||
"peer_id": msg.sender.id,
|
||||
}
|
||||
|
||||
plugin = ChannelPluginRegistry.get(msg.channel_type)
|
||||
|
||||
safe_content, _detection = sanitize_external_content(msg.content)
|
||||
query = safe_content
|
||||
|
||||
if isinstance(plugin, AgentPromptProtocol):
|
||||
context = AgentPromptContext(
|
||||
channel_type=msg.channel_type,
|
||||
account_id=msg.account_id,
|
||||
peer_id=msg.sender.id,
|
||||
peer_name=msg.sender.display_name,
|
||||
group_id=msg.group.id if msg.group else None,
|
||||
group_name=msg.group.name if msg.group else None,
|
||||
thread_id=msg.group.thread_id if msg.group else None,
|
||||
)
|
||||
system_prompt = plugin.build_system_prompt(context)
|
||||
if system_prompt:
|
||||
meta["channel_system_prompt"] = system_prompt
|
||||
context_note = plugin.build_context_note(context)
|
||||
if context_note:
|
||||
query = f"{context_note}\n{safe_content}"
|
||||
|
||||
channel_tools = self.collect_channel_tools(msg.channel_type)
|
||||
if channel_tools:
|
||||
meta["channel_tools"] = channel_tools
|
||||
|
||||
return meta, query
|
||||
|
||||
# ── invoke (原始字节流) ─────────────────────────────────
|
||||
|
||||
async def invoke(
|
||||
self,
|
||||
msg: UnifiedMessage,
|
||||
dispatch_result: DispatchResult,
|
||||
agent_config_id: int,
|
||||
current_user,
|
||||
db,
|
||||
) -> AsyncIterator[bytes]:
|
||||
if self._channel_manager:
|
||||
self._channel_manager.inc_active_runs(msg.channel_type, msg.account_id)
|
||||
|
||||
try:
|
||||
meta, query = self._build_meta(msg)
|
||||
|
||||
async for chunk in self._stream_fn(
|
||||
query=query,
|
||||
agent_config_id=agent_config_id,
|
||||
thread_id=dispatch_result.thread_id,
|
||||
meta=meta,
|
||||
image_content=msg.image_base64,
|
||||
current_user=current_user,
|
||||
db=db,
|
||||
):
|
||||
yield chunk
|
||||
finally:
|
||||
if self._channel_manager:
|
||||
self._channel_manager.dec_active_runs(msg.channel_type, msg.account_id)
|
||||
|
||||
# ── invoke_stream (结构化流) ────────────────────────────
|
||||
|
||||
async def invoke_stream(
|
||||
self,
|
||||
msg: UnifiedMessage,
|
||||
dispatch_result: DispatchResult,
|
||||
agent_config_id: int,
|
||||
current_user,
|
||||
db,
|
||||
*,
|
||||
cancel_token: StreamCancellationToken | None = None,
|
||||
on_chunk: Callable[[StreamingChunk], None] | None = None,
|
||||
timeout: float = _DEFAULT_TIMEOUT_SECONDS,
|
||||
) -> AsyncIterator[StreamingChunk]:
|
||||
if self._channel_manager:
|
||||
self._channel_manager.inc_active_runs(msg.channel_type, msg.account_id)
|
||||
|
||||
try:
|
||||
meta, query = self._build_meta(msg)
|
||||
|
||||
async def _stream_all() -> AsyncIterator[bytes]:
|
||||
async for raw in self._stream_fn(
|
||||
query=query,
|
||||
agent_config_id=agent_config_id,
|
||||
thread_id=dispatch_result.thread_id,
|
||||
meta=meta,
|
||||
image_content=msg.image_base64,
|
||||
current_user=current_user,
|
||||
db=db,
|
||||
):
|
||||
yield raw
|
||||
|
||||
stream_iter = _stream_all()
|
||||
|
||||
async with contextlib.aclosing(stream_iter):
|
||||
while True:
|
||||
if cancel_token is not None and cancel_token.is_cancelled:
|
||||
raise StreamCancelledError("stream cancelled by cancel token")
|
||||
|
||||
try:
|
||||
raw_chunk = await asyncio.wait_for(
|
||||
stream_iter.__anext__(),
|
||||
timeout=timeout,
|
||||
)
|
||||
except TimeoutError:
|
||||
raise StreamTimeoutError(
|
||||
f"No chunk received within {timeout}s"
|
||||
)
|
||||
except StopAsyncIteration:
|
||||
break
|
||||
|
||||
parsed = self._parse_chunk(raw_chunk)
|
||||
if parsed is None:
|
||||
continue
|
||||
|
||||
if on_chunk is not None:
|
||||
on_chunk(parsed)
|
||||
|
||||
yield parsed
|
||||
|
||||
yield StreamingChunk(
|
||||
stage=ReplyStage.FINAL,
|
||||
chunk_type=ChunkType.STATUS,
|
||||
content="finished",
|
||||
)
|
||||
finally:
|
||||
if self._channel_manager:
|
||||
self._channel_manager.dec_active_runs(msg.channel_type, msg.account_id)
|
||||
|
||||
# ── invoke_with_result (结构化结果) ─────────────────────
|
||||
|
||||
async def invoke_with_result(
|
||||
self,
|
||||
msg: UnifiedMessage,
|
||||
dispatch_result: DispatchResult,
|
||||
agent_config_id: int,
|
||||
current_user,
|
||||
db,
|
||||
*,
|
||||
cancel_token: StreamCancellationToken | None = None,
|
||||
timeout: float = _DEFAULT_TIMEOUT_SECONDS,
|
||||
) -> StreamResult:
|
||||
result = StreamResult(start_time=time.monotonic())
|
||||
accumulated: list[str] = []
|
||||
|
||||
try:
|
||||
async for chunk in self.invoke_stream(
|
||||
msg,
|
||||
dispatch_result,
|
||||
agent_config_id,
|
||||
current_user,
|
||||
db,
|
||||
cancel_token=cancel_token,
|
||||
timeout=timeout,
|
||||
):
|
||||
result.chunk_count += 1
|
||||
|
||||
if chunk.chunk_type == ChunkType.TEXT_DELTA:
|
||||
result.text_chunk_count += 1
|
||||
if chunk.content:
|
||||
accumulated.append(chunk.content)
|
||||
elif chunk.chunk_type == ChunkType.REASONING:
|
||||
result.reasoning_chunk_count += 1
|
||||
elif chunk.chunk_type == ChunkType.TOOL_CALL:
|
||||
result.tool_chunk_count += 1
|
||||
elif chunk.chunk_type == ChunkType.TOOL_RESULT:
|
||||
result.tool_chunk_count += 1
|
||||
elif chunk.chunk_type == ChunkType.HEARTBEAT:
|
||||
result.heartbeat_count += 1
|
||||
elif chunk.chunk_type == ChunkType.CITATION:
|
||||
result.citation_count += 1
|
||||
elif chunk.chunk_type == ChunkType.ERROR:
|
||||
result.error_count += 1
|
||||
if chunk.content and not result.error_message:
|
||||
result.error_message = chunk.content
|
||||
|
||||
if chunk.metadata and "token_usage" in chunk.metadata:
|
||||
result.token_usage = chunk.metadata["token_usage"]
|
||||
|
||||
result.accumulated_text = "".join(accumulated)
|
||||
result.status = StreamStatus.COMPLETED
|
||||
|
||||
except StreamCancelledError:
|
||||
result.status = StreamStatus.CANCELLED
|
||||
result.accumulated_text = "".join(accumulated)
|
||||
result.error_message = "stream cancelled"
|
||||
logger.debug("Agent stream cancelled for channel=%s peer=%s", msg.channel_type, msg.sender.id)
|
||||
except StreamTimeoutError:
|
||||
result.status = StreamStatus.TIMEOUT
|
||||
result.accumulated_text = "".join(accumulated)
|
||||
result.error_message = "stream timeout"
|
||||
logger.warning("Agent stream timeout for channel=%s peer=%s", msg.channel_type, msg.sender.id)
|
||||
except Exception as e:
|
||||
result.status = StreamStatus.ERROR
|
||||
result.accumulated_text = "".join(accumulated)
|
||||
result.error_message = f"{type(e).__name__}: {e}"
|
||||
logger.exception("Agent stream error for channel=%s peer=%s", msg.channel_type, msg.sender.id)
|
||||
finally:
|
||||
result.end_time = time.monotonic()
|
||||
|
||||
return result
|
||||
|
||||
# ── invoke_and_collect (兼容方法) ────────────────────────
|
||||
|
||||
async def invoke_and_collect(
|
||||
self,
|
||||
msg: UnifiedMessage,
|
||||
dispatch_result: DispatchResult,
|
||||
agent_config_id: int,
|
||||
current_user,
|
||||
db,
|
||||
) -> str | None:
|
||||
result = await self.invoke_with_result(msg, dispatch_result, agent_config_id, current_user, db)
|
||||
return result.accumulated_text or None
|
||||
|
||||
# ── chunk 解析 ─────────────────────────────────────────
|
||||
|
||||
def _parse_chunk(self, raw: bytes) -> StreamingChunk | None:
|
||||
try:
|
||||
data = json.loads(raw.decode("utf-8"))
|
||||
except (json.JSONDecodeError, UnicodeDecodeError):
|
||||
logger.debug("Failed to parse stream chunk (len=%d): %s", len(raw), raw[:200])
|
||||
return None
|
||||
|
||||
status = data.get("status", "")
|
||||
content = data.get("response", "")
|
||||
|
||||
# 流式文本块
|
||||
if status in _STREAMING_TEXT_STATUSES:
|
||||
return StreamingChunk(
|
||||
stage=ReplyStage.BLOCK,
|
||||
chunk_type=ChunkType.TEXT_DELTA,
|
||||
content=content,
|
||||
)
|
||||
|
||||
# 推理/思考块
|
||||
if status in _STREAMING_REASONING_STATUSES:
|
||||
return StreamingChunk(
|
||||
stage=ReplyStage.TOOL,
|
||||
chunk_type=ChunkType.REASONING,
|
||||
content=content,
|
||||
)
|
||||
|
||||
# 工具调用
|
||||
if status == "tool_call":
|
||||
tool_name = data.get("tool_name", "")
|
||||
tool_input = data.get("tool_input", {})
|
||||
return StreamingChunk(
|
||||
stage=ReplyStage.TOOL,
|
||||
chunk_type=ChunkType.TOOL_CALL,
|
||||
content=content,
|
||||
tool_name=tool_name,
|
||||
tool_input=tool_input if isinstance(tool_input, dict) else {},
|
||||
)
|
||||
|
||||
# 工具结果
|
||||
if status == "tool_result":
|
||||
tool_name = data.get("tool_name", "")
|
||||
tool_output = data.get("tool_output", content)
|
||||
return StreamingChunk(
|
||||
stage=ReplyStage.TOOL,
|
||||
chunk_type=ChunkType.TOOL_RESULT,
|
||||
content=content,
|
||||
tool_name=tool_name,
|
||||
tool_output=tool_output,
|
||||
)
|
||||
|
||||
# 心跳
|
||||
if status == "heartbeat":
|
||||
return StreamingChunk(
|
||||
stage=ReplyStage.BLOCK,
|
||||
chunk_type=ChunkType.HEARTBEAT,
|
||||
content=content,
|
||||
)
|
||||
|
||||
# 引用
|
||||
if status == "citation":
|
||||
return StreamingChunk(
|
||||
stage=ReplyStage.BLOCK,
|
||||
chunk_type=ChunkType.CITATION,
|
||||
content=content,
|
||||
)
|
||||
|
||||
# 错误
|
||||
if status == "error":
|
||||
return StreamingChunk(
|
||||
stage=ReplyStage.FINAL,
|
||||
chunk_type=ChunkType.ERROR,
|
||||
content=content,
|
||||
)
|
||||
|
||||
# 终止态 / 交互态 — 不产出 chunk
|
||||
if status in _TERMINAL_STATUSES | _INTERACTIVE_STATUSES:
|
||||
return None
|
||||
|
||||
logger.debug("Unknown chunk status: %s", status)
|
||||
return None
|
||||
112
backend/package/yuxi/channel/message/circuit_breaker.py
Normal file
112
backend/package/yuxi/channel/message/circuit_breaker.py
Normal file
@ -0,0 +1,112 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from enum import StrEnum
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class CircuitState(StrEnum):
|
||||
CLOSED = "closed"
|
||||
OPEN = "open"
|
||||
HALF_OPEN = "half_open"
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class CircuitBreakerConfig:
|
||||
failure_threshold: int = 5
|
||||
recovery_timeout_sec: float = 30.0
|
||||
half_open_max_requests: int = 1
|
||||
consecutive_successes_to_close: int = 2
|
||||
|
||||
|
||||
@dataclass
|
||||
class CircuitBreaker:
|
||||
"""简单熔断器。
|
||||
|
||||
状态机:
|
||||
- CLOSED: 正常通行,累计失败计数。
|
||||
- OPEN: 拒绝请求,持续 recovery_timeout_sec 后转 HALF_OPEN。
|
||||
- HALF_OPEN: 允许少量探测请求。连续成功则转 CLOSED,任何失败则回 OPEN。
|
||||
"""
|
||||
|
||||
name: str
|
||||
config: CircuitBreakerConfig = field(default_factory=CircuitBreakerConfig)
|
||||
|
||||
_state: CircuitState = CircuitState.CLOSED
|
||||
_failure_count: int = 0
|
||||
_success_count: int = 0
|
||||
_last_failure_time: float = 0.0
|
||||
_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
|
||||
|
||||
async def call(self, coro) -> object:
|
||||
if not await self._allow_request():
|
||||
raise CircuitBreakerOpenError(
|
||||
f"Circuit breaker '{self.name}' is OPEN"
|
||||
)
|
||||
try:
|
||||
result = await coro
|
||||
except Exception:
|
||||
await self._record_failure()
|
||||
raise
|
||||
else:
|
||||
await self._record_success()
|
||||
return result
|
||||
|
||||
async def _allow_request(self) -> bool:
|
||||
async with self._lock:
|
||||
if self._state == CircuitState.CLOSED:
|
||||
return True
|
||||
if self._state == CircuitState.OPEN:
|
||||
if time.monotonic() - self._last_failure_time >= self.config.recovery_timeout_sec:
|
||||
self._state = CircuitState.HALF_OPEN
|
||||
self._success_count = 0
|
||||
logger.info("Circuit '%s' transitioning OPEN -> HALF_OPEN", self.name)
|
||||
return True
|
||||
return False
|
||||
if self._state == CircuitState.HALF_OPEN:
|
||||
return True
|
||||
return True
|
||||
|
||||
async def _record_success(self) -> None:
|
||||
async with self._lock:
|
||||
if self._state == CircuitState.HALF_OPEN:
|
||||
self._success_count += 1
|
||||
if self._success_count >= self.config.consecutive_successes_to_close:
|
||||
self._state = CircuitState.CLOSED
|
||||
self._failure_count = 0
|
||||
self._success_count = 0
|
||||
logger.info("Circuit '%s' transitioning HALF_OPEN -> CLOSED", self.name)
|
||||
|
||||
async def _record_failure(self) -> None:
|
||||
async with self._lock:
|
||||
self._failure_count += 1
|
||||
self._last_failure_time = time.monotonic()
|
||||
if self._state == CircuitState.HALF_OPEN:
|
||||
self._state = CircuitState.OPEN
|
||||
logger.warning(
|
||||
"Circuit '%s' transitioning HALF_OPEN -> OPEN (probe failed)",
|
||||
self.name,
|
||||
)
|
||||
elif self._state == CircuitState.CLOSED and self._failure_count >= self.config.failure_threshold:
|
||||
self._state = CircuitState.OPEN
|
||||
logger.warning(
|
||||
"Circuit '%s' transitioning CLOSED -> OPEN after %d failures",
|
||||
self.name,
|
||||
self._failure_count,
|
||||
)
|
||||
|
||||
@property
|
||||
def state(self) -> CircuitState:
|
||||
return self._state
|
||||
|
||||
async def reset(self) -> None:
|
||||
async with self._lock:
|
||||
self._state = CircuitState.CLOSED
|
||||
self._failure_count = 0
|
||||
self._success_count = 0
|
||||
|
||||
|
||||
class CircuitBreakerOpenError(Exception):
|
||||
pass
|
||||
53
backend/package/yuxi/channel/message/conversation_fence.py
Normal file
53
backend/package/yuxi/channel/message/conversation_fence.py
Normal file
@ -0,0 +1,53 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
|
||||
from yuxi.channel.message.models import UnifiedMessage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_DEFAULT_FENCE_TTL = 300.0
|
||||
|
||||
|
||||
class ConversationFence:
|
||||
def __init__(self, ttl: float = _DEFAULT_FENCE_TTL):
|
||||
self._ttl = ttl
|
||||
self._versions: dict[str, int] = {}
|
||||
self._version_ts: dict[str, float] = {}
|
||||
self._locks: dict[str, asyncio.Lock] = {}
|
||||
self._abort_events: dict[str, asyncio.Event] = {}
|
||||
|
||||
@staticmethod
|
||||
def key_for(msg: UnifiedMessage) -> str:
|
||||
if msg.group and msg.group.id:
|
||||
return f"{msg.channel_type}:{msg.account_id}:group:{msg.group.id}"
|
||||
return f"{msg.channel_type}:{msg.account_id}:dm:{msg.sender.id}"
|
||||
|
||||
def enter(self, key: str) -> tuple[int, asyncio.Event]:
|
||||
now = time.monotonic()
|
||||
self._gc(now)
|
||||
self._versions[key] = self._versions.get(key, 0) + 1
|
||||
self._version_ts[key] = now
|
||||
|
||||
old_event = self._abort_events.get(key)
|
||||
if old_event is not None and not old_event.is_set():
|
||||
old_event.set()
|
||||
logger.debug("Foreground fence: aborting previous run for %s", key)
|
||||
|
||||
new_event = asyncio.Event()
|
||||
self._abort_events[key] = new_event
|
||||
return self._versions[key], new_event
|
||||
|
||||
def lock_for(self, key: str) -> asyncio.Lock:
|
||||
return self._locks.setdefault(key, asyncio.Lock())
|
||||
|
||||
def current_version(self, key: str) -> int:
|
||||
return self._versions.get(key, 0)
|
||||
|
||||
def _gc(self, now: float) -> None:
|
||||
expired = [k for k, ts in self._version_ts.items() if now - ts > self._ttl]
|
||||
for k in expired:
|
||||
self._versions.pop(k, None)
|
||||
self._version_ts.pop(k, None)
|
||||
self._locks.pop(k, None)
|
||||
self._abort_events.pop(k, None)
|
||||
407
backend/package/yuxi/channel/message/dispatch.py
Normal file
407
backend/package/yuxi/channel/message/dispatch.py
Normal file
@ -0,0 +1,407 @@
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from yuxi.channel.message.models import DispatchResult, PeerKind, UnifiedMessage
|
||||
from yuxi.channel.protocols import (
|
||||
ApprovalAction,
|
||||
ApprovalDecision,
|
||||
ConfigProtocol,
|
||||
MessagingProtocol,
|
||||
SessionResolution,
|
||||
)
|
||||
from yuxi.channel.plugins.registry import ChannelPluginRegistry
|
||||
from yuxi.channel.routing.matcher import MessageContext, RouteMatcher, VerificationContext
|
||||
from yuxi.channel.routing.models import RouteBinding
|
||||
from yuxi.channel.routing.session_key import SessionKeyBuilder, session_key_to_thread_id
|
||||
from yuxi.channel.security.allowlist import AllowlistChecker
|
||||
from yuxi.channel.security.pairing import PairingManager
|
||||
from yuxi.repositories.channel_thread_mapping_repo import ChannelThreadMappingRepository
|
||||
from yuxi.repositories.channel_user_mapping_repo import ChannelUserMappingRepository
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from yuxi.channel.message.langfuse_trace import ChannelTrace
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
BeforeDispatchHook = Callable[[UnifiedMessage], Awaitable[bool | None]]
|
||||
AfterDispatchHook = Callable[[UnifiedMessage, DispatchResult], Awaitable[None]]
|
||||
|
||||
|
||||
class MessageDispatcher:
|
||||
def __init__(
|
||||
self,
|
||||
bindings: list[RouteBinding],
|
||||
allowlist: AllowlistChecker,
|
||||
pairing: PairingManager,
|
||||
*,
|
||||
matcher: RouteMatcher | None = None,
|
||||
session_key_builder: SessionKeyBuilder | None = None,
|
||||
approval_engine=None,
|
||||
user_mapping_repo: ChannelUserMappingRepository | None = None,
|
||||
thread_mapping_repo: ChannelThreadMappingRepository | None = None,
|
||||
on_before_dispatch: BeforeDispatchHook | None = None,
|
||||
on_after_dispatch: AfterDispatchHook | None = None,
|
||||
):
|
||||
self._bindings = bindings
|
||||
self._matcher = matcher
|
||||
self._allowlist = allowlist
|
||||
self._pairing = pairing
|
||||
self._session_key_builder = session_key_builder or SessionKeyBuilder()
|
||||
self._approval_engine = approval_engine
|
||||
self._user_mapping_repo = user_mapping_repo
|
||||
self._thread_mapping_repo = thread_mapping_repo
|
||||
self._on_before_dispatch = on_before_dispatch
|
||||
self._on_after_dispatch = on_after_dispatch
|
||||
|
||||
async def dispatch(self, msg: UnifiedMessage, *, trace: "ChannelTrace | None" = None) -> DispatchResult:
|
||||
command_result = await self._try_handle_command(msg)
|
||||
if command_result is not None:
|
||||
return command_result
|
||||
|
||||
try:
|
||||
return await self._dispatch_inner(msg, trace=trace)
|
||||
except Exception:
|
||||
logger.exception("Dispatch failed for message: %s", msg.msg_id)
|
||||
return DispatchResult(success=False, error="dispatch_error")
|
||||
|
||||
@staticmethod
|
||||
def _plugin_has_commands(plugin) -> bool:
|
||||
return callable(getattr(plugin, "get_commands", None)) and callable(getattr(plugin, "handle_command", None))
|
||||
|
||||
async def _try_handle_command(self, msg: UnifiedMessage) -> DispatchResult | None:
|
||||
text = (msg.content or "").strip()
|
||||
if not text.startswith("/"):
|
||||
return None
|
||||
|
||||
parts = text[1:].split()
|
||||
if not parts:
|
||||
return None
|
||||
|
||||
command_name = parts[0].lower()
|
||||
args = parts[1:]
|
||||
|
||||
plugin = ChannelPluginRegistry.get(msg.channel_type)
|
||||
if plugin is None or not self._plugin_has_commands(plugin):
|
||||
return None
|
||||
|
||||
config = {}
|
||||
if isinstance(plugin, ConfigProtocol):
|
||||
try:
|
||||
account_cfg = await plugin.resolve_account(msg.account_id)
|
||||
config = account_cfg if account_cfg else {}
|
||||
except Exception:
|
||||
logger.debug("Failed to resolve account config for command: %s", command_name)
|
||||
|
||||
try:
|
||||
response = await plugin.handle_command(
|
||||
config,
|
||||
command_name,
|
||||
args,
|
||||
msg,
|
||||
None,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Command handling failed: %s", command_name)
|
||||
return DispatchResult(
|
||||
success=True,
|
||||
handled_by="command",
|
||||
command_response="命令执行失败,请稍后重试。",
|
||||
)
|
||||
|
||||
return DispatchResult(
|
||||
success=True,
|
||||
handled_by="command",
|
||||
command_response=response,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _plugin_has_approval(plugin) -> bool:
|
||||
return (
|
||||
callable(getattr(plugin, "check_approval_required", None))
|
||||
and callable(getattr(plugin, "create_approval_request", None))
|
||||
and callable(getattr(plugin, "check_approval_status", None))
|
||||
)
|
||||
|
||||
async def _should_check_approval(self, msg: UnifiedMessage) -> bool:
|
||||
if self._approval_engine is None:
|
||||
return False
|
||||
plugin = ChannelPluginRegistry.get(msg.channel_type)
|
||||
return plugin is not None and self._plugin_has_approval(plugin)
|
||||
|
||||
async def _dispatch_inner(self, msg: UnifiedMessage, *, trace: "ChannelTrace | None" = None) -> DispatchResult:
|
||||
plugin = ChannelPluginRegistry.get(msg.channel_type)
|
||||
|
||||
sec_span_id = None
|
||||
if trace:
|
||||
sec_span_id = await trace.add_span("security_admission", metadata={"peer_kind": msg.sender.kind.value})
|
||||
|
||||
if msg.sender.kind == PeerKind.DIRECT:
|
||||
try:
|
||||
dm_result = self._allowlist.check_dm(msg.sender.id)
|
||||
if not dm_result.allowed:
|
||||
if dm_result.pairing_required:
|
||||
req = await self._pairing.upsert_code(msg.channel_type, msg.sender.id, msg.account_id)
|
||||
if req is not None:
|
||||
if sec_span_id:
|
||||
await trace.end_span(sec_span_id, level="ERROR", status_message="pairing_required")
|
||||
return DispatchResult(success=False, error=f"pairing_required:{req.code}")
|
||||
if sec_span_id:
|
||||
await trace.end_span(sec_span_id, level="ERROR", status_message="pairing_check_error")
|
||||
return DispatchResult(success=False, error="pairing_check_error")
|
||||
logger.info(
|
||||
"Sender not in DM allowlist: %s/%s",
|
||||
msg.channel_type,
|
||||
msg.sender.id,
|
||||
)
|
||||
if sec_span_id:
|
||||
await trace.end_span(sec_span_id, level="ERROR", status_message="not_in_allowlist")
|
||||
return DispatchResult(success=False, error="not_in_allowlist")
|
||||
except Exception:
|
||||
logger.exception("DM allowlist check failed, blocking message")
|
||||
if sec_span_id:
|
||||
await trace.end_span(sec_span_id, level="ERROR", status_message="allowlist_check_error")
|
||||
return DispatchResult(success=False, error="allowlist_check_error")
|
||||
auth_scope = "dm"
|
||||
auth_verified_by = "AllowlistChecker.dm"
|
||||
else:
|
||||
try:
|
||||
if msg.group and msg.group.id:
|
||||
group_result = self._allowlist.check_group(msg.group.id)
|
||||
if not group_result.allowed:
|
||||
logger.info(
|
||||
"Group not in allowlist: %s/%s",
|
||||
msg.channel_type,
|
||||
msg.group.id,
|
||||
)
|
||||
if sec_span_id:
|
||||
await trace.end_span(sec_span_id, level="ERROR", status_message="not_in_allowlist")
|
||||
return DispatchResult(success=False, error="not_in_allowlist")
|
||||
from_result = self._allowlist.check_group_allow_from(msg.sender.id)
|
||||
if not from_result.allowed:
|
||||
logger.info(
|
||||
"Sender not in group allow_from: %s/%s",
|
||||
msg.channel_type,
|
||||
msg.sender.id,
|
||||
)
|
||||
if sec_span_id:
|
||||
await trace.end_span(sec_span_id, level="ERROR", status_message="not_in_allowlist")
|
||||
return DispatchResult(success=False, error="not_in_allowlist")
|
||||
except Exception:
|
||||
logger.exception("Group allowlist check failed, blocking message")
|
||||
if sec_span_id:
|
||||
await trace.end_span(sec_span_id, level="ERROR", status_message="allowlist_check_error")
|
||||
return DispatchResult(success=False, error="allowlist_check_error")
|
||||
auth_scope = "group"
|
||||
auth_verified_by = "AllowlistChecker.group"
|
||||
|
||||
if sec_span_id:
|
||||
await trace.end_span(sec_span_id, status_message="ok")
|
||||
|
||||
if self._matcher is None:
|
||||
self._matcher = RouteMatcher()
|
||||
await self._matcher.set_bindings(self._bindings)
|
||||
|
||||
explicit_target = None
|
||||
if isinstance(plugin, MessagingProtocol):
|
||||
session = plugin.resolve_session(msg)
|
||||
explicit_target = plugin.parse_explicit_target(msg.content)
|
||||
else:
|
||||
if msg.sender.kind == PeerKind.DIRECT:
|
||||
session = SessionResolution(kind="direct", conversation_id=msg.sender.id)
|
||||
else:
|
||||
gid = msg.group.id if msg.group else "unknown"
|
||||
session = SessionResolution(kind="group", conversation_id=gid)
|
||||
|
||||
try:
|
||||
peer_kind = PeerKind(session.kind)
|
||||
except ValueError:
|
||||
peer_kind = msg.sender.kind
|
||||
|
||||
ctx = MessageContext(
|
||||
channel=msg.channel_type,
|
||||
account_id=msg.account_id,
|
||||
peer_kind=peer_kind,
|
||||
peer_id=msg.sender.id,
|
||||
channel_config_id=msg.channel_config_id,
|
||||
member_role_ids=msg.member_role_ids,
|
||||
)
|
||||
ctx._verification_ctx = VerificationContext(
|
||||
verified_by=auth_verified_by,
|
||||
scope=auth_scope,
|
||||
timestamp=time.time(),
|
||||
)
|
||||
if msg.group:
|
||||
ctx.guild_id = msg.group.guild_id
|
||||
ctx.team_id = msg.group.team_id
|
||||
ctx.roles = msg.group.roles
|
||||
|
||||
if msg.group.thread_id or msg.message_thread_id:
|
||||
if msg.group.route_peer_kind:
|
||||
try:
|
||||
ctx.parent_peer_kind = PeerKind(msg.group.route_peer_kind)
|
||||
except ValueError:
|
||||
pass
|
||||
ctx.parent_peer_id = msg.group.route_peer_id or msg.group.id
|
||||
|
||||
route_span_id = None
|
||||
if trace:
|
||||
route_span_id = await trace.add_span("route_resolution")
|
||||
|
||||
try:
|
||||
route = await self._matcher.resolve(ctx)
|
||||
except Exception:
|
||||
logger.exception("Route resolution failed")
|
||||
if route_span_id:
|
||||
await trace.end_span(route_span_id, level="ERROR", status_message="route_error")
|
||||
return DispatchResult(success=False, error="route_error")
|
||||
|
||||
if route.agent_config_id == 0:
|
||||
if route_span_id:
|
||||
await trace.end_span(route_span_id, status_message="no_route")
|
||||
return DispatchResult(success=False, error="no_route")
|
||||
|
||||
if route_span_id:
|
||||
await trace.end_span(route_span_id, metadata={"agent_config_id": str(route.agent_config_id)})
|
||||
|
||||
if self._on_before_dispatch is not None:
|
||||
try:
|
||||
handled = await self._on_before_dispatch(msg)
|
||||
if handled:
|
||||
return DispatchResult(
|
||||
success=True,
|
||||
handled_by="hook:before_dispatch",
|
||||
thread_id=None,
|
||||
agent_config_id=None,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("before_dispatch hook failed, continuing")
|
||||
|
||||
if await self._should_check_approval(msg):
|
||||
try:
|
||||
request = await self._approval_engine.request_approval(
|
||||
msg.channel_type,
|
||||
msg.account_id,
|
||||
{},
|
||||
ApprovalAction.APPROVE_EXEC,
|
||||
msg.sender.id,
|
||||
description=f"Agent 执行: {msg.content[:100]}",
|
||||
context={"full_text": msg.content},
|
||||
)
|
||||
if request is not None:
|
||||
result = await self._approval_engine.wait_for_decision(
|
||||
msg.channel_type,
|
||||
{},
|
||||
request.id,
|
||||
)
|
||||
if result.decision != ApprovalDecision.APPROVED:
|
||||
return DispatchResult(
|
||||
success=False,
|
||||
error=f"approval_denied:{result.reason or '审批未通过'}",
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Approval check failed, blocking message")
|
||||
return DispatchResult(
|
||||
success=False,
|
||||
error="approval_check_error",
|
||||
)
|
||||
|
||||
agent_id_str = str(route.agent_config_id)
|
||||
|
||||
sk_span_id = None
|
||||
if trace:
|
||||
sk_span_id = await trace.add_span("session_key_build")
|
||||
|
||||
internal_user_id, thread_id = await self._resolve_thread_id(
|
||||
msg,
|
||||
agent_id_str,
|
||||
route.dm_scope,
|
||||
)
|
||||
|
||||
if sk_span_id:
|
||||
await trace.end_span(sk_span_id, metadata={
|
||||
"thread_id": thread_id,
|
||||
"has_internal_user": internal_user_id is not None,
|
||||
})
|
||||
|
||||
reply_sent = False
|
||||
if explicit_target:
|
||||
try:
|
||||
target_plugin = ChannelPluginRegistry.get(explicit_target)
|
||||
if target_plugin is not None and isinstance(target_plugin, MessagingProtocol):
|
||||
logger.info(
|
||||
"Message %s forwarded to explicit target: %s",
|
||||
msg.msg_id,
|
||||
explicit_target,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Explicit target forwarding failed: %s", explicit_target)
|
||||
|
||||
result = DispatchResult(
|
||||
success=True,
|
||||
thread_id=thread_id,
|
||||
agent_config_id=route.agent_config_id,
|
||||
internal_user_id=internal_user_id,
|
||||
reply_sent=reply_sent,
|
||||
)
|
||||
|
||||
if self._on_after_dispatch is not None:
|
||||
try:
|
||||
await self._on_after_dispatch(msg, result)
|
||||
except Exception:
|
||||
logger.exception("after_dispatch hook failed")
|
||||
|
||||
return result
|
||||
|
||||
async def _resolve_thread_id(
|
||||
self,
|
||||
msg: UnifiedMessage,
|
||||
agent_id_str: str,
|
||||
dm_scope: str,
|
||||
) -> tuple[str | None, str]:
|
||||
internal_user_id = None
|
||||
if self._user_mapping_repo is not None:
|
||||
try:
|
||||
user_mapping = await self._user_mapping_repo.resolve_user(
|
||||
msg.channel_type,
|
||||
msg.sender.id,
|
||||
)
|
||||
internal_user_id = user_mapping.internal_user_id
|
||||
except Exception:
|
||||
logger.exception("Failed to resolve user mapping")
|
||||
|
||||
user_id = internal_user_id or msg.sender.id
|
||||
|
||||
if msg.sender.kind == PeerKind.DIRECT:
|
||||
session_key = self._session_key_builder.build_dm_session(
|
||||
agent_id_str,
|
||||
msg.channel_type,
|
||||
msg.account_id,
|
||||
user_id,
|
||||
dm_scope=dm_scope,
|
||||
)
|
||||
else:
|
||||
session_key = self._session_key_builder.build_group_session(
|
||||
agent_id_str,
|
||||
msg.channel_type,
|
||||
msg.sender.kind.value,
|
||||
user_id,
|
||||
)
|
||||
|
||||
thread_id = session_key_to_thread_id(session_key)
|
||||
|
||||
if internal_user_id and self._thread_mapping_repo is not None:
|
||||
try:
|
||||
channel_chat_id = msg.group.id if msg.group and msg.group.id else msg.sender.id
|
||||
await self._thread_mapping_repo.resolve_thread(
|
||||
msg.channel_type,
|
||||
channel_chat_id,
|
||||
internal_user_id,
|
||||
agent_id=agent_id_str,
|
||||
thread_id=thread_id,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Failed to persist thread mapping")
|
||||
|
||||
return internal_user_id, thread_id
|
||||
346
backend/package/yuxi/channel/message/durable_send.py
Normal file
346
backend/package/yuxi/channel/message/durable_send.py
Normal file
@ -0,0 +1,346 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import random
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Callable, Coroutine
|
||||
from dataclasses import dataclass, field
|
||||
from enum import StrEnum
|
||||
from typing import Any
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DurableStrategy(StrEnum):
|
||||
REQUIRED = "required"
|
||||
BEST_EFFORT = "best_effort"
|
||||
DISABLED = "disabled"
|
||||
|
||||
|
||||
class MessageSendState(StrEnum):
|
||||
IDLE = "idle"
|
||||
RENDERING = "rendering"
|
||||
PREVIEWING = "previewing"
|
||||
SENDING = "sending"
|
||||
SENT = "sent"
|
||||
SUPPRESSED = "suppressed"
|
||||
PARTIAL_FAILED = "partial_failed"
|
||||
FAILED = "failed"
|
||||
UNKNOWN_AFTER_SEND = "unknown_after_send"
|
||||
FINALIZING = "finalizing"
|
||||
EDITING = "editing"
|
||||
EDITED = "edited"
|
||||
DELETING = "deleting"
|
||||
DELETED = "deleted"
|
||||
CANCELLED = "cancelled"
|
||||
|
||||
|
||||
class MessageReceiptPartKind(StrEnum):
|
||||
TEXT = "text"
|
||||
MEDIA = "media"
|
||||
VOICE = "voice"
|
||||
CARD = "card"
|
||||
PREVIEW = "preview"
|
||||
UNKNOWN = "unknown"
|
||||
|
||||
|
||||
@dataclass
|
||||
class MessageReceiptPart:
|
||||
platform_message_id: str
|
||||
kind: MessageReceiptPartKind = MessageReceiptPartKind.UNKNOWN
|
||||
index: int = 0
|
||||
thread_id: str | None = None
|
||||
reply_to_id: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class DurableMessageReceipt:
|
||||
primary_platform_message_id: str = ""
|
||||
platform_message_ids: list[str] = field(default_factory=list)
|
||||
parts: list[MessageReceiptPart] = field(default_factory=list)
|
||||
thread_id: str | None = None
|
||||
reply_to_id: str | None = None
|
||||
edit_token: str | None = None
|
||||
delete_token: str | None = None
|
||||
sent_at: float = 0.0
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
def __post_init__(self):
|
||||
if self.sent_at == 0.0:
|
||||
self.sent_at = time.time()
|
||||
if not self.platform_message_ids and self.primary_platform_message_id:
|
||||
self.platform_message_ids = [self.primary_platform_message_id]
|
||||
|
||||
|
||||
@dataclass
|
||||
class MessageSendContext:
|
||||
target_id: str
|
||||
content: str
|
||||
id: str = ""
|
||||
channel: str = ""
|
||||
account_id: str | None = None
|
||||
reply_to_id: str | None = None
|
||||
thread_id: str | None = None
|
||||
strategy: DurableStrategy = DurableStrategy.BEST_EFFORT
|
||||
receipt: DurableMessageReceipt | None = None
|
||||
previous_receipt: DurableMessageReceipt | None = None
|
||||
state: MessageSendState = MessageSendState.IDLE
|
||||
error: str | None = None
|
||||
attempt: int = 1
|
||||
retry_count: int = 0
|
||||
max_retries: int = 3
|
||||
min_delay_ms: int = 300
|
||||
max_delay_ms: int = 30_000
|
||||
jitter: float = 0.0
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
parts: list["MessageSendContext"] = field(default_factory=list)
|
||||
_on_commit: Callable[..., Coroutine[Any, Any, None]] | None = field(default=None, repr=False)
|
||||
_on_fail: Callable[..., Coroutine[Any, Any, None]] | None = field(default=None, repr=False)
|
||||
|
||||
def __post_init__(self):
|
||||
if not self.id:
|
||||
self.id = f"{self.channel or 'msg'}:{self.target_id}:{uuid.uuid4().hex[:8]}"
|
||||
|
||||
@property
|
||||
def is_terminal(self) -> bool:
|
||||
return self.state in (
|
||||
MessageSendState.SENT,
|
||||
MessageSendState.SUPPRESSED,
|
||||
MessageSendState.PARTIAL_FAILED,
|
||||
MessageSendState.FAILED,
|
||||
MessageSendState.CANCELLED,
|
||||
)
|
||||
|
||||
async def render(self) -> str:
|
||||
self.state = MessageSendState.RENDERING
|
||||
return self.content
|
||||
|
||||
def _backoff_delay(self, attempt: int) -> int:
|
||||
base = self.min_delay_ms * (2 ** (attempt - 1))
|
||||
delay = min(base, self.max_delay_ms)
|
||||
if self.jitter > 0:
|
||||
offset = (random.random() * 2 - 1) * self.jitter
|
||||
delay = int(delay * (1 + offset))
|
||||
return max(0, delay)
|
||||
|
||||
async def send(self, send_fn: Callable[..., Coroutine[Any, Any, str | None]]) -> DurableMessageReceipt | None:
|
||||
self.state = MessageSendState.SENDING
|
||||
last_error = None
|
||||
total_attempts = self.max_retries + 1
|
||||
for attempt_idx in range(total_attempts):
|
||||
self.attempt = attempt_idx + 1
|
||||
try:
|
||||
message_id = await send_fn(self.content)
|
||||
if message_id:
|
||||
self.state = MessageSendState.SENT
|
||||
self.receipt = DurableMessageReceipt(primary_platform_message_id=message_id)
|
||||
return self.receipt
|
||||
self.state = MessageSendState.SENT
|
||||
return None
|
||||
except Exception as e:
|
||||
last_error = str(e)
|
||||
self.retry_count = attempt_idx + 1
|
||||
logger.warning(
|
||||
"Message send attempt %d/%d failed: %s",
|
||||
attempt_idx + 1,
|
||||
total_attempts,
|
||||
e,
|
||||
)
|
||||
if attempt_idx < total_attempts - 1:
|
||||
delay = self._backoff_delay(attempt_idx + 1)
|
||||
if delay > 0:
|
||||
await asyncio.sleep(delay / 1000)
|
||||
|
||||
self.state = MessageSendState.FAILED
|
||||
self.error = last_error
|
||||
return None
|
||||
|
||||
async def send_batch(
|
||||
self,
|
||||
contents: list[str],
|
||||
send_fn: Callable[..., Coroutine[Any, Any, str | None]],
|
||||
) -> list[DurableMessageReceipt | None]:
|
||||
if not contents:
|
||||
return []
|
||||
|
||||
results: list[DurableMessageReceipt | None] = []
|
||||
self.parts.clear()
|
||||
failed_count = 0
|
||||
|
||||
for i, content in enumerate(contents):
|
||||
part = MessageSendContext(
|
||||
target_id=self.target_id,
|
||||
content=content,
|
||||
id=f"{self.id}#{i}",
|
||||
channel=self.channel,
|
||||
account_id=self.account_id,
|
||||
reply_to_id=self.reply_to_id,
|
||||
thread_id=self.thread_id,
|
||||
strategy=self.strategy,
|
||||
max_retries=self.max_retries,
|
||||
min_delay_ms=self.min_delay_ms,
|
||||
max_delay_ms=self.max_delay_ms,
|
||||
jitter=self.jitter,
|
||||
)
|
||||
receipt = await part.send(send_fn)
|
||||
results.append(receipt)
|
||||
self.parts.append(part)
|
||||
if receipt is None and part.state == MessageSendState.FAILED:
|
||||
failed_count += 1
|
||||
|
||||
total = len(contents)
|
||||
if failed_count == total:
|
||||
self.state = MessageSendState.FAILED
|
||||
self.error = f"All {total} parts failed"
|
||||
elif failed_count > 0:
|
||||
self.state = MessageSendState.PARTIAL_FAILED
|
||||
self.error = f"{failed_count}/{total} parts failed"
|
||||
else:
|
||||
self.state = MessageSendState.SENT
|
||||
return results
|
||||
|
||||
def mark_suppressed(self, reason: str = "") -> None:
|
||||
self.state = MessageSendState.SUPPRESSED
|
||||
self.error = reason
|
||||
|
||||
async def edit(
|
||||
self, edit_fn: Callable[..., Coroutine[Any, Any, str | None]], new_content: str
|
||||
) -> DurableMessageReceipt | None:
|
||||
if self.receipt is None:
|
||||
logger.warning("Cannot edit message without receipt")
|
||||
return None
|
||||
self.state = MessageSendState.EDITING
|
||||
try:
|
||||
new_id = await edit_fn(self.receipt.primary_platform_message_id, new_content)
|
||||
if new_id:
|
||||
self.receipt.primary_platform_message_id = new_id
|
||||
if new_id not in self.receipt.platform_message_ids:
|
||||
self.receipt.platform_message_ids.append(new_id)
|
||||
self.state = MessageSendState.EDITED
|
||||
return self.receipt
|
||||
except Exception as e:
|
||||
self.state = MessageSendState.FAILED
|
||||
self.error = str(e)
|
||||
return None
|
||||
|
||||
async def delete(self, delete_fn: Callable[..., Coroutine[Any, Any, None]]) -> bool:
|
||||
if self.receipt is None:
|
||||
return False
|
||||
self.state = MessageSendState.DELETING
|
||||
try:
|
||||
await delete_fn(self.receipt.primary_platform_message_id)
|
||||
self.state = MessageSendState.DELETED
|
||||
return True
|
||||
except Exception as e:
|
||||
self.state = MessageSendState.FAILED
|
||||
self.error = str(e)
|
||||
return False
|
||||
|
||||
def mark_cancelled(self) -> None:
|
||||
self.state = MessageSendState.CANCELLED
|
||||
|
||||
def mark_unknown_after_send(self) -> None:
|
||||
self.state = MessageSendState.UNKNOWN_AFTER_SEND
|
||||
|
||||
async def commit(self) -> None:
|
||||
if self._on_commit:
|
||||
await self._on_commit(self.receipt)
|
||||
|
||||
async def fail(self, error: Exception | None = None) -> None:
|
||||
if self._on_fail:
|
||||
if error is None:
|
||||
error = Exception(self.error or "send failed")
|
||||
await self._on_fail(error)
|
||||
|
||||
|
||||
class OutboundBridge:
|
||||
def __init__(
|
||||
self,
|
||||
send_text_fn: Callable[..., Coroutine[Any, Any, str | None]] | None = None,
|
||||
send_media_fn: Callable[..., Coroutine[Any, Any, str | None]] | None = None,
|
||||
send_payload_fn: Callable[..., Coroutine[Any, Any, str | None]] | None = None,
|
||||
):
|
||||
self._send_text = send_text_fn
|
||||
self._send_media = send_media_fn
|
||||
self._send_payload = send_payload_fn
|
||||
|
||||
async def text(
|
||||
self,
|
||||
target_id: str,
|
||||
content: str,
|
||||
*,
|
||||
reply_to_id: str | None = None,
|
||||
thread_id: str | None = None,
|
||||
) -> str | None:
|
||||
if not self._send_text:
|
||||
raise RuntimeError("OutboundBridge: send_text not configured")
|
||||
return await self._send_text(target_id, content, reply_to_id=reply_to_id, thread_id=thread_id)
|
||||
|
||||
async def media(
|
||||
self,
|
||||
target_id: str,
|
||||
media_url: str,
|
||||
text: str = "",
|
||||
*,
|
||||
reply_to_id: str | None = None,
|
||||
thread_id: str | None = None,
|
||||
audio_as_voice: bool = False,
|
||||
) -> str | None:
|
||||
if not self._send_media:
|
||||
raise RuntimeError("OutboundBridge: send_media not configured")
|
||||
return await self._send_media(
|
||||
target_id, media_url, text, reply_to_id=reply_to_id, thread_id=thread_id, audio_as_voice=audio_as_voice
|
||||
)
|
||||
|
||||
async def payload(
|
||||
self,
|
||||
target_id: str,
|
||||
payload: Any,
|
||||
*,
|
||||
reply_to_id: str | None = None,
|
||||
thread_id: str | None = None,
|
||||
) -> str | None:
|
||||
if not self._send_payload:
|
||||
raise RuntimeError("OutboundBridge: send_payload not configured")
|
||||
return await self._send_payload(target_id, payload, reply_to_id=reply_to_id, thread_id=thread_id)
|
||||
|
||||
|
||||
class DurableSendContextManager:
|
||||
def __init__(
|
||||
self,
|
||||
ctx: "MessageSendContext",
|
||||
*,
|
||||
on_commit: Callable[..., Coroutine[Any, Any, None]] | None = None,
|
||||
on_fail: Callable[..., Coroutine[Any, Any, None]] | None = None,
|
||||
):
|
||||
self.ctx = ctx
|
||||
self.ctx._on_commit = on_commit
|
||||
self.ctx._on_fail = on_fail
|
||||
|
||||
async def __aenter__(self) -> "MessageSendContext":
|
||||
return self.ctx
|
||||
|
||||
async def __aexit__(self, exc_type, exc_val, exc_tb) -> bool:
|
||||
if exc_type is not None:
|
||||
await self.ctx.fail(exc_val)
|
||||
return False
|
||||
|
||||
|
||||
async def send_durable_message_batch(
|
||||
ctx: "MessageSendContext",
|
||||
send_fn: Callable[..., Coroutine[Any, Any, str | None]],
|
||||
*,
|
||||
on_commit: Callable[..., Coroutine[Any, Any, None]] | None = None,
|
||||
on_fail: Callable[..., Coroutine[Any, Any, None]] | None = None,
|
||||
) -> "DurableMessageReceipt | None":
|
||||
if on_commit:
|
||||
ctx._on_commit = on_commit
|
||||
if on_fail:
|
||||
ctx._on_fail = on_fail
|
||||
await ctx.render()
|
||||
result = await ctx.send(send_fn)
|
||||
if result is not None and ctx.state in (MessageSendState.SENT, MessageSendState.SUPPRESSED):
|
||||
await ctx.commit()
|
||||
else:
|
||||
await ctx.fail()
|
||||
return result
|
||||
183
backend/package/yuxi/channel/message/handler.py
Normal file
183
backend/package/yuxi/channel/message/handler.py
Normal file
@ -0,0 +1,183 @@
|
||||
import logging
|
||||
|
||||
from sqlalchemy import select
|
||||
|
||||
from yuxi.channel.message.bridge import AgentBridge
|
||||
from yuxi.channel.message.block_reply_pipeline import create_block_reply_pipeline
|
||||
from yuxi.channel.message.models import ChunkType, DispatchResult, ReplyStage, StreamingChunk, UnifiedMessage
|
||||
from yuxi.channel.protocols import OutboundProtocol
|
||||
from yuxi.channel.plugins.registry import ChannelPluginRegistry
|
||||
from yuxi.channel.session import SessionRecorder
|
||||
from yuxi.repositories.conversation_repository import ConversationRepository
|
||||
from yuxi.services.chat_service import stream_agent_chat
|
||||
from yuxi.storage.postgres.manager import pg_manager
|
||||
from yuxi.storage.postgres.models_business import User
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def channel_message_handler(
|
||||
msg: UnifiedMessage,
|
||||
result: DispatchResult,
|
||||
channel_manager=None,
|
||||
) -> None:
|
||||
channel_type = msg.channel_type
|
||||
account_id = msg.account_id
|
||||
|
||||
logger.info(
|
||||
"channel_turn start: msg_id=%s channel=%s agent_config_id=%s",
|
||||
msg.msg_id,
|
||||
channel_type,
|
||||
result.agent_config_id,
|
||||
)
|
||||
|
||||
user = await _resolve_user(result, msg.msg_id, channel_type)
|
||||
if user is None:
|
||||
logger.error(
|
||||
"channel_turn drop: msg_id=%s channel=%s reason=no_user",
|
||||
msg.msg_id,
|
||||
channel_type,
|
||||
)
|
||||
return
|
||||
|
||||
if channel_manager:
|
||||
channel_manager.inc_active_runs(channel_type, account_id)
|
||||
channel_manager.update_last_message_at(channel_type, account_id)
|
||||
|
||||
try:
|
||||
bridge = AgentBridge(stream_fn=stream_agent_chat, channel_manager=channel_manager)
|
||||
|
||||
async with pg_manager.get_async_session_context() as db:
|
||||
await _record_session(msg, result, user, db)
|
||||
|
||||
pipeline = create_block_reply_pipeline(
|
||||
send_fn=_make_send_fn(msg, result, channel_manager),
|
||||
msg=msg,
|
||||
group_msg_id=msg.group.id if msg.group else None,
|
||||
)
|
||||
|
||||
async for chunk in bridge.invoke_stream(
|
||||
msg,
|
||||
result,
|
||||
result.agent_config_id,
|
||||
user,
|
||||
db,
|
||||
):
|
||||
await pipeline.enqueue(chunk)
|
||||
|
||||
await pipeline.wait_idle()
|
||||
|
||||
if result.thread_id:
|
||||
await _touch_session(msg, result, db)
|
||||
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"channel_turn error: msg_id=%s channel=%s",
|
||||
msg.msg_id,
|
||||
channel_type,
|
||||
)
|
||||
finally:
|
||||
if channel_manager:
|
||||
channel_manager.dec_active_runs(channel_type, account_id)
|
||||
|
||||
|
||||
def _make_send_fn(msg: UnifiedMessage, result: DispatchResult, channel_manager):
|
||||
channel_type = msg.channel_type
|
||||
plugin = ChannelPluginRegistry.get(channel_type)
|
||||
if plugin is None or not isinstance(plugin, OutboundProtocol):
|
||||
return _noop_send_fn
|
||||
|
||||
target_id = msg.group.id if msg.group and msg.group.id else msg.sender.id
|
||||
|
||||
async def send_fn(chunk: StreamingChunk) -> None:
|
||||
try:
|
||||
if chunk.stage == ReplyStage.BLOCK and chunk.chunk_type == ChunkType.TEXT_DELTA:
|
||||
await plugin.send_text(
|
||||
target_id,
|
||||
chunk.content,
|
||||
reply_to_id=result.thread_id,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"channel_turn deliver_reply error: msg_id=%s channel=%s",
|
||||
msg.msg_id,
|
||||
channel_type,
|
||||
)
|
||||
|
||||
return send_fn
|
||||
|
||||
|
||||
async def _noop_send_fn(chunk: StreamingChunk) -> None:
|
||||
pass
|
||||
|
||||
|
||||
async def _record_session(
|
||||
msg: UnifiedMessage,
|
||||
result: DispatchResult,
|
||||
user: User,
|
||||
db,
|
||||
) -> None:
|
||||
session_key = result.session_key or f"{msg.channel_type}:{msg.account_id}:{msg.sender.id}"
|
||||
try:
|
||||
conv_repo = ConversationRepository(db)
|
||||
recorder = SessionRecorder(conv_repo)
|
||||
await recorder.record_or_update(
|
||||
msg,
|
||||
result,
|
||||
session_key,
|
||||
user_id=str(user.id),
|
||||
title=msg.sender.display_name or str(msg.sender.id),
|
||||
)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"channel_turn session record failed: msg_id=%s channel=%s",
|
||||
msg.msg_id,
|
||||
msg.channel_type,
|
||||
)
|
||||
|
||||
|
||||
async def _touch_session(
|
||||
msg: UnifiedMessage,
|
||||
result: DispatchResult,
|
||||
db,
|
||||
) -> None:
|
||||
try:
|
||||
conv_repo = ConversationRepository(db)
|
||||
recorder = SessionRecorder(conv_repo)
|
||||
await recorder.touch(msg, result.thread_id)
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"channel_turn session touch failed: msg_id=%s thread=%s",
|
||||
msg.msg_id,
|
||||
result.thread_id,
|
||||
)
|
||||
|
||||
|
||||
async def _resolve_user(
|
||||
result: DispatchResult,
|
||||
msg_id: str,
|
||||
channel_type: str,
|
||||
) -> User | None:
|
||||
async with pg_manager.get_async_session_context() as db:
|
||||
if result.internal_user_id:
|
||||
try:
|
||||
uid = int(result.internal_user_id)
|
||||
except (ValueError, TypeError):
|
||||
uid = None
|
||||
if uid is not None:
|
||||
user = (await db.execute(select(User).where(User.id == uid))).scalar_one_or_none()
|
||||
if user is not None:
|
||||
logger.info(
|
||||
"channel_turn resolve_user: msg_id=%s channel=%s resolved_by=internal_user_id",
|
||||
msg_id,
|
||||
channel_type,
|
||||
)
|
||||
return user
|
||||
|
||||
logger.warning(
|
||||
"channel_turn resolve_user: msg_id=%s channel=%s internal_user_id=%s no match",
|
||||
msg_id,
|
||||
channel_type,
|
||||
result.internal_user_id,
|
||||
)
|
||||
return None
|
||||
340
backend/package/yuxi/channel/message/idempotency.py
Normal file
340
backend/package/yuxi/channel/message/idempotency.py
Normal file
@ -0,0 +1,340 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from abc import ABC, abstractmethod
|
||||
from enum import StrEnum
|
||||
|
||||
from cachetools import LRUCache
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_DEFAULT_TTL_SECONDS = 300
|
||||
|
||||
_IDEMPOTENCY_KEY_SEPARATOR = ":"
|
||||
_IDEMPOTENCY_NONCE_MARKER = ":nonce:"
|
||||
|
||||
|
||||
class ClaimStatus(StrEnum):
|
||||
INVALID = "invalid"
|
||||
DUPLICATE = "duplicate"
|
||||
INFLIGHT = "inflight"
|
||||
CLAIMED = "claimed"
|
||||
|
||||
|
||||
class IdempotencyBackend(ABC):
|
||||
"""幂等性后端抽象基类"""
|
||||
|
||||
@abstractmethod
|
||||
async def claim(self, msg_id: str, ttl: int = _DEFAULT_TTL_SECONDS) -> bool: ...
|
||||
|
||||
@abstractmethod
|
||||
async def claim_with_status(self, msg_id: str) -> ClaimStatus: ...
|
||||
|
||||
@abstractmethod
|
||||
async def check_and_set(self, msg_id: str, ttl: int = _DEFAULT_TTL_SECONDS) -> bool: ...
|
||||
|
||||
@abstractmethod
|
||||
async def commit(self, msg_id: str) -> None: ...
|
||||
|
||||
@abstractmethod
|
||||
async def release(self, msg_id: str) -> None: ...
|
||||
|
||||
@abstractmethod
|
||||
async def release_and_forget(self, msg_id: str) -> None: ...
|
||||
|
||||
@abstractmethod
|
||||
async def clear_inflight(self, msg_id: str | None = None) -> None: ...
|
||||
|
||||
@abstractmethod
|
||||
async def reset(self) -> None: ...
|
||||
|
||||
|
||||
class InMemoryBackend(IdempotencyBackend):
|
||||
def __init__(self) -> None:
|
||||
self._cache: LRUCache = LRUCache(maxsize=10_000)
|
||||
self._inflight: set[str] = set()
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
def _prune_expired(self, now: float) -> None:
|
||||
expired = [mid for mid, (ts, _ttl) in self._cache.items() if now - ts >= _ttl]
|
||||
for mid in expired:
|
||||
del self._cache[mid]
|
||||
self._inflight.discard(mid)
|
||||
|
||||
async def check_and_set(self, msg_id: str, ttl: int = _DEFAULT_TTL_SECONDS) -> bool:
|
||||
if not msg_id:
|
||||
return False
|
||||
async with self._lock:
|
||||
now = time.monotonic()
|
||||
if msg_id in self._cache:
|
||||
timestamp, entry_ttl = self._cache[msg_id]
|
||||
if now - timestamp < entry_ttl:
|
||||
return False
|
||||
self._prune_expired(now)
|
||||
self._cache[msg_id] = (now, ttl)
|
||||
return True
|
||||
|
||||
async def claim(self, msg_id: str, ttl: int = _DEFAULT_TTL_SECONDS) -> bool:
|
||||
if not msg_id:
|
||||
return False
|
||||
async with self._lock:
|
||||
self._prune_expired(time.monotonic())
|
||||
if msg_id in self._inflight:
|
||||
return False
|
||||
now = time.monotonic()
|
||||
if msg_id in self._cache:
|
||||
timestamp, entry_ttl = self._cache[msg_id]
|
||||
if now - timestamp < entry_ttl:
|
||||
return False
|
||||
self._cache[msg_id] = (now, ttl)
|
||||
self._inflight.add(msg_id)
|
||||
return True
|
||||
|
||||
async def claim_with_status(self, msg_id: str) -> ClaimStatus:
|
||||
if not msg_id:
|
||||
return ClaimStatus.INVALID
|
||||
async with self._lock:
|
||||
self._prune_expired(time.monotonic())
|
||||
if msg_id in self._inflight:
|
||||
return ClaimStatus.INFLIGHT
|
||||
now = time.monotonic()
|
||||
if msg_id in self._cache:
|
||||
timestamp, entry_ttl = self._cache[msg_id]
|
||||
if now - timestamp < entry_ttl:
|
||||
return ClaimStatus.DUPLICATE
|
||||
self._cache[msg_id] = (now, _DEFAULT_TTL_SECONDS)
|
||||
self._inflight.add(msg_id)
|
||||
return ClaimStatus.CLAIMED
|
||||
|
||||
async def commit(self, msg_id: str) -> None:
|
||||
async with self._lock:
|
||||
self._inflight.discard(msg_id)
|
||||
|
||||
async def release(self, msg_id: str) -> None:
|
||||
async with self._lock:
|
||||
self._inflight.discard(msg_id)
|
||||
|
||||
async def release_and_forget(self, msg_id: str) -> None:
|
||||
async with self._lock:
|
||||
self._inflight.discard(msg_id)
|
||||
self._cache.pop(msg_id, None)
|
||||
|
||||
async def clear_inflight(self, msg_id: str | None = None) -> None:
|
||||
async with self._lock:
|
||||
if msg_id is not None:
|
||||
self._inflight.discard(msg_id)
|
||||
else:
|
||||
self._inflight.clear()
|
||||
|
||||
async def reset(self) -> None:
|
||||
async with self._lock:
|
||||
self._cache.clear()
|
||||
self._inflight.clear()
|
||||
|
||||
|
||||
class RedisBackend(IdempotencyBackend):
|
||||
"""基于 Redis 的分布式幂等性后端
|
||||
|
||||
使用 SET NX EX 实现原子性检查和设置,结合 SADD/SREM 管理 inflight 状态。
|
||||
"""
|
||||
|
||||
_INFLIGHT_SUFFIX = ":inflight"
|
||||
|
||||
def __init__(self, redis_client, key_prefix: str = "yuxi:idempotency") -> None:
|
||||
self._redis = redis_client
|
||||
self._key_prefix = key_prefix
|
||||
|
||||
def _cache_key(self, msg_id: str) -> str:
|
||||
return f"{self._key_prefix}:msg:{msg_id}"
|
||||
|
||||
def _inflight_key(self, msg_id: str) -> str:
|
||||
return f"{self._key_prefix}{self._INFLIGHT_SUFFIX}"
|
||||
|
||||
async def check_and_set(self, msg_id: str, ttl: int = _DEFAULT_TTL_SECONDS) -> bool:
|
||||
if not msg_id:
|
||||
return False
|
||||
key = self._cache_key(msg_id)
|
||||
return await self._redis.set(key, "1", nx=True, ex=ttl) or False
|
||||
|
||||
async def claim(self, msg_id: str, ttl: int = _DEFAULT_TTL_SECONDS) -> bool:
|
||||
if not msg_id:
|
||||
return False
|
||||
inflight_key = self._inflight_key(msg_id)
|
||||
added = await self._redis.sadd(inflight_key, msg_id)
|
||||
if added == 0:
|
||||
return False
|
||||
key = self._cache_key(msg_id)
|
||||
set_ok = await self._redis.set(key, "1", nx=True, ex=ttl)
|
||||
if not set_ok:
|
||||
await self._redis.srem(inflight_key, msg_id)
|
||||
return False
|
||||
return True
|
||||
|
||||
async def claim_with_status(self, msg_id: str) -> ClaimStatus:
|
||||
if not msg_id:
|
||||
return ClaimStatus.INVALID
|
||||
inflight_key = self._inflight_key(msg_id)
|
||||
if await self._redis.sismember(inflight_key, msg_id):
|
||||
return ClaimStatus.INFLIGHT
|
||||
key = self._cache_key(msg_id)
|
||||
if await self._redis.exists(key):
|
||||
return ClaimStatus.DUPLICATE
|
||||
added = await self._redis.sadd(inflight_key, msg_id)
|
||||
if added == 0:
|
||||
return ClaimStatus.INFLIGHT
|
||||
await self._redis.set(key, "1", ex=_DEFAULT_TTL_SECONDS)
|
||||
return ClaimStatus.CLAIMED
|
||||
|
||||
async def commit(self, msg_id: str) -> None:
|
||||
inflight_key = self._inflight_key(msg_id)
|
||||
await self._redis.srem(inflight_key, msg_id)
|
||||
|
||||
async def release(self, msg_id: str) -> None:
|
||||
inflight_key = self._inflight_key(msg_id)
|
||||
await self._redis.srem(inflight_key, msg_id)
|
||||
|
||||
async def release_and_forget(self, msg_id: str) -> None:
|
||||
inflight_key = self._inflight_key(msg_id)
|
||||
cache_key = self._cache_key(msg_id)
|
||||
await self._redis.srem(inflight_key, msg_id)
|
||||
await self._redis.delete(cache_key)
|
||||
|
||||
async def clear_inflight(self, msg_id: str | None = None) -> None:
|
||||
if msg_id is not None:
|
||||
inflight_key = self._inflight_key(msg_id)
|
||||
await self._redis.srem(inflight_key, msg_id)
|
||||
else:
|
||||
keys = await self._redis.keys(f"{self._key_prefix}{self._INFLIGHT_SUFFIX}:*")
|
||||
for key in keys:
|
||||
await self._redis.delete(key)
|
||||
|
||||
async def reset(self) -> None:
|
||||
keys = await self._redis.keys(f"{self._key_prefix}:*")
|
||||
for key in keys:
|
||||
await self._redis.delete(key)
|
||||
|
||||
|
||||
_backend: IdempotencyBackend | None = None
|
||||
_backend_lock = asyncio.Lock()
|
||||
|
||||
|
||||
def _get_redis_client():
|
||||
redis_url = os.getenv("REDIS_URL", os.getenv("YUXI_REDIS_URL", ""))
|
||||
if not redis_url:
|
||||
return None
|
||||
try:
|
||||
from redis.asyncio import Redis
|
||||
|
||||
return Redis.from_url(redis_url, decode_responses=True)
|
||||
except Exception:
|
||||
logger.warning("Failed to create Redis client for idempotency, falling back to in-memory")
|
||||
return None
|
||||
|
||||
|
||||
async def get_backend() -> IdempotencyBackend:
|
||||
global _backend
|
||||
if _backend is not None:
|
||||
return _backend
|
||||
async with _backend_lock:
|
||||
if _backend is not None:
|
||||
return _backend
|
||||
redis_client = _get_redis_client()
|
||||
if redis_client is not None:
|
||||
_backend = RedisBackend(redis_client)
|
||||
logger.info("Idempotency backend: Redis")
|
||||
else:
|
||||
_backend = InMemoryBackend()
|
||||
logger.info("Idempotency backend: InMemory (single-process)")
|
||||
return _backend
|
||||
|
||||
|
||||
async def _ensure_backend() -> IdempotencyBackend:
|
||||
if _backend is not None:
|
||||
return _backend
|
||||
return await get_backend()
|
||||
|
||||
|
||||
async def check_and_set(msg_id: str, ttl: int = _DEFAULT_TTL_SECONDS) -> bool:
|
||||
backend = await _ensure_backend()
|
||||
return await backend.check_and_set(msg_id, ttl)
|
||||
|
||||
|
||||
async def claim(msg_id: str, ttl: int = _DEFAULT_TTL_SECONDS) -> bool:
|
||||
backend = await _ensure_backend()
|
||||
return await backend.claim(msg_id, ttl)
|
||||
|
||||
|
||||
async def claim_with_status(msg_id: str) -> ClaimStatus:
|
||||
backend = await _ensure_backend()
|
||||
return await backend.claim_with_status(msg_id)
|
||||
|
||||
|
||||
async def commit(msg_id: str) -> None:
|
||||
backend = await _ensure_backend()
|
||||
await backend.commit(msg_id)
|
||||
|
||||
|
||||
async def release(msg_id: str) -> None:
|
||||
backend = await _ensure_backend()
|
||||
await backend.release(msg_id)
|
||||
|
||||
|
||||
async def release_and_forget(msg_id: str) -> None:
|
||||
backend = await _ensure_backend()
|
||||
await backend.release_and_forget(msg_id)
|
||||
|
||||
|
||||
async def consume(msg_id: str) -> bool:
|
||||
if not await claim(msg_id):
|
||||
return False
|
||||
await release(msg_id)
|
||||
return True
|
||||
|
||||
|
||||
def build_idempotency_key(prefix: str, key: str, nonce: str | None = None) -> str:
|
||||
base = f"{prefix}{_IDEMPOTENCY_KEY_SEPARATOR}{key}"
|
||||
if nonce:
|
||||
return f"{base}{_IDEMPOTENCY_NONCE_MARKER}{nonce}"
|
||||
return base
|
||||
|
||||
|
||||
def parse_idempotency_key(idempotency_key: str) -> tuple[str, str] | None:
|
||||
if not idempotency_key:
|
||||
return None
|
||||
sep_idx = idempotency_key.find(_IDEMPOTENCY_KEY_SEPARATOR)
|
||||
if sep_idx < 0:
|
||||
return None
|
||||
prefix = idempotency_key[:sep_idx]
|
||||
body = idempotency_key[sep_idx + 1 :]
|
||||
nonce_marker = body.rfind(_IDEMPOTENCY_NONCE_MARKER)
|
||||
if nonce_marker >= 0:
|
||||
return prefix, body[:nonce_marker]
|
||||
return prefix, body
|
||||
|
||||
|
||||
async def clear_inflight(msg_id: str | None = None) -> None:
|
||||
backend = await _ensure_backend()
|
||||
await backend.clear_inflight(msg_id)
|
||||
|
||||
|
||||
async def reset() -> None:
|
||||
backend = await _ensure_backend()
|
||||
await backend.reset()
|
||||
|
||||
|
||||
def reset_sync() -> None:
|
||||
global _backend
|
||||
if _backend is not None:
|
||||
if isinstance(_backend, InMemoryBackend):
|
||||
_backend._cache.clear()
|
||||
_backend._inflight.clear()
|
||||
else:
|
||||
import asyncio
|
||||
|
||||
asyncio.get_event_loop().run_until_complete(reset())
|
||||
else:
|
||||
_backend = InMemoryBackend()
|
||||
|
||||
|
||||
reset_for_tests = reset_sync
|
||||
202
backend/package/yuxi/channel/message/langfuse_trace.py
Normal file
202
backend/package/yuxi/channel/message/langfuse_trace.py
Normal file
@ -0,0 +1,202 @@
|
||||
import logging
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
from yuxi.channel.message.models import DispatchResult, UnifiedMessage
|
||||
from yuxi.services.langfuse_service import (
|
||||
get_langfuse_client,
|
||||
is_langfuse_enabled,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ChannelTrace:
|
||||
"""渠道消息处理链路的 Langfuse trace 包装器。
|
||||
|
||||
复用现有 langfuse_service 的 Langfuse 客户端,创建实际 Langfuse
|
||||
span 以追踪消息处理每个阶段(分发、校验、规则匹配等)的耗时和状态。
|
||||
"""
|
||||
|
||||
def __init__(self, msg: UnifiedMessage, dispatch_result: DispatchResult | None = None):
|
||||
self._msg = msg
|
||||
self._dispatch_result = dispatch_result
|
||||
self._client = get_langfuse_client()
|
||||
self._enabled = is_langfuse_enabled() and self._client is not None
|
||||
self._trace_id: str | None = None
|
||||
self._trace_obj: Any = None
|
||||
self._trace_start_time: float | None = None
|
||||
self._live_spans: dict[str, Any] = {}
|
||||
|
||||
@property
|
||||
def trace_name(self) -> str:
|
||||
peer_kind = self._msg.sender.kind.value
|
||||
return f"channel:{self._msg.channel_type}:{peer_kind}"
|
||||
|
||||
@property
|
||||
def trace_id(self) -> str | None:
|
||||
return self._trace_id
|
||||
|
||||
@property
|
||||
def span_count(self) -> int:
|
||||
return len(self._live_spans)
|
||||
|
||||
async def start(self) -> "ChannelTrace":
|
||||
if not self._enabled:
|
||||
return self
|
||||
|
||||
try:
|
||||
self._trace_start_time = time.monotonic()
|
||||
self._trace_id = self._client.create_trace_id()
|
||||
self._trace_obj = self._client.trace(
|
||||
id=self._trace_id,
|
||||
name=self.trace_name,
|
||||
metadata=self._build_base_metadata(),
|
||||
tags=self._build_base_tags(),
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Failed to create channel trace: %s", self.trace_name)
|
||||
self._enabled = False
|
||||
|
||||
return self
|
||||
|
||||
async def add_span(
|
||||
self,
|
||||
name: str,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
input_data: Any = None,
|
||||
output_data: Any = None,
|
||||
level: str = "DEFAULT",
|
||||
status_message: str | None = None,
|
||||
parent_span_id: str | None = None,
|
||||
) -> str | None:
|
||||
"""创建 Langfuse span 并追踪其生命周期。
|
||||
|
||||
返回 span_id 供后续 end_span() 结束使用。
|
||||
调用者必须在完成处理后调用 end_span() 或依赖 finish() 统一结束。
|
||||
"""
|
||||
if not self._enabled or self._trace_obj is None:
|
||||
return None
|
||||
|
||||
try:
|
||||
parent = self._live_spans.get(parent_span_id) if parent_span_id else self._trace_obj
|
||||
span = parent.span(
|
||||
name=name,
|
||||
input=input_data,
|
||||
output=output_data,
|
||||
metadata=metadata or {},
|
||||
level=level,
|
||||
status_message=status_message,
|
||||
)
|
||||
span_id: str = span.id
|
||||
self._live_spans[span_id] = span
|
||||
return span_id
|
||||
except Exception:
|
||||
logger.exception("Failed to create span '%s' for trace: %s", name, self.trace_name)
|
||||
return None
|
||||
|
||||
async def end_span(
|
||||
self,
|
||||
span_id: str,
|
||||
output_data: Any = None,
|
||||
status_message: str | None = None,
|
||||
level: str | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
"""结束指定 span,记录输出与最终状态。"""
|
||||
span = self._live_spans.pop(span_id, None)
|
||||
if span is None:
|
||||
return
|
||||
|
||||
try:
|
||||
update_kwargs: dict[str, Any] = {}
|
||||
if output_data is not None:
|
||||
update_kwargs["output"] = output_data
|
||||
if level is not None:
|
||||
update_kwargs["level"] = level
|
||||
if status_message is not None:
|
||||
update_kwargs["status_message"] = status_message
|
||||
if metadata is not None:
|
||||
update_kwargs["metadata"] = metadata
|
||||
if update_kwargs:
|
||||
span.update(**update_kwargs)
|
||||
span.end()
|
||||
except Exception:
|
||||
logger.exception("Failed to end span '%s' for trace: %s", span_id, self.trace_name)
|
||||
|
||||
async def finish(self, error: str | None = None) -> None:
|
||||
if not self._enabled or self._trace_obj is None:
|
||||
return
|
||||
|
||||
try:
|
||||
pending = dict(self._live_spans)
|
||||
self._live_spans.clear()
|
||||
for span_id, span in pending.items():
|
||||
try:
|
||||
span.end()
|
||||
except Exception:
|
||||
logger.exception("Failed to end pending span '%s'", span_id)
|
||||
|
||||
duration_ms = None
|
||||
if self._trace_start_time is not None:
|
||||
duration_ms = (time.monotonic() - self._trace_start_time) * 1000
|
||||
|
||||
metadata: dict[str, Any] = {
|
||||
**self._build_base_metadata(),
|
||||
"span_count": len(pending),
|
||||
"dispatch_success": error is None,
|
||||
**(self._build_dispatch_metadata() if self._dispatch_result else {}),
|
||||
}
|
||||
if duration_ms is not None:
|
||||
metadata["duration_ms"] = round(duration_ms, 2)
|
||||
|
||||
self._trace_obj.update(
|
||||
output=error or "success",
|
||||
metadata=metadata,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Failed to finish channel trace: %s", self.trace_name)
|
||||
|
||||
def _build_base_metadata(self) -> dict[str, Any]:
|
||||
return {
|
||||
"source": "channel",
|
||||
"channel_type": self._msg.channel_type,
|
||||
"account_id": self._msg.account_id,
|
||||
"peer_kind": self._msg.sender.kind.value,
|
||||
"peer_id": self._msg.sender.id,
|
||||
"msg_id": self._msg.msg_id,
|
||||
"message_type": self._msg.message_type.value,
|
||||
"feature": "channel_dispatch",
|
||||
}
|
||||
|
||||
def _build_dispatch_metadata(self) -> dict[str, Any]:
|
||||
if self._dispatch_result is None:
|
||||
return {}
|
||||
return {
|
||||
"agent_config_id": str(self._dispatch_result.agent_config_id)
|
||||
if self._dispatch_result.agent_config_id
|
||||
else None,
|
||||
"session_key": getattr(self._dispatch_result, "session_key", None),
|
||||
"thread_id": self._dispatch_result.thread_id,
|
||||
"matched_by": getattr(self._dispatch_result, "matched_by", None),
|
||||
}
|
||||
|
||||
def _build_base_tags(self) -> list[str]:
|
||||
tags = [
|
||||
"yuxi",
|
||||
"channel",
|
||||
f"channel:{self._msg.channel_type}",
|
||||
f"peer_kind:{self._msg.sender.kind.value}",
|
||||
]
|
||||
if self._dispatch_result and self._dispatch_result.agent_config_id:
|
||||
tags.append(f"agent_config:{self._dispatch_result.agent_config_id}")
|
||||
return tags
|
||||
|
||||
|
||||
async def create_channel_trace(
|
||||
msg: UnifiedMessage,
|
||||
dispatch_result: DispatchResult | None = None,
|
||||
) -> ChannelTrace:
|
||||
trace = ChannelTrace(msg, dispatch_result)
|
||||
await trace.start()
|
||||
return trace
|
||||
427
backend/package/yuxi/channel/message/media.py
Normal file
427
backend/package/yuxi/channel/message/media.py
Normal file
@ -0,0 +1,427 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import io
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from email.message import EmailMessage
|
||||
from pathlib import Path, PurePosixPath
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import httpx
|
||||
from PIL import Image
|
||||
|
||||
from yuxi.channel.message.media_cleaner import MediaCleaner
|
||||
from yuxi.channel.message.media_store import MediaStore
|
||||
from yuxi.config import config
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_media_store: MediaStore | None = None
|
||||
_media_store_initialized: bool = False
|
||||
_media_cleaner: MediaCleaner | None = None
|
||||
_background_tasks: set[asyncio.Task[object]] = set()
|
||||
|
||||
MAX_IMAGE_BYTES = 6 * 1024 * 1024 # 6MB
|
||||
_MAX_STORE_CONCURRENCY = 5
|
||||
_store_semaphore = asyncio.Semaphore(_MAX_STORE_CONCURRENCY)
|
||||
|
||||
_SVG_MIME_TYPES = frozenset({"image/svg+xml", "image/svg"})
|
||||
_IMAGE_MAGIC_BYTES: dict[str, tuple[bytes, ...]] = {
|
||||
"image/jpeg": (b"\xff\xd8\xff",),
|
||||
"image/png": (b"\x89PNG\r\n\x1a\n",),
|
||||
"image/gif": (b"GIF87a", b"GIF89a"),
|
||||
"image/webp": (b"RIFF",),
|
||||
"image/bmp": (b"BM",),
|
||||
"image/tiff": (b"II*\x00", b"MM\x00*"),
|
||||
"image/x-icon": (b"\x00\x00\x01\x00",),
|
||||
}
|
||||
_SVG_SCRIPT_MARKER = b"<script"
|
||||
|
||||
_media_client: httpx.AsyncClient | None = None
|
||||
|
||||
|
||||
def _get_media_client() -> httpx.AsyncClient:
|
||||
global _media_client
|
||||
if _media_client is None:
|
||||
_media_client = httpx.AsyncClient(limits=httpx.Limits(max_connections=10, max_keepalive_connections=5))
|
||||
return _media_client
|
||||
|
||||
|
||||
async def _cleanup_media_client() -> None:
|
||||
global _media_client
|
||||
if _media_client is not None:
|
||||
await _media_client.aclose()
|
||||
_media_client = None
|
||||
|
||||
_MAX_RETRIES = 2
|
||||
_RETRY_DELAY = 1.0
|
||||
_RETRYABLE_STATUSES = frozenset({408, 429, 500, 502, 503, 504})
|
||||
|
||||
|
||||
@dataclass
|
||||
class MediaDownloadResult:
|
||||
data: bytes
|
||||
content_type: str | None
|
||||
filename: str | None
|
||||
|
||||
|
||||
def _parse_content_disposition_filename(header: str | None) -> str | None:
|
||||
if not header:
|
||||
return None
|
||||
msg = EmailMessage()
|
||||
msg["Content-Disposition"] = header
|
||||
filename = msg.get_filename()
|
||||
if filename:
|
||||
return filename
|
||||
return None
|
||||
|
||||
|
||||
def _is_transient_error(exc: Exception) -> bool:
|
||||
if isinstance(exc, httpx.HTTPStatusError):
|
||||
return exc.response.status_code in _RETRYABLE_STATUSES
|
||||
if isinstance(exc, httpx.TimeoutException):
|
||||
return True
|
||||
if isinstance(exc, httpx.NetworkError):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _resolve_filename(final_url: str, response: httpx.Response, file_path_hint: str | None) -> str | None:
|
||||
header_filename = _parse_content_disposition_filename(response.headers.get("content-disposition"))
|
||||
if header_filename:
|
||||
return header_filename
|
||||
if file_path_hint:
|
||||
return PurePosixPath(file_path_hint).name or None
|
||||
try:
|
||||
path = urlparse(final_url).path
|
||||
name = PurePosixPath(path).name
|
||||
return name or None
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
async def _read_with_limit(response: httpx.Response, max_bytes: int) -> bytes:
|
||||
chunks: list[bytes] = []
|
||||
total = 0
|
||||
async for chunk in response.aiter_bytes(chunk_size=65536):
|
||||
total += len(chunk)
|
||||
if total > max_bytes:
|
||||
raise ValueError(f"Payload exceeds maxBytes {max_bytes}")
|
||||
chunks.append(chunk)
|
||||
return b"".join(chunks)
|
||||
|
||||
|
||||
def _validate_image_magic_bytes(data: bytes, content_type: str | None) -> bool:
|
||||
ct = (content_type or "").split(";")[0].strip().lower()
|
||||
if ct in _SVG_MIME_TYPES:
|
||||
if data.lower().find(_SVG_SCRIPT_MARKER) != -1:
|
||||
logger.warning("SVG contains script tag, rejecting for XSS safety")
|
||||
return False
|
||||
return True
|
||||
expected_magics = _IMAGE_MAGIC_BYTES.get(ct)
|
||||
if expected_magics is None:
|
||||
return True
|
||||
return any(data.startswith(magic) for magic in expected_magics)
|
||||
|
||||
|
||||
async def _fetch_media(
|
||||
url: str,
|
||||
*,
|
||||
max_bytes: int,
|
||||
timeout: float,
|
||||
file_path_hint: str | None = None,
|
||||
) -> MediaDownloadResult | None:
|
||||
parsed = urlparse(url)
|
||||
if parsed.scheme not in ("http", "https"):
|
||||
logger.warning("Unsupported media URL scheme: %s", parsed.scheme)
|
||||
return None
|
||||
|
||||
for attempt in range(_MAX_RETRIES + 1):
|
||||
try:
|
||||
client = _get_media_client()
|
||||
response = await client.get(url, follow_redirects=True, timeout=timeout)
|
||||
|
||||
content_length_raw = response.headers.get("content-length")
|
||||
if content_length_raw:
|
||||
try:
|
||||
content_length = int(content_length_raw)
|
||||
if content_length > max_bytes:
|
||||
logger.error(
|
||||
"Content-Length %d exceeds maxBytes %d for url=%s",
|
||||
content_length,
|
||||
max_bytes,
|
||||
url,
|
||||
)
|
||||
return None
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
response.raise_for_status()
|
||||
|
||||
content_type = response.headers.get("content-type")
|
||||
data = await _read_with_limit(response, max_bytes)
|
||||
|
||||
if not _validate_image_magic_bytes(data, content_type):
|
||||
logger.warning(
|
||||
"Media magic bytes mismatch: content-type=%s, url=%s",
|
||||
content_type,
|
||||
url,
|
||||
)
|
||||
return None
|
||||
|
||||
filename = _resolve_filename(str(response.url), response, file_path_hint)
|
||||
|
||||
return MediaDownloadResult(
|
||||
data=data,
|
||||
content_type=content_type,
|
||||
filename=filename,
|
||||
)
|
||||
|
||||
except (httpx.HTTPStatusError, httpx.TimeoutException, httpx.NetworkError) as e:
|
||||
if attempt < _MAX_RETRIES and _is_transient_error(e):
|
||||
delay = _RETRY_DELAY * (2**attempt)
|
||||
logger.warning(
|
||||
"Transient error fetching media (attempt %d/%d), retrying in %.1fs: url=%s, err=%s",
|
||||
attempt + 1,
|
||||
_MAX_RETRIES + 1,
|
||||
delay,
|
||||
url,
|
||||
e,
|
||||
)
|
||||
await asyncio.sleep(delay)
|
||||
continue
|
||||
if isinstance(e, httpx.HTTPStatusError):
|
||||
logger.error(
|
||||
"HTTP error downloading media: url=%s, status=%d",
|
||||
url,
|
||||
e.response.status_code,
|
||||
)
|
||||
elif isinstance(e, httpx.TimeoutException):
|
||||
logger.error("Timeout downloading media: url=%s", url)
|
||||
else:
|
||||
logger.error("Network error downloading media: url=%s, err=%s", url, e)
|
||||
return None
|
||||
|
||||
except Exception:
|
||||
logger.exception("Failed to download media: url=%s", url)
|
||||
return None
|
||||
|
||||
logger.error("All retry attempts exhausted for media: url=%s", url)
|
||||
return None
|
||||
|
||||
|
||||
async def download_image_to_base64(
|
||||
url: str,
|
||||
timeout: float = 30.0,
|
||||
) -> str | None:
|
||||
result = await _fetch_media(url, max_bytes=MAX_IMAGE_BYTES, timeout=timeout)
|
||||
if result is None:
|
||||
return None
|
||||
|
||||
if not (result.content_type or "").startswith("image/"):
|
||||
logger.warning(
|
||||
"Non-image response: content-type=%s, url=%s",
|
||||
result.content_type,
|
||||
url,
|
||||
)
|
||||
return None
|
||||
|
||||
_maybe_store(result.data, result.content_type, result.filename, url)
|
||||
|
||||
return base64.b64encode(result.data).decode("utf-8")
|
||||
|
||||
|
||||
async def resolve_image_base64(media_urls: list[str]) -> str | None:
|
||||
if not media_urls:
|
||||
return None
|
||||
for url in media_urls:
|
||||
result = await download_image_to_base64(url)
|
||||
if result:
|
||||
return result
|
||||
return None
|
||||
|
||||
|
||||
def get_media_store() -> MediaStore | None:
|
||||
global _media_store, _media_store_initialized
|
||||
if _media_store_initialized:
|
||||
return _media_store
|
||||
_media_store_initialized = True
|
||||
if not config.media_store_enabled:
|
||||
return None
|
||||
media_dir = Path(config.save_dir) / "data" / "media"
|
||||
_media_store = MediaStore(
|
||||
base_dir=media_dir,
|
||||
max_total_bytes=config.media_store_max_total_bytes,
|
||||
max_file_bytes=config.media_store_max_file_bytes,
|
||||
ttl_seconds=config.media_store_ttl_seconds,
|
||||
)
|
||||
return _media_store
|
||||
|
||||
|
||||
def get_media_cleaner() -> MediaCleaner | None:
|
||||
global _media_cleaner
|
||||
store = get_media_store()
|
||||
if store is None:
|
||||
return None
|
||||
if _media_cleaner is None:
|
||||
_media_cleaner = MediaCleaner(store, config.media_cleanup_interval_seconds)
|
||||
return _media_cleaner
|
||||
|
||||
|
||||
async def start_media_cleaner() -> None:
|
||||
cleaner = get_media_cleaner()
|
||||
if cleaner is not None:
|
||||
await cleaner.start()
|
||||
|
||||
|
||||
async def stop_media_cleaner() -> None:
|
||||
cleaner = get_media_cleaner()
|
||||
if cleaner is not None:
|
||||
await cleaner.stop()
|
||||
await _cleanup_media_client()
|
||||
|
||||
|
||||
async def shutdown_media() -> None:
|
||||
await stop_media_cleaner()
|
||||
await _cancel_background_tasks()
|
||||
await _cleanup_media_client()
|
||||
|
||||
|
||||
async def _cancel_background_tasks() -> None:
|
||||
_cleanup_completed_tasks()
|
||||
pending = list(_background_tasks)
|
||||
if not pending:
|
||||
return
|
||||
for t in pending:
|
||||
t.cancel()
|
||||
results = await asyncio.gather(*pending, return_exceptions=True)
|
||||
for r in results:
|
||||
if isinstance(r, Exception) and not isinstance(r, asyncio.CancelledError):
|
||||
logger.warning("Background store task error during shutdown: %s", r)
|
||||
_background_tasks.clear()
|
||||
|
||||
|
||||
def _cleanup_completed_tasks() -> None:
|
||||
done = {t for t in _background_tasks if t.done()}
|
||||
_background_tasks.difference_update(done)
|
||||
if done:
|
||||
logger.debug("Cleaned up %d completed background store tasks, %d remaining", len(done), len(_background_tasks))
|
||||
|
||||
|
||||
async def _store_with_semaphore(store: MediaStore, data: bytes, content_type: str | None, filename: str | None, source_url: str) -> None:
|
||||
async with _store_semaphore:
|
||||
await store.store(data, content_type, filename, source_url)
|
||||
|
||||
|
||||
def _maybe_store(data: bytes, content_type: str | None, filename: str | None, source_url: str) -> None:
|
||||
store = get_media_store()
|
||||
if store is None:
|
||||
return
|
||||
_cleanup_completed_tasks()
|
||||
if len(_background_tasks) >= _MAX_STORE_CONCURRENCY * 2:
|
||||
logger.warning(
|
||||
"Background store tasks piling up: %d pending, concurrency limit=%d",
|
||||
len(_background_tasks),
|
||||
_MAX_STORE_CONCURRENCY,
|
||||
)
|
||||
task = asyncio.create_task(_store_with_semaphore(store, data, content_type, filename, source_url))
|
||||
task.add_done_callback(_handle_store_task_done)
|
||||
_background_tasks.add(task)
|
||||
|
||||
|
||||
def _handle_store_task_done(task: asyncio.Task[object]) -> None:
|
||||
try:
|
||||
task.result()
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
except Exception:
|
||||
logger.exception("Failed to store media in background task")
|
||||
finally:
|
||||
_background_tasks.discard(task)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ImageProcessResult:
|
||||
original_file_id: str | None
|
||||
thumbnail_file_id: str | None
|
||||
width: int
|
||||
height: int
|
||||
format: str
|
||||
mime_type: str
|
||||
size_bytes: int
|
||||
|
||||
|
||||
async def download_and_process_image(
|
||||
url: str,
|
||||
*,
|
||||
timeout: float = 30.0,
|
||||
thumbnail_size: tuple[int, int] | None = None,
|
||||
convert_to: str | None = None,
|
||||
) -> ImageProcessResult | None:
|
||||
result = await _fetch_media(url, max_bytes=MAX_IMAGE_BYTES, timeout=timeout)
|
||||
if result is None:
|
||||
return None
|
||||
|
||||
if not (result.content_type or "").startswith("image/"):
|
||||
logger.warning(
|
||||
"Non-image response for download_and_process_image: content-type=%s, url=%s",
|
||||
result.content_type,
|
||||
url,
|
||||
)
|
||||
return None
|
||||
|
||||
from yuxi.utils.image_processor import image_processor as img_proc
|
||||
|
||||
image_data = result.data
|
||||
if convert_to and convert_to in img_proc.CONVERTIBLE_FORMATS:
|
||||
try:
|
||||
image_data = img_proc.convert_format(result.data, convert_to)
|
||||
except Exception:
|
||||
logger.exception("Failed to convert image format to %s for url=%s", convert_to, url)
|
||||
|
||||
store = get_media_store()
|
||||
original_file_id: str | None = None
|
||||
thumbnail_file_id: str | None = None
|
||||
|
||||
if store is not None:
|
||||
original_file_id = await store.store(
|
||||
image_data,
|
||||
content_type=result.content_type,
|
||||
filename=result.filename,
|
||||
source_url=url,
|
||||
)
|
||||
|
||||
try:
|
||||
thumbnail_data = img_proc.generate_thumbnail(image_data, size=thumbnail_size)
|
||||
except Exception:
|
||||
logger.exception("Failed to generate thumbnail for url=%s", url)
|
||||
thumbnail_data = None
|
||||
|
||||
if thumbnail_data and store is not None:
|
||||
thumbnail_file_id = await store.store(
|
||||
thumbnail_data,
|
||||
content_type="image/jpeg",
|
||||
filename=None,
|
||||
source_url=url,
|
||||
)
|
||||
|
||||
try:
|
||||
with io.BytesIO(image_data) as buf:
|
||||
with Image.open(buf) as img:
|
||||
width, height = img.size
|
||||
fmt = img.format or "JPEG"
|
||||
except Exception:
|
||||
width, height = 0, 0
|
||||
fmt = "JPEG"
|
||||
|
||||
return ImageProcessResult(
|
||||
original_file_id=original_file_id,
|
||||
thumbnail_file_id=thumbnail_file_id,
|
||||
width=width,
|
||||
height=height,
|
||||
format=fmt,
|
||||
mime_type=f"image/{fmt.lower()}",
|
||||
size_bytes=len(image_data),
|
||||
)
|
||||
63
backend/package/yuxi/channel/message/media_cleaner.py
Normal file
63
backend/package/yuxi/channel/message/media_cleaner.py
Normal file
@ -0,0 +1,63 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
|
||||
from yuxi.channel.message.media_store import MediaStore
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class MediaCleaner:
|
||||
def __init__(self, media_store: MediaStore, interval_seconds: int):
|
||||
self._media_store = media_store
|
||||
self._interval = interval_seconds
|
||||
self._task: asyncio.Task | None = None
|
||||
self._stop_event = asyncio.Event()
|
||||
|
||||
@property
|
||||
def is_running(self) -> bool:
|
||||
return self._task is not None and not self._task.done()
|
||||
|
||||
async def start(self) -> None:
|
||||
if self.is_running:
|
||||
return
|
||||
self._stop_event.clear()
|
||||
self._task = asyncio.create_task(self._loop(), name="media-cleaner")
|
||||
logger.info(
|
||||
"MediaCleaner started (interval=%ds, ttl=%ds)",
|
||||
self._interval,
|
||||
self._media_store.ttl_seconds,
|
||||
)
|
||||
|
||||
async def stop(self) -> None:
|
||||
if not self.is_running:
|
||||
return
|
||||
self._stop_event.set()
|
||||
try:
|
||||
await asyncio.wait_for(self._task, timeout=10)
|
||||
except TimeoutError:
|
||||
logger.warning("MediaCleaner stop timed out, cancelling task")
|
||||
self._task.cancel()
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
self._task = None
|
||||
logger.info("MediaCleaner stopped")
|
||||
|
||||
async def cleanup_once(self) -> int:
|
||||
return await self._media_store.cleanup_expired()
|
||||
|
||||
async def _loop(self) -> None:
|
||||
while not self._stop_event.is_set():
|
||||
try:
|
||||
await asyncio.wait_for(self._stop_event.wait(), timeout=self._interval)
|
||||
break
|
||||
except TimeoutError:
|
||||
pass
|
||||
|
||||
try:
|
||||
deleted = await self._media_store.cleanup_expired()
|
||||
if deleted:
|
||||
logger.info("MediaCleaner cycle: deleted %d expired files", deleted)
|
||||
except Exception:
|
||||
logger.exception("MediaCleaner cycle error")
|
||||
242
backend/package/yuxi/channel/message/media_store.py
Normal file
242
backend/package/yuxi/channel/message/media_store.py
Normal file
@ -0,0 +1,242 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import uuid
|
||||
from datetime import UTC, datetime
|
||||
from pathlib import Path
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_META_SUFFIX = ".meta.json"
|
||||
|
||||
|
||||
def _now_iso() -> str:
|
||||
return datetime.now(UTC).isoformat()
|
||||
|
||||
|
||||
def _date_dir_name() -> str:
|
||||
return datetime.now(UTC).strftime("%Y-%m-%d")
|
||||
|
||||
|
||||
class MediaStore:
|
||||
def __init__(
|
||||
self,
|
||||
base_dir: Path,
|
||||
max_total_bytes: int,
|
||||
max_file_bytes: int,
|
||||
ttl_seconds: int,
|
||||
):
|
||||
self.base_dir = base_dir
|
||||
self.max_total_bytes = max_total_bytes
|
||||
self.max_file_bytes = max_file_bytes
|
||||
self.ttl_seconds = ttl_seconds
|
||||
self._size_initialized = False
|
||||
self._running_total: int = 0
|
||||
self._size_lock = asyncio.Lock()
|
||||
|
||||
def _ensure_dir(self, dir_path: Path) -> None:
|
||||
dir_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def _make_file_id(self, date_dir: str, uid: str, ext: str) -> str:
|
||||
return f"{date_dir}/{uid}{ext}"
|
||||
|
||||
def _get_file_path(self, file_id: str) -> Path:
|
||||
return self.base_dir / file_id
|
||||
|
||||
def _get_meta_path(self, file_id: str) -> Path:
|
||||
return self.base_dir / f"{file_id}{_META_SUFFIX}"
|
||||
|
||||
def _dir_size(self, dir_path: Path) -> int:
|
||||
total = 0
|
||||
try:
|
||||
for entry in dir_path.rglob("*"):
|
||||
if entry.is_file():
|
||||
try:
|
||||
total += entry.stat().st_size
|
||||
except OSError:
|
||||
logger.debug("Failed to stat file during size calculation: %s", entry)
|
||||
except OSError:
|
||||
logger.debug("Failed to calculate directory size: %s", dir_path)
|
||||
return total
|
||||
|
||||
async def _check_quota(self, incoming_bytes: int) -> bool:
|
||||
if incoming_bytes > self.max_file_bytes:
|
||||
logger.warning(
|
||||
"File size %d exceeds per-file limit %d",
|
||||
incoming_bytes,
|
||||
self.max_file_bytes,
|
||||
)
|
||||
return False
|
||||
current_total = await self._get_running_total()
|
||||
if current_total + incoming_bytes > self.max_total_bytes:
|
||||
logger.warning(
|
||||
"Storage quota exceeded: current=%d, incoming=%d, max=%d",
|
||||
current_total,
|
||||
incoming_bytes,
|
||||
self.max_total_bytes,
|
||||
)
|
||||
return False
|
||||
return True
|
||||
|
||||
async def _get_running_total(self) -> int:
|
||||
async with self._size_lock:
|
||||
if not self._size_initialized:
|
||||
self._running_total = self._dir_size(self.base_dir)
|
||||
self._size_initialized = True
|
||||
return self._running_total
|
||||
|
||||
def _infer_ext(self, content_type: str | None, filename: str | None) -> str:
|
||||
if filename:
|
||||
suffix = Path(filename).suffix
|
||||
if suffix:
|
||||
return suffix.lower()
|
||||
if content_type:
|
||||
ct = content_type.split(";")[0].strip().lower()
|
||||
mapping = {
|
||||
"image/jpeg": ".jpg",
|
||||
"image/png": ".png",
|
||||
"image/gif": ".gif",
|
||||
"image/webp": ".webp",
|
||||
"image/bmp": ".bmp",
|
||||
"image/svg+xml": ".svg",
|
||||
"audio/mpeg": ".mp3",
|
||||
"audio/wav": ".wav",
|
||||
"audio/ogg": ".ogg",
|
||||
"video/mp4": ".mp4",
|
||||
"application/pdf": ".pdf",
|
||||
}
|
||||
if ct in mapping:
|
||||
return mapping[ct]
|
||||
return ".bin"
|
||||
|
||||
async def store(
|
||||
self,
|
||||
data: bytes,
|
||||
content_type: str | None = None,
|
||||
filename: str | None = None,
|
||||
source_url: str | None = None,
|
||||
) -> str | None:
|
||||
if not await self._check_quota(len(data)):
|
||||
return None
|
||||
|
||||
date_dir = _date_dir_name()
|
||||
uid = uuid.uuid4().hex
|
||||
ext = self._infer_ext(content_type, filename)
|
||||
file_id = self._make_file_id(date_dir, uid, ext)
|
||||
|
||||
file_path = self._get_file_path(file_id)
|
||||
meta_path = self._get_meta_path(file_id)
|
||||
|
||||
self._ensure_dir(file_path.parent)
|
||||
|
||||
try:
|
||||
file_path.write_bytes(data)
|
||||
except OSError:
|
||||
logger.exception("Failed to write media file: %s", file_path)
|
||||
return None
|
||||
|
||||
meta = {
|
||||
"original_filename": filename,
|
||||
"content_type": content_type,
|
||||
"source_url": source_url,
|
||||
"stored_at": _now_iso(),
|
||||
"ttl_seconds": self.ttl_seconds,
|
||||
"size": len(data),
|
||||
}
|
||||
try:
|
||||
meta_path.write_text(json.dumps(meta, ensure_ascii=False), encoding="utf-8")
|
||||
except OSError:
|
||||
logger.exception("Failed to write media metadata: %s", meta_path)
|
||||
try:
|
||||
file_path.unlink(missing_ok=True)
|
||||
except OSError:
|
||||
logger.debug("Failed to clean up orphaned media file: %s", file_path)
|
||||
return None
|
||||
|
||||
logger.info("Stored media: file_id=%s, size=%d", file_id, len(data))
|
||||
async with self._size_lock:
|
||||
self._running_total += len(data)
|
||||
return file_id
|
||||
|
||||
def resolve_path(self, file_id: str) -> Path | None:
|
||||
file_path = self._get_file_path(file_id)
|
||||
if file_path.exists():
|
||||
return file_path
|
||||
return None
|
||||
|
||||
def read_meta(self, file_id: str) -> dict | None:
|
||||
meta_path = self._get_meta_path(file_id)
|
||||
try:
|
||||
return json.loads(meta_path.read_text(encoding="utf-8"))
|
||||
except (OSError, json.JSONDecodeError):
|
||||
return None
|
||||
|
||||
async def delete(self, file_id: str) -> bool:
|
||||
file_path = self._get_file_path(file_id)
|
||||
meta_path = self._get_meta_path(file_id)
|
||||
deleted = False
|
||||
for p in (file_path, meta_path):
|
||||
if not p.exists():
|
||||
continue
|
||||
try:
|
||||
p.unlink()
|
||||
deleted = True
|
||||
except OSError:
|
||||
logger.debug("Failed to delete file during media cleanup: %s", p)
|
||||
if deleted:
|
||||
logger.info("Deleted media: file_id=%s", file_id)
|
||||
return deleted
|
||||
|
||||
def get_total_size(self) -> int:
|
||||
return self._dir_size(self.base_dir)
|
||||
|
||||
async def cleanup_expired(self) -> int:
|
||||
deleted_count = 0
|
||||
now = datetime.now(UTC)
|
||||
try:
|
||||
for entry in sorted(self.base_dir.rglob(f"*{_META_SUFFIX}")):
|
||||
try:
|
||||
meta = json.loads(entry.read_text(encoding="utf-8"))
|
||||
except (OSError, json.JSONDecodeError):
|
||||
continue
|
||||
|
||||
stored_at_str = meta.get("stored_at")
|
||||
ttl_s = meta.get("ttl_seconds", self.ttl_seconds)
|
||||
if not stored_at_str:
|
||||
continue
|
||||
|
||||
try:
|
||||
stored_at = datetime.fromisoformat(stored_at_str)
|
||||
except ValueError:
|
||||
continue
|
||||
|
||||
if (now - stored_at).total_seconds() <= ttl_s:
|
||||
continue
|
||||
|
||||
file_id = str(entry.relative_to(self.base_dir)).removesuffix(_META_SUFFIX)
|
||||
if await self.delete(file_id):
|
||||
deleted_count += 1
|
||||
|
||||
self._remove_empty_dirs()
|
||||
except OSError:
|
||||
logger.exception("Error during media cleanup scan")
|
||||
|
||||
if deleted_count:
|
||||
logger.info("Media cleanup: deleted %d expired files", deleted_count)
|
||||
return deleted_count
|
||||
|
||||
def _remove_empty_dirs(self) -> None:
|
||||
try:
|
||||
for dirpath, dirnames, filenames in os.walk(self.base_dir, topdown=False):
|
||||
if dirpath == str(self.base_dir):
|
||||
continue
|
||||
if not dirnames and not filenames:
|
||||
try:
|
||||
os.rmdir(dirpath)
|
||||
except OSError:
|
||||
logger.debug("Failed to remove empty directory: %s", dirpath)
|
||||
except OSError:
|
||||
logger.debug("Failed to walk directory for cleanup: %s", self.base_dir)
|
||||
211
backend/package/yuxi/channel/message/metrics.py
Normal file
211
backend/package/yuxi/channel/message/metrics.py
Normal file
@ -0,0 +1,211 @@
|
||||
import time
|
||||
from collections import defaultdict
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
from enum import StrEnum
|
||||
|
||||
|
||||
class MetricType(StrEnum):
|
||||
COUNTER = "counter"
|
||||
GAUGE = "gauge"
|
||||
HISTOGRAM = "histogram"
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class _Counter:
|
||||
value: int = 0
|
||||
|
||||
def inc(self, amount: int = 1) -> None:
|
||||
self.value += amount
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class _Gauge:
|
||||
value: float = 0.0
|
||||
|
||||
def set(self, value: float) -> None:
|
||||
self.value = value
|
||||
|
||||
def inc(self, amount: float = 1.0) -> None:
|
||||
self.value += amount
|
||||
|
||||
def dec(self, amount: float = 1.0) -> None:
|
||||
self.value -= amount
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class _Histogram:
|
||||
buckets: list[float]
|
||||
values: list[int] = field(default_factory=list)
|
||||
_sum: float = 0.0
|
||||
_count: int = 0
|
||||
|
||||
def observe(self, value: float) -> None:
|
||||
self._sum += value
|
||||
self._count += 1
|
||||
while len(self.values) < len(self.buckets):
|
||||
self.values.append(0)
|
||||
for i, bound in enumerate(self.buckets):
|
||||
if value <= bound:
|
||||
self.values[i] += 1
|
||||
return
|
||||
|
||||
def quantile(self, q: float) -> float | None:
|
||||
if self._count == 0:
|
||||
return None
|
||||
if not self.values:
|
||||
return None
|
||||
|
||||
target_rank = q * self._count
|
||||
|
||||
cumulative = 0
|
||||
for i, (bound, count) in enumerate(zip(self.buckets, self.values)):
|
||||
cumulative += count
|
||||
if cumulative >= target_rank:
|
||||
if i == 0:
|
||||
lower_bound = 0.0
|
||||
else:
|
||||
lower_bound = self.buckets[i - 1]
|
||||
upper_bound = bound
|
||||
prev_cumulative = cumulative - count
|
||||
fraction = (target_rank - prev_cumulative) / max(count, 1)
|
||||
return lower_bound + fraction * (upper_bound - lower_bound)
|
||||
|
||||
return float(self.buckets[-1]) if self.buckets else None
|
||||
|
||||
|
||||
class MetricsRegistry:
|
||||
def __init__(self) -> None:
|
||||
self._label_values: dict[str, dict[tuple[str, ...], object]] = defaultdict(dict)
|
||||
|
||||
def counter(self, name: str, label_keys: tuple[str, ...] = ()) -> Callable[..., None]:
|
||||
def inc(labels: dict[str, str] | None = None, amount: int = 1) -> None:
|
||||
key = self._resolve_key(labels, label_keys)
|
||||
counter = self._label_values[name].get(key)
|
||||
if counter is None:
|
||||
counter = _Counter()
|
||||
self._label_values[name][key] = counter
|
||||
counter.inc(amount)
|
||||
|
||||
return inc
|
||||
|
||||
def gauge(self, name: str, label_keys: tuple[str, ...] = ()) -> Callable[..., None]:
|
||||
def set(value: float, labels: dict[str, str] | None = None) -> None:
|
||||
key = self._resolve_key(labels, label_keys)
|
||||
gauge = self._label_values[name].get(key)
|
||||
if gauge is None:
|
||||
gauge = _Gauge()
|
||||
self._label_values[name][key] = gauge
|
||||
gauge.set(value)
|
||||
|
||||
return set
|
||||
|
||||
def histogram(self, name: str, buckets: list[float], label_keys: tuple[str, ...] = ()) -> Callable[..., None]:
|
||||
def observe(value: float, labels: dict[str, str] | None = None) -> None:
|
||||
key = self._resolve_key(labels, label_keys)
|
||||
hist = self._label_values[name].get(key)
|
||||
if hist is None:
|
||||
hist = _Histogram(buckets=buckets)
|
||||
self._label_values[name][key] = hist
|
||||
hist.observe(value)
|
||||
|
||||
return observe
|
||||
|
||||
@staticmethod
|
||||
def _resolve_key(labels: dict[str, str] | None, label_keys: tuple[str, ...]) -> tuple[str, ...]:
|
||||
if not labels:
|
||||
return tuple("" for _ in label_keys)
|
||||
return tuple(labels.get(k, "") for k in label_keys)
|
||||
|
||||
def snapshot(self) -> dict:
|
||||
result: dict = {}
|
||||
for metric_name, entries in self._label_values.items():
|
||||
result[metric_name] = {}
|
||||
for label_tuple, metric in entries.items():
|
||||
if isinstance(metric, _Counter):
|
||||
result[metric_name][str(label_tuple)] = {"type": "counter", "value": metric.value}
|
||||
elif isinstance(metric, _Gauge):
|
||||
result[metric_name][str(label_tuple)] = {"type": "gauge", "value": metric.value}
|
||||
elif isinstance(metric, _Histogram):
|
||||
result[metric_name][str(label_tuple)] = {
|
||||
"type": "histogram",
|
||||
"buckets": metric.buckets,
|
||||
"values": metric.values,
|
||||
"sum": metric._sum,
|
||||
"count": metric._count,
|
||||
"p50": metric.quantile(0.50),
|
||||
"p95": metric.quantile(0.95),
|
||||
"p99": metric.quantile(0.99),
|
||||
}
|
||||
return result
|
||||
|
||||
|
||||
registry = MetricsRegistry()
|
||||
|
||||
channel_messages_total = registry.counter(
|
||||
"channel_messages_total",
|
||||
("channel_type", "status"),
|
||||
)
|
||||
|
||||
channel_dispatch_duration_ms = registry.histogram(
|
||||
"channel_dispatch_duration_ms",
|
||||
buckets=[5, 25, 50, 100, 250, 500, 1000, 2500, 5000],
|
||||
label_keys=("channel_type",),
|
||||
)
|
||||
|
||||
channel_agent_duration_ms = registry.histogram(
|
||||
"channel_agent_duration_ms",
|
||||
buckets=[100, 500, 1000, 2500, 5000, 10000, 30000, 60000],
|
||||
label_keys=("channel_type",),
|
||||
)
|
||||
|
||||
channel_rate_limit_rejects_total = registry.counter(
|
||||
"channel_rate_limit_rejects_total",
|
||||
("channel_type",),
|
||||
)
|
||||
|
||||
channel_messages_inflight = registry.gauge(
|
||||
"channel_messages_inflight",
|
||||
)
|
||||
|
||||
|
||||
def record_message(channel_type: str, status: str) -> None:
|
||||
channel_messages_total(labels={"channel_type": channel_type, "status": status})
|
||||
|
||||
|
||||
def record_dispatch_duration_ms(channel_type: str, duration_ms: float) -> None:
|
||||
channel_dispatch_duration_ms(duration_ms, labels={"channel_type": channel_type})
|
||||
|
||||
|
||||
def record_agent_duration_ms(channel_type: str, duration_ms: float) -> None:
|
||||
channel_agent_duration_ms(duration_ms, labels={"channel_type": channel_type})
|
||||
|
||||
|
||||
def record_rate_limit_reject(channel_type: str) -> None:
|
||||
channel_rate_limit_rejects_total(labels={"channel_type": channel_type})
|
||||
|
||||
|
||||
def set_inflight(count: int) -> None:
|
||||
channel_messages_inflight(count)
|
||||
|
||||
|
||||
class MetricsTimer:
|
||||
def __init__(self, on_finish: Callable[[float], None]) -> None:
|
||||
self._on_finish = on_finish
|
||||
self._start = 0.0
|
||||
|
||||
def __enter__(self) -> "MetricsTimer":
|
||||
self._start = time.monotonic()
|
||||
return self
|
||||
|
||||
def __exit__(self, *args) -> None:
|
||||
elapsed = (time.monotonic() - self._start) * 1000
|
||||
self._on_finish(elapsed)
|
||||
|
||||
async def __aenter__(self) -> "MetricsTimer":
|
||||
self._start = time.monotonic()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args) -> None:
|
||||
elapsed = (time.monotonic() - self._start) * 1000
|
||||
self._on_finish(elapsed)
|
||||
283
backend/package/yuxi/channel/message/models.py
Normal file
283
backend/package/yuxi/channel/message/models.py
Normal file
@ -0,0 +1,283 @@
|
||||
import hashlib
|
||||
import json
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from enum import StrEnum
|
||||
from typing import Any
|
||||
|
||||
from yuxi.channel.routing.models import PeerKind
|
||||
|
||||
|
||||
class MessageType(StrEnum):
|
||||
TEXT = "text"
|
||||
IMAGE = "image"
|
||||
VOICE = "voice"
|
||||
FILE = "file"
|
||||
EVENT = "event"
|
||||
|
||||
|
||||
class ReplyStage(StrEnum):
|
||||
TOOL = "tool"
|
||||
BLOCK = "block"
|
||||
FINAL = "final"
|
||||
|
||||
|
||||
class ChunkType(StrEnum):
|
||||
TEXT_DELTA = "text-delta"
|
||||
REASONING = "reasoning"
|
||||
TOOL_CALL = "tool-call"
|
||||
TOOL_RESULT = "tool-result"
|
||||
HEARTBEAT = "heartbeat"
|
||||
CITATION = "citation"
|
||||
STATUS = "status"
|
||||
ERROR = "error"
|
||||
|
||||
|
||||
class MentionSource(StrEnum):
|
||||
EXPLICIT_BOT = "explicit_bot"
|
||||
SUBTEAM = "subteam"
|
||||
MENTION_PATTERN = "mention_pattern"
|
||||
IMPLICIT_THREAD = "implicit_thread"
|
||||
COMMAND_BYPASS = "command_bypass"
|
||||
NONE = "none"
|
||||
|
||||
|
||||
@dataclass
|
||||
class PeerInfo:
|
||||
kind: PeerKind
|
||||
id: str
|
||||
display_name: str | None = None
|
||||
username: str | None = None
|
||||
tag: str | None = None
|
||||
is_bot: bool = False
|
||||
is_self: bool = False
|
||||
display_label: str | None = None
|
||||
roles: list[str] = field(default_factory=list)
|
||||
|
||||
|
||||
@dataclass
|
||||
class GroupContext:
|
||||
id: str | None = None
|
||||
name: str | None = None
|
||||
guild_id: str | None = None
|
||||
team_id: str | None = None
|
||||
thread_id: str | None = None
|
||||
roles: list[str] = field(default_factory=list)
|
||||
kind: str | None = None
|
||||
space_id: str | None = None
|
||||
parent_id: str | None = None
|
||||
native_channel_id: str | None = None
|
||||
route_peer_kind: str | None = None
|
||||
route_peer_id: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class UnifiedMessage:
|
||||
msg_id: str
|
||||
channel_type: str
|
||||
account_id: str
|
||||
content: str
|
||||
sender: PeerInfo
|
||||
channel_config_id: str | None = None
|
||||
message_type: MessageType = MessageType.TEXT
|
||||
media_urls: list[str] = field(default_factory=list)
|
||||
image_base64: str | None = None
|
||||
group: GroupContext | None = None
|
||||
timestamp: datetime | None = None
|
||||
raw_payload: dict[str, Any] = field(default_factory=dict)
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
body_for_agent: str | None = None
|
||||
command_body: str | None = None
|
||||
body_for_commands: str | None = None
|
||||
inbound_history: list[dict[str, Any]] = field(default_factory=list)
|
||||
|
||||
reply_chain: list[dict[str, Any]] = field(default_factory=list)
|
||||
reply_to_id: str | None = None
|
||||
reply_to_id_full: str | None = None
|
||||
root_message_id: str | None = None
|
||||
message_thread_id: str | None = None
|
||||
reply_to_tag: bool = False
|
||||
reply_to_current: bool = False
|
||||
|
||||
forwarded_from: str | None = None
|
||||
forwarded_from_type: str | None = None
|
||||
forwarded_from_id: str | None = None
|
||||
forwarded_from_username: str | None = None
|
||||
forwarded_date: float | None = None
|
||||
|
||||
media_path: str | None = None
|
||||
media_paths: list[str] = field(default_factory=list)
|
||||
media_types: list[str] = field(default_factory=list)
|
||||
|
||||
conversation_label: str | None = None
|
||||
group_channel: str | None = None
|
||||
group_space: str | None = None
|
||||
group_members: str | None = None
|
||||
group_system_prompt: str | None = None
|
||||
member_role_ids: list[str] = field(default_factory=list)
|
||||
|
||||
was_mentioned: bool = False
|
||||
explicitly_mentioned_bot: bool = False
|
||||
mentioned_user_ids: list[str] = field(default_factory=list)
|
||||
mention_source: MentionSource | None = None
|
||||
|
||||
surface: str | None = None
|
||||
originating_channel: str | None = None
|
||||
originating_to: str | None = None
|
||||
native_channel_id: str | None = None
|
||||
native_direct_user_id: str | None = None
|
||||
thread_parent_id: str | None = None
|
||||
|
||||
location_lat: float | None = None
|
||||
location_lon: float | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class DispatchResult:
|
||||
success: bool
|
||||
thread_id: str | None = None
|
||||
agent_config_id: int | None = None
|
||||
error: str | None = None
|
||||
session_key: str | None = None
|
||||
matched_by: str | None = None
|
||||
handled_by: str | None = None
|
||||
command_response: str | None = None
|
||||
internal_user_id: str | None = None
|
||||
reply_sent: bool = False
|
||||
explicit_target: str | None = None
|
||||
queued_final: bool = False
|
||||
counts: dict[str, int] = field(default_factory=dict)
|
||||
failed_counts: dict[str, int] = field(default_factory=dict)
|
||||
source_reply_delivery_mode: str | None = None
|
||||
before_agent_run_blocked: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class StreamingChunk:
|
||||
"""Agent 推理流式块"""
|
||||
|
||||
stage: ReplyStage = ReplyStage.BLOCK
|
||||
chunk_type: ChunkType = ChunkType.TEXT_DELTA
|
||||
content: str | None = None
|
||||
tool_name: str | None = None
|
||||
tool_input: dict[str, Any] | None = None
|
||||
tool_output: str | None = None
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
agent_run_id: str | None = None
|
||||
is_final: bool = False
|
||||
item_id: str | None = None
|
||||
title: str | None = None
|
||||
phase: str | None = None
|
||||
steps: list[str] | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class ReplyPayload:
|
||||
"""出站回复载荷(经过去重和合并后)"""
|
||||
|
||||
target_id: str
|
||||
content: str
|
||||
reply_to_id: str | None = None
|
||||
thread_id: str | None = None
|
||||
media_urls: list[str] = field(default_factory=list)
|
||||
presentation: dict[str, Any] | None = None
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
btw: dict[str, str] | None = None
|
||||
audio_as_voice: bool = False
|
||||
spoken_text: str | None = None
|
||||
is_error: bool = False
|
||||
is_reasoning: bool = False
|
||||
is_compaction_notice: bool = False
|
||||
channel_data: dict[str, Any] | None = None
|
||||
trusted_local_media: bool = False
|
||||
sensitive_media: bool = False
|
||||
delivery: dict[str, Any] | None = None
|
||||
reply_to_tag: bool = False
|
||||
reply_to_current: bool = False
|
||||
|
||||
@property
|
||||
def payload_key(self) -> str:
|
||||
key = f"{self.content}|{','.join(sorted(self.media_urls))}|{self.reply_to_id or ''}"
|
||||
if self.presentation:
|
||||
stable_hash = hashlib.md5(
|
||||
json.dumps(self.presentation, sort_keys=True, ensure_ascii=False).encode()
|
||||
).hexdigest()
|
||||
key += f"|{stable_hash}"
|
||||
return key
|
||||
|
||||
@property
|
||||
def content_key(self) -> str:
|
||||
return f"{self.content}|{','.join(sorted(self.media_urls))}"
|
||||
|
||||
|
||||
class StreamStatus(StrEnum):
|
||||
"""流执行终止状态"""
|
||||
|
||||
COMPLETED = "completed"
|
||||
CANCELLED = "cancelled"
|
||||
ERROR = "error"
|
||||
TIMEOUT = "timeout"
|
||||
|
||||
|
||||
@dataclass
|
||||
class StreamResult:
|
||||
"""流执行完成后的结构化结果,对应 openclaw 的 MessageReceipt 概念"""
|
||||
|
||||
status: StreamStatus = StreamStatus.COMPLETED
|
||||
accumulated_text: str = ""
|
||||
chunk_count: int = 0
|
||||
text_chunk_count: int = 0
|
||||
tool_chunk_count: int = 0
|
||||
reasoning_chunk_count: int = 0
|
||||
heartbeat_count: int = 0
|
||||
citation_count: int = 0
|
||||
error_count: int = 0
|
||||
start_time: float = 0.0
|
||||
end_time: float = 0.0
|
||||
token_usage: dict[str, int] | None = None
|
||||
error_message: str | None = None
|
||||
|
||||
@property
|
||||
def elapsed_ms(self) -> float:
|
||||
return (self.end_time - self.start_time) * 1000 if self.start_time > 0 else 0.0
|
||||
|
||||
@property
|
||||
def is_success(self) -> bool:
|
||||
return self.status == StreamStatus.COMPLETED
|
||||
|
||||
|
||||
class StreamCancellationToken:
|
||||
"""流取消令牌,对应 openclaw 的 AbortSignal"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._cancelled = False
|
||||
|
||||
def cancel(self) -> None:
|
||||
self._cancelled = True
|
||||
|
||||
@property
|
||||
def is_cancelled(self) -> bool:
|
||||
return self._cancelled
|
||||
|
||||
def throw_if_cancelled(self) -> None:
|
||||
if self._cancelled:
|
||||
raise StreamCancelledError()
|
||||
|
||||
|
||||
class StreamCancelledError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class StreamTimeoutError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class MessageReceipt:
|
||||
platform_message_id: str
|
||||
thread_id: str | None = None
|
||||
reply_to_id: str | None = None
|
||||
sent_at: float | None = None
|
||||
parts: list[dict[str, Any]] = field(default_factory=list)
|
||||
raw: list[dict[str, Any]] = field(default_factory=list)
|
||||
685
backend/package/yuxi/channel/message/processor.py
Normal file
685
backend/package/yuxi/channel/message/processor.py
Normal file
@ -0,0 +1,685 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import AsyncIterator, Callable
|
||||
|
||||
from sqlalchemy import select
|
||||
|
||||
from yuxi.channel.hooks.lifecycle import HookEvent, LifecycleHookRegistry
|
||||
from yuxi.channel.message.block_reply_pipeline import create_block_reply_pipeline
|
||||
from yuxi.channel.message.bridge import AgentBridge
|
||||
from yuxi.channel.message.conversation_fence import ConversationFence
|
||||
from yuxi.channel.message.dispatch import MessageDispatcher
|
||||
from yuxi.channel.message.idempotency import claim, release
|
||||
from yuxi.channel.message.langfuse_trace import ChannelTrace, create_channel_trace
|
||||
from yuxi.channel.message.metrics import (
|
||||
record_dispatch_duration_ms,
|
||||
record_message,
|
||||
record_rate_limit_reject,
|
||||
set_inflight,
|
||||
)
|
||||
from yuxi.channel.message.models import (
|
||||
PeerKind,
|
||||
StreamCancelledError,
|
||||
StreamStatus,
|
||||
StreamTimeoutError,
|
||||
UnifiedMessage,
|
||||
)
|
||||
from yuxi.channel.message.rate_limiter import rate_limit_manager
|
||||
from yuxi.channel.protocols import (
|
||||
InboundMessageHook,
|
||||
MessageSendingHook,
|
||||
OutboundProtocol,
|
||||
SendLifecycleHook,
|
||||
)
|
||||
from yuxi.channel.plugins.registry import ChannelPluginRegistry
|
||||
from yuxi.channel.routing.matcher import RouteMatcher
|
||||
from yuxi.channel.routing.models import RouteBinding
|
||||
from yuxi.channel.routing.session_key import SessionKeyBuilder
|
||||
from yuxi.channel.security.allowlist import AllowlistChecker
|
||||
from yuxi.channel.security.identity_link import IdentityLinkResolver
|
||||
from yuxi.channel.security.pairing import InMemoryPairingStore, PairingManager
|
||||
from yuxi.channel.message.circuit_breaker import CircuitBreaker
|
||||
from yuxi.repositories.channel_binding_repo import ChannelBindingRepository, to_route_binding
|
||||
from yuxi.repositories.channel_msg_record_repo import ChannelMsgRecordRepository
|
||||
from yuxi.repositories.channel_thread_mapping_repo import ChannelThreadMappingRepository
|
||||
from yuxi.repositories.channel_user_mapping_repo import ChannelUserMappingRepository
|
||||
from yuxi.storage.postgres.manager import pg_manager
|
||||
from yuxi.storage.postgres.models_business import User
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_BINDINGS_CACHE_TTL = 60.0
|
||||
_IDEMPOTENCY_TTL = 300
|
||||
|
||||
_FALLBACK_MESSAGES: dict[str, str] = {
|
||||
"no_route": "暂不支持该渠道的消息处理。",
|
||||
"route_error": "消息路由异常,请稍后重试。",
|
||||
"dispatch_error": "消息处理异常,请稍后重试。",
|
||||
"approval_denied": "操作未通过审批。",
|
||||
"approval_check_error": "审批校验异常,请稍后重试。",
|
||||
"rate_limited": "当前消息量较大,请稍后重试。",
|
||||
}
|
||||
|
||||
|
||||
class MessageProcessor:
|
||||
def __init__(
|
||||
self,
|
||||
stream_fn: Callable[..., AsyncIterator[bytes]],
|
||||
channel_manager=None,
|
||||
*,
|
||||
msg_repo=None,
|
||||
binding_repo=None,
|
||||
user_mapping_repo=None,
|
||||
thread_mapping_repo=None,
|
||||
allowlist=None,
|
||||
pairing=None,
|
||||
identity_link_resolver=None,
|
||||
response_prefix: str = "",
|
||||
human_delay: tuple[float, float] | None = None,
|
||||
):
|
||||
self._bridge = AgentBridge(stream_fn, channel_manager)
|
||||
self._msg_repo = msg_repo or ChannelMsgRecordRepository()
|
||||
self._binding_repo = binding_repo or ChannelBindingRepository()
|
||||
self._user_mapping_repo = user_mapping_repo or ChannelUserMappingRepository()
|
||||
self._thread_mapping_repo = thread_mapping_repo or ChannelThreadMappingRepository()
|
||||
self._hooks = LifecycleHookRegistry()
|
||||
self._allowlist = allowlist or AllowlistChecker()
|
||||
if pairing:
|
||||
self._pairing = pairing
|
||||
else:
|
||||
self._pairing = PairingManager(InMemoryPairingStore(), allowlist=self._allowlist)
|
||||
|
||||
if allowlist is None:
|
||||
self._allowlist.load_dm_allowlist(["*"])
|
||||
self._allowlist.load_group_allowlist(["*"])
|
||||
|
||||
self._response_prefix = response_prefix
|
||||
self._human_delay = human_delay
|
||||
|
||||
self._identity_links = identity_link_resolver or IdentityLinkResolver()
|
||||
self._matcher = RouteMatcher()
|
||||
self._session_key_builder = SessionKeyBuilder()
|
||||
|
||||
self._active_message_count = 0
|
||||
self._active_count_lock = asyncio.Lock()
|
||||
|
||||
self._cached_bindings: list[RouteBinding] = []
|
||||
self._bindings_ts: float = 0.0
|
||||
|
||||
self._fence = ConversationFence()
|
||||
|
||||
self._self_user_ids: dict[str, set[str]] = {}
|
||||
|
||||
def set_self_user_ids(self, channel_type: str, account_id: str, user_ids: set[str]) -> None:
|
||||
key = f"{channel_type}:{account_id}"
|
||||
self._self_user_ids[key] = user_ids
|
||||
|
||||
@property
|
||||
def hooks(self) -> LifecycleHookRegistry:
|
||||
return self._hooks
|
||||
|
||||
async def _get_bindings(self) -> list[RouteBinding]:
|
||||
now = time.monotonic()
|
||||
if now - self._bindings_ts < _BINDINGS_CACHE_TTL and self._cached_bindings:
|
||||
return self._cached_bindings
|
||||
db_bindings = await self._binding_repo.list_all()
|
||||
self._cached_bindings = [to_route_binding(b) for b in db_bindings]
|
||||
self._bindings_ts = now
|
||||
await self._matcher.set_bindings(self._cached_bindings)
|
||||
return self._cached_bindings
|
||||
|
||||
async def _get_user_for_channel(self, internal_user_id: str | None) -> User | None:
|
||||
if not internal_user_id:
|
||||
return None
|
||||
try:
|
||||
uid = int(internal_user_id)
|
||||
except (ValueError, TypeError):
|
||||
return None
|
||||
async with pg_manager.get_async_session_context() as db:
|
||||
return (await db.execute(select(User).where(User.id == uid))).scalar_one_or_none()
|
||||
|
||||
def _is_echo(self, msg: UnifiedMessage) -> bool:
|
||||
key = f"{msg.channel_type}:{msg.account_id}"
|
||||
self_ids = self._self_user_ids.get(key)
|
||||
if not self_ids:
|
||||
return False
|
||||
return msg.sender.id in self_ids
|
||||
|
||||
async def _ensure_security_loaded(self) -> None:
|
||||
await self._identity_links.ensure_loaded()
|
||||
self._session_key_builder = SessionKeyBuilder(
|
||||
identity_links=self._identity_links.list_links()
|
||||
)
|
||||
|
||||
async def process(self, msg: UnifiedMessage) -> None:
|
||||
if self._is_echo(msg):
|
||||
logger.debug("Skipping echo message: %s/%s", msg.channel_type, msg.msg_id)
|
||||
return
|
||||
|
||||
limiter = await rate_limit_manager.get_limiter(msg.channel_type, msg.account_id)
|
||||
acquire_result = await limiter.try_acquire()
|
||||
if not acquire_result:
|
||||
record_rate_limit_reject(msg.channel_type)
|
||||
logger.warning(
|
||||
"Rate limited: channel=%s account=%s retry_after=%.2fs",
|
||||
msg.channel_type, msg.account_id, acquire_result.retry_after_sec,
|
||||
)
|
||||
await self._send_fallback(msg, _FALLBACK_MESSAGES["rate_limited"])
|
||||
return
|
||||
|
||||
if not await claim(msg.msg_id, ttl=_IDEMPOTENCY_TTL):
|
||||
logger.debug("Duplicate/inflight message ignored: %s", msg.msg_id)
|
||||
return
|
||||
|
||||
try:
|
||||
async with self._active_count_lock:
|
||||
self._active_message_count += 1
|
||||
set_inflight(self._active_message_count)
|
||||
await self._ensure_security_loaded()
|
||||
await self._process_impl(msg, limiter)
|
||||
finally:
|
||||
await release(msg.msg_id)
|
||||
async with self._active_count_lock:
|
||||
self._active_message_count = max(0, self._active_message_count - 1)
|
||||
set_inflight(self._active_message_count)
|
||||
|
||||
async def _process_impl(self, msg: UnifiedMessage, limiter) -> None:
|
||||
trace = await create_channel_trace(msg)
|
||||
await trace.add_span("idempotency_check", metadata={"claimed": True})
|
||||
|
||||
conversation_key = self._fence.key_for(msg)
|
||||
version, stop_event = self._fence.enter(conversation_key)
|
||||
lock = self._fence.lock_for(conversation_key)
|
||||
|
||||
logger.debug(
|
||||
"processor: phase=enter msg_id=%s channel=%s conv_key=%s version=%d",
|
||||
msg.msg_id, msg.channel_type, conversation_key, version,
|
||||
)
|
||||
|
||||
async with lock:
|
||||
if version != self._fence.current_version(conversation_key):
|
||||
logger.debug("Foreground fence: message superseded for %s", conversation_key)
|
||||
record_message(msg.channel_type, "superseded")
|
||||
await trace.finish(error="superseded")
|
||||
return
|
||||
|
||||
record_message(msg.channel_type, "received")
|
||||
|
||||
plugin = ChannelPluginRegistry.get(msg.channel_type)
|
||||
|
||||
if isinstance(plugin, InboundMessageHook):
|
||||
try:
|
||||
result = await plugin.on_message_received(msg)
|
||||
if result is None:
|
||||
logger.info("Message dropped by inbound hook: %s", msg.msg_id)
|
||||
await trace.finish(error="dropped_by_hook")
|
||||
return
|
||||
except Exception:
|
||||
logger.exception("Inbound hook failed, proceeding: %s", msg.msg_id)
|
||||
|
||||
try:
|
||||
await self._hooks.fire(HookEvent.MESSAGE_RECEIVED, msg)
|
||||
except Exception:
|
||||
logger.exception("MESSAGE_RECEIVED hook failed: %s", msg.msg_id)
|
||||
|
||||
chat_id = msg.group.id if msg.group and msg.group.id else msg.sender.id
|
||||
chat_type = "group" if msg.sender.kind != PeerKind.DIRECT else "direct"
|
||||
|
||||
record_data = {
|
||||
"channel_id": msg.account_id,
|
||||
"channel_type": msg.channel_type,
|
||||
"message_id": msg.msg_id,
|
||||
"chat_id": chat_id,
|
||||
"chat_type": chat_type,
|
||||
"content_type": msg.message_type.value if msg.message_type else "text",
|
||||
"sender_user_id": msg.sender.id,
|
||||
"content_preview": (msg.content or "")[:500],
|
||||
"reply_to_message_id": msg.metadata.get("reply_to_message_id") if msg.metadata else None,
|
||||
"extra_data": msg.raw_payload,
|
||||
}
|
||||
|
||||
record = None
|
||||
start_time = time.monotonic()
|
||||
|
||||
try:
|
||||
record = await self._msg_repo.create_record(record_data)
|
||||
if stop_event.is_set():
|
||||
logger.debug("Aborted before dispatch for %s", conversation_key)
|
||||
elapsed_ms = int((time.monotonic() - start_time) * 1000)
|
||||
await self._msg_repo.mark_error(record.id, "aborted_by_fence", elapsed_ms)
|
||||
await trace.finish(error="aborted_by_fence")
|
||||
return
|
||||
|
||||
bindings = await self._get_bindings()
|
||||
dispatcher = MessageDispatcher(
|
||||
bindings,
|
||||
self._allowlist,
|
||||
self._pairing,
|
||||
matcher=self._matcher,
|
||||
session_key_builder=self._session_key_builder,
|
||||
user_mapping_repo=self._user_mapping_repo,
|
||||
thread_mapping_repo=self._thread_mapping_repo,
|
||||
)
|
||||
|
||||
try:
|
||||
await self._hooks.fire(HookEvent.BEFORE_AGENT_RUN, msg)
|
||||
except Exception:
|
||||
logger.exception("BEFORE_AGENT_RUN hook failed: %s", msg.msg_id)
|
||||
|
||||
result = await dispatcher.dispatch(msg, trace=trace)
|
||||
record_dispatch_duration_ms(
|
||||
msg.channel_type,
|
||||
(time.monotonic() - start_time) * 1000,
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
"processor: phase=dispatched msg_id=%s handled_by=%s success=%s agent_config_id=%s",
|
||||
msg.msg_id, result.handled_by, result.success, result.agent_config_id,
|
||||
)
|
||||
|
||||
if result.handled_by == "command":
|
||||
if result.command_response:
|
||||
await self._send_reply(msg, result.command_response)
|
||||
try:
|
||||
await self._hooks.fire(HookEvent.AFTER_AGENT_RUN, msg, result, result.command_response)
|
||||
except Exception:
|
||||
logger.exception("AFTER_AGENT_RUN hook failed: %s", msg.msg_id)
|
||||
elapsed_ms = int((time.monotonic() - start_time) * 1000)
|
||||
await self._msg_repo.mark_success(
|
||||
record.id,
|
||||
"",
|
||||
(result.command_response or "")[:500],
|
||||
elapsed_ms,
|
||||
)
|
||||
await trace.finish()
|
||||
return
|
||||
|
||||
if not result.success:
|
||||
try:
|
||||
await self._hooks.fire(HookEvent.AFTER_AGENT_RUN, msg, result, None)
|
||||
except Exception:
|
||||
logger.exception("AFTER_AGENT_RUN hook failed: %s", msg.msg_id)
|
||||
|
||||
error_msg = result.error or "dispatch_failed"
|
||||
|
||||
fallback = _FALLBACK_MESSAGES.get(error_msg)
|
||||
if fallback is None and result.error and result.error.startswith("approval_denied:"):
|
||||
fallback = _FALLBACK_MESSAGES["approval_denied"]
|
||||
|
||||
if fallback:
|
||||
await self._send_fallback(msg, fallback)
|
||||
|
||||
elapsed_ms = int((time.monotonic() - start_time) * 1000)
|
||||
await self._msg_repo.mark_error(record.id, error_msg, elapsed_ms)
|
||||
await trace.finish(error=error_msg)
|
||||
return
|
||||
|
||||
if stop_event.is_set():
|
||||
logger.debug("Aborted after dispatch for %s", conversation_key)
|
||||
try:
|
||||
await self._hooks.fire(HookEvent.AFTER_AGENT_RUN, msg, result, None)
|
||||
except Exception:
|
||||
logger.exception("AFTER_AGENT_RUN hook failed: %s", msg.msg_id)
|
||||
elapsed_ms = int((time.monotonic() - start_time) * 1000)
|
||||
await self._msg_repo.mark_error(record.id, "aborted_by_fence", elapsed_ms)
|
||||
await trace.finish(error="aborted_by_fence")
|
||||
return
|
||||
|
||||
logger.debug(
|
||||
"processor: phase=reply_start msg_id=%s channel=%s agent_config_id=%s",
|
||||
msg.msg_id, msg.channel_type, result.agent_config_id,
|
||||
)
|
||||
await self._handle_agent_reply(msg, result, plugin, stop_event, record, start_time, conversation_key, trace, limiter)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception("Message processing failed: %s", msg.msg_id)
|
||||
await trace.finish(error=f"{type(e).__name__}: {e}")
|
||||
try:
|
||||
await self._hooks.fire(HookEvent.AFTER_AGENT_RUN, msg, None, None)
|
||||
except Exception:
|
||||
logger.exception("AFTER_AGENT_RUN hook failed: %s", msg.msg_id)
|
||||
elapsed_ms = int((time.monotonic() - start_time) * 1000)
|
||||
error_text = f"{type(e).__name__}: {e}"
|
||||
if record is not None:
|
||||
await self._msg_repo.mark_error(record.id, error_text[:500], elapsed_ms)
|
||||
record_message(msg.channel_type, "error")
|
||||
|
||||
async def _handle_agent_reply(
|
||||
self,
|
||||
msg: UnifiedMessage,
|
||||
result,
|
||||
plugin,
|
||||
stop_event: asyncio.Event | None,
|
||||
record,
|
||||
start_time: float,
|
||||
conversation_key: str,
|
||||
trace: ChannelTrace,
|
||||
limiter,
|
||||
) -> None:
|
||||
user = await self._get_user_for_channel(result.internal_user_id)
|
||||
if user is None:
|
||||
logger.error(
|
||||
"channel_turn skip: msg_id=%s channel=%s reason=no_user_for_internal_id:%s",
|
||||
msg.msg_id,
|
||||
msg.channel_type,
|
||||
result.internal_user_id,
|
||||
)
|
||||
elapsed_ms = int((time.monotonic() - start_time) * 1000)
|
||||
await self._msg_repo.mark_error(record.id, "no_user_for_internal_id", elapsed_ms)
|
||||
await trace.finish(error="no_user_for_internal_id")
|
||||
return
|
||||
|
||||
agent_config_id = result.agent_config_id
|
||||
chat_id = msg.group.id if msg.group and msg.group.id else msg.sender.id
|
||||
|
||||
try:
|
||||
await self._hooks.fire(HookEvent.AGENT_BOOTSTRAP, msg, result)
|
||||
except Exception:
|
||||
logger.exception("AGENT_BOOTSTRAP hook failed: %s", msg.msg_id)
|
||||
|
||||
stream_status: StreamStatus | None = None
|
||||
|
||||
if plugin is not None and isinstance(plugin, OutboundProtocol):
|
||||
try:
|
||||
await self._hooks.fire(HookEvent.MESSAGE_SENDING, msg, "")
|
||||
except Exception:
|
||||
logger.exception("MESSAGE_SENDING hook failed: %s", msg.msg_id)
|
||||
|
||||
reply_text, stream_status = await self._process_with_pipeline(
|
||||
msg,
|
||||
result,
|
||||
agent_config_id,
|
||||
user,
|
||||
plugin,
|
||||
stop_event,
|
||||
limiter=limiter,
|
||||
trace=trace,
|
||||
)
|
||||
else:
|
||||
async with pg_manager.get_async_session_context() as db:
|
||||
stream_result = await self._bridge.invoke_with_result(
|
||||
msg,
|
||||
result,
|
||||
agent_config_id,
|
||||
user,
|
||||
db,
|
||||
)
|
||||
reply_text = stream_result.accumulated_text or None
|
||||
stream_status = stream_result.status
|
||||
|
||||
if reply_text:
|
||||
if plugin is not None:
|
||||
reply_text = await self._apply_sending_hooks(msg, plugin, chat_id, reply_text)
|
||||
if reply_text is None:
|
||||
elapsed_ms = int((time.monotonic() - start_time) * 1000)
|
||||
await self._msg_repo.mark_error(record.id, "blocked_by_sending_hook", elapsed_ms)
|
||||
await trace.finish(error="blocked_by_sending_hook")
|
||||
return
|
||||
success = await self._send_with_retry(plugin, chat_id, reply_text, reply_to_id=msg.msg_id)
|
||||
if not success:
|
||||
logger.exception("Failed to send reply via %s after retries", msg.channel_type)
|
||||
else:
|
||||
logger.warning("No outbound plugin for channel: %s", msg.channel_type)
|
||||
elapsed_ms = int((time.monotonic() - start_time) * 1000)
|
||||
await self._msg_repo.mark_error(record.id, "no_outbound_plugin", elapsed_ms)
|
||||
await trace.finish(error="no_outbound_plugin")
|
||||
return
|
||||
|
||||
if not reply_text and stream_status in (StreamStatus.TIMEOUT, StreamStatus.ERROR, StreamStatus.CANCELLED):
|
||||
fallback_map = {
|
||||
StreamStatus.TIMEOUT: "处理超时,请稍后重试。",
|
||||
StreamStatus.ERROR: "处理异常,已通知管理员。",
|
||||
StreamStatus.CANCELLED: "处理已取消。",
|
||||
}
|
||||
fallback = fallback_map.get(stream_status)
|
||||
if fallback:
|
||||
await self._send_fallback(msg, fallback)
|
||||
elapsed_ms = int((time.monotonic() - start_time) * 1000)
|
||||
await self._msg_repo.mark_error(record.id, stream_status.value, elapsed_ms)
|
||||
await trace.finish(error=stream_status.value)
|
||||
return
|
||||
|
||||
if isinstance(plugin, SendLifecycleHook):
|
||||
try:
|
||||
await plugin.after_send_success(
|
||||
chat_id,
|
||||
{"content": reply_text, "msg_id": msg.msg_id},
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Send lifecycle hook failed: %s", msg.msg_id)
|
||||
|
||||
try:
|
||||
await self._hooks.fire(HookEvent.MESSAGE_SENT, msg, reply_text)
|
||||
except Exception:
|
||||
logger.exception("MESSAGE_SENT hook failed: %s", msg.msg_id)
|
||||
|
||||
try:
|
||||
await self._hooks.fire(HookEvent.AFTER_AGENT_RUN, msg, result, reply_text)
|
||||
except Exception:
|
||||
logger.exception("AFTER_AGENT_RUN hook failed: %s", msg.msg_id)
|
||||
|
||||
elapsed_ms = int((time.monotonic() - start_time) * 1000)
|
||||
await self._msg_repo.mark_success(
|
||||
record.id,
|
||||
"",
|
||||
(reply_text or "")[:500],
|
||||
elapsed_ms,
|
||||
)
|
||||
record_message(msg.channel_type, "success")
|
||||
await trace.finish()
|
||||
|
||||
async def _apply_sending_hooks(
|
||||
self,
|
||||
msg: UnifiedMessage,
|
||||
plugin,
|
||||
chat_id: str,
|
||||
reply_text: str,
|
||||
) -> str | None:
|
||||
if isinstance(plugin, MessageSendingHook):
|
||||
try:
|
||||
modified = await plugin.on_message_sending(msg, reply_text)
|
||||
if modified is None:
|
||||
logger.info("Message blocked by MessageSendingHook: %s", msg.msg_id)
|
||||
return None
|
||||
reply_text = modified
|
||||
except Exception:
|
||||
logger.exception("MessageSendingHook failed: %s", msg.msg_id)
|
||||
|
||||
if isinstance(plugin, SendLifecycleHook):
|
||||
try:
|
||||
await plugin.before_send_attempt(chat_id, reply_text)
|
||||
except Exception:
|
||||
logger.exception("before_send_attempt failed: %s", msg.msg_id)
|
||||
|
||||
try:
|
||||
result = await self._hooks.fire_sequential(
|
||||
HookEvent.MESSAGE_SENDING,
|
||||
initial=reply_text,
|
||||
msg=msg,
|
||||
)
|
||||
reply_text = result if result is not None else reply_text
|
||||
except Exception:
|
||||
logger.exception("MESSAGE_SENDING hook failed: %s", msg.msg_id)
|
||||
|
||||
return reply_text
|
||||
|
||||
async def _process_with_pipeline(
|
||||
self,
|
||||
msg: UnifiedMessage,
|
||||
result,
|
||||
agent_config_id: int,
|
||||
user,
|
||||
plugin,
|
||||
stop_event: asyncio.Event | None = None,
|
||||
*,
|
||||
limiter=None,
|
||||
trace: "ChannelTrace | None" = None,
|
||||
) -> tuple[str | None, StreamStatus | None]:
|
||||
target_id = msg.group.id if msg.group and msg.group.id else msg.sender.id
|
||||
|
||||
async def send_fn(content: str) -> str | None:
|
||||
if stop_event and stop_event.is_set():
|
||||
return None
|
||||
if not content:
|
||||
return None
|
||||
|
||||
send_content = content
|
||||
if isinstance(plugin, MessageSendingHook):
|
||||
try:
|
||||
modified = await plugin.on_message_sending(msg, send_content)
|
||||
if modified is None:
|
||||
logger.info("Pipeline send blocked by MessageSendingHook: %s", msg.msg_id)
|
||||
return None
|
||||
send_content = modified
|
||||
except Exception:
|
||||
logger.exception("Pipeline MessageSendingHook failed: %s", msg.msg_id)
|
||||
|
||||
if isinstance(plugin, SendLifecycleHook):
|
||||
try:
|
||||
await plugin.before_send_attempt(target_id, send_content)
|
||||
except Exception:
|
||||
logger.exception("Pipeline before_send_attempt failed: %s", msg.msg_id)
|
||||
|
||||
success = await self._send_with_retry(plugin, target_id, send_content, reply_to_id=msg.msg_id)
|
||||
if not success:
|
||||
logger.exception("Pipeline send failed via %s after retries", msg.channel_type)
|
||||
return None
|
||||
|
||||
pipeline = create_block_reply_pipeline(
|
||||
send_fn=send_fn,
|
||||
msg=msg,
|
||||
response_prefix=self._response_prefix,
|
||||
human_delay=self._human_delay,
|
||||
)
|
||||
|
||||
accumulated: list[str] = []
|
||||
stream_error: StreamStatus | None = None
|
||||
|
||||
exec_span_id = None
|
||||
if trace:
|
||||
exec_span_id = await trace.add_span(
|
||||
"agent_execution",
|
||||
metadata={
|
||||
"agent_config_id": agent_config_id,
|
||||
"thread_id": result.thread_id,
|
||||
},
|
||||
)
|
||||
|
||||
async with pg_manager.get_async_session_context() as db:
|
||||
try:
|
||||
stream_iter = self._bridge.invoke_stream(
|
||||
msg,
|
||||
result,
|
||||
agent_config_id,
|
||||
user,
|
||||
db,
|
||||
)
|
||||
|
||||
if limiter is not None:
|
||||
stream_iter = limiter.run_with_limit(stream_iter)
|
||||
|
||||
async for chunk in stream_iter:
|
||||
if stop_event and stop_event.is_set():
|
||||
await pipeline.abort()
|
||||
logger.debug("Pipeline aborted by foreground fence for %s", self._fence.key_for(msg))
|
||||
break
|
||||
await pipeline.enqueue(chunk)
|
||||
if chunk.content:
|
||||
accumulated.append(chunk.content)
|
||||
except StreamTimeoutError:
|
||||
stream_error = StreamStatus.TIMEOUT
|
||||
logger.warning("Agent stream timeout for channel=%s peer=%s", msg.channel_type, msg.sender.id)
|
||||
except StreamCancelledError:
|
||||
stream_error = StreamStatus.CANCELLED
|
||||
logger.debug("Agent stream cancelled for channel=%s peer=%s", msg.channel_type, msg.sender.id)
|
||||
|
||||
await pipeline.wait_idle()
|
||||
|
||||
if exec_span_id:
|
||||
status_msg = stream_error.value if stream_error else "ok"
|
||||
await trace.end_span(
|
||||
exec_span_id,
|
||||
level="ERROR" if stream_error else "DEFAULT",
|
||||
status_message=status_msg,
|
||||
metadata={"stream_status": status_msg, "chunks": len(accumulated)},
|
||||
)
|
||||
|
||||
if stop_event and stop_event.is_set():
|
||||
return None, None
|
||||
|
||||
return ("".join(accumulated) if accumulated else None, stream_error)
|
||||
|
||||
async def _send_fallback(self, msg: UnifiedMessage, fallback_text: str) -> None:
|
||||
plugin = ChannelPluginRegistry.get(msg.channel_type)
|
||||
if plugin is None or not isinstance(plugin, OutboundProtocol):
|
||||
return
|
||||
try:
|
||||
target_id = msg.group.id if msg.group and msg.group.id else msg.sender.id
|
||||
await plugin.send_text(target_id, fallback_text, reply_to_id=msg.msg_id)
|
||||
except Exception:
|
||||
logger.exception("Failed to send fallback reply via %s", msg.channel_type)
|
||||
|
||||
async def _send_with_retry(
|
||||
self,
|
||||
plugin,
|
||||
target_id: str,
|
||||
content: str,
|
||||
*,
|
||||
reply_to_id: str | None = None,
|
||||
max_retries: int = 3,
|
||||
circuit_breaker: "CircuitBreaker | None" = None,
|
||||
) -> bool:
|
||||
last_error = None
|
||||
total_attempts = max_retries + 1
|
||||
|
||||
for attempt in range(total_attempts):
|
||||
if circuit_breaker is not None and not circuit_breaker.allow_request():
|
||||
logger.warning(
|
||||
"Circuit breaker OPEN for %s, skipping send",
|
||||
plugin.id if hasattr(plugin, "id") else "unknown",
|
||||
)
|
||||
return False
|
||||
|
||||
try:
|
||||
await plugin.send_text(target_id, content, reply_to_id=reply_to_id)
|
||||
if circuit_breaker is not None:
|
||||
circuit_breaker.record_success()
|
||||
return True
|
||||
except Exception as e:
|
||||
last_error = e
|
||||
if circuit_breaker is not None:
|
||||
circuit_breaker.record_failure()
|
||||
logger.warning(
|
||||
"Send attempt %d/%d failed (target=%s): %s",
|
||||
attempt + 1, total_attempts, target_id, e,
|
||||
)
|
||||
if attempt < max_retries:
|
||||
delay = min(0.3 * (2 ** attempt), 30.0)
|
||||
await asyncio.sleep(delay)
|
||||
|
||||
logger.error("All %d send attempts failed (target=%s): %s", total_attempts, target_id, last_error)
|
||||
if isinstance(plugin, SendLifecycleHook):
|
||||
try:
|
||||
await plugin.after_send_failure(target_id, str(last_error))
|
||||
except Exception:
|
||||
pass
|
||||
return False
|
||||
|
||||
async def _send_reply(self, msg: UnifiedMessage, content: str) -> str | None:
|
||||
plugin = ChannelPluginRegistry.get(msg.channel_type)
|
||||
if plugin is None or not isinstance(plugin, OutboundProtocol):
|
||||
logger.warning("No outbound plugin for channel: %s", msg.channel_type)
|
||||
return None
|
||||
|
||||
target_id = msg.group.id if msg.group and msg.group.id else msg.sender.id
|
||||
reply_to_id = msg.msg_id
|
||||
|
||||
send_content = await self._apply_sending_hooks(msg, plugin, target_id, content)
|
||||
if send_content is None:
|
||||
return None
|
||||
|
||||
success = await self._send_with_retry(plugin, target_id, send_content, reply_to_id=reply_to_id)
|
||||
if not success:
|
||||
logger.exception("Failed to send reply via %s after retries", msg.channel_type)
|
||||
return None
|
||||
217
backend/package/yuxi/channel/message/rate_limiter.py
Normal file
217
backend/package/yuxi/channel/message/rate_limiter.py
Normal file
@ -0,0 +1,217 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import AsyncIterator, Callable
|
||||
from dataclasses import dataclass
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_global_semaphore = asyncio.Semaphore(50)
|
||||
|
||||
|
||||
def get_global_semaphore() -> asyncio.Semaphore:
|
||||
return _global_semaphore
|
||||
|
||||
DEFAULT_RATE = 10.0
|
||||
DEFAULT_CAPACITY = 20
|
||||
DEFAULT_PER_ACCOUNT_CONCURRENCY = 5
|
||||
DEFAULT_QUEUE_DEPTH = 20
|
||||
DEFAULT_MAX_LIMITERS = 10_000
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class AcquireResult:
|
||||
allowed: bool
|
||||
retry_after_sec: float = 0.0
|
||||
|
||||
def __bool__(self) -> bool:
|
||||
return self.allowed
|
||||
|
||||
|
||||
class TokenBucket:
|
||||
"""令牌桶算法 — 控制单账户消息速率。
|
||||
|
||||
参数:
|
||||
- rate: 每秒补充的令牌数(默认 10)
|
||||
- capacity: 桶容量上限(默认 20,允许短时突发)
|
||||
- time_func: 可选时钟注入,便于测试(默认 time.monotonic)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
rate: float = DEFAULT_RATE,
|
||||
capacity: int = DEFAULT_CAPACITY,
|
||||
time_func: Callable[[], float] = time.monotonic,
|
||||
):
|
||||
self._rate = rate
|
||||
self._capacity = capacity
|
||||
self._tokens = float(capacity)
|
||||
self._time = time_func
|
||||
self._last_refill = self._time()
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
async def acquire(self) -> AcquireResult:
|
||||
async with self._lock:
|
||||
now = self._time()
|
||||
elapsed = now - self._last_refill
|
||||
self._tokens = min(self._capacity, self._tokens + elapsed * self._rate)
|
||||
self._last_refill = now
|
||||
|
||||
if self._tokens >= 1.0:
|
||||
self._tokens -= 1.0
|
||||
return AcquireResult(allowed=True)
|
||||
|
||||
if self._rate <= 0.0:
|
||||
return AcquireResult(allowed=False, retry_after_sec=float("inf"))
|
||||
|
||||
retry_after = (1.0 - self._tokens) / self._rate
|
||||
return AcquireResult(allowed=False, retry_after_sec=retry_after)
|
||||
|
||||
def reset(self) -> None:
|
||||
self._tokens = float(self._capacity)
|
||||
self._last_refill = self._time()
|
||||
|
||||
|
||||
class ChannelRateLimiter:
|
||||
"""单渠道账户的限流控制器,组合令牌桶 + 并发信号量。
|
||||
|
||||
线程安全:所有方法在 asyncio 上下文中调用。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
rate: float = DEFAULT_RATE,
|
||||
per_account_concurrency: int = DEFAULT_PER_ACCOUNT_CONCURRENCY,
|
||||
queue_depth: int = DEFAULT_QUEUE_DEPTH,
|
||||
):
|
||||
self._rate = rate
|
||||
self._per_account_concurrency = per_account_concurrency
|
||||
self._bucket = TokenBucket(rate=rate, capacity=DEFAULT_CAPACITY)
|
||||
self._semaphore = asyncio.Semaphore(per_account_concurrency)
|
||||
self._queue_depth = queue_depth
|
||||
self._waiting: int = 0
|
||||
|
||||
async def try_acquire(self) -> AcquireResult:
|
||||
"""尝试获取消息处理许可。
|
||||
|
||||
返回 AcquireResult(allowed=..., retry_after_sec=...)。
|
||||
"""
|
||||
result = await self._bucket.acquire()
|
||||
if result.allowed:
|
||||
return result
|
||||
|
||||
if self._waiting >= self._queue_depth:
|
||||
logger.warning(
|
||||
"Rate limit queue full (depth=%d), rejecting message",
|
||||
self._queue_depth,
|
||||
)
|
||||
return AcquireResult(allowed=False, retry_after_sec=result.retry_after_sec)
|
||||
|
||||
self._waiting += 1
|
||||
try:
|
||||
await asyncio.sleep(result.retry_after_sec)
|
||||
retry = await self._bucket.acquire()
|
||||
if not retry.allowed:
|
||||
return AcquireResult(allowed=False, retry_after_sec=retry.retry_after_sec)
|
||||
return retry
|
||||
finally:
|
||||
self._waiting -= 1
|
||||
|
||||
async def run_with_limit(self, coro) -> AsyncIterator:
|
||||
"""在并发限制内执行 Agent 调用。
|
||||
|
||||
先获取账户级信号量,再获取全局信号量。
|
||||
"""
|
||||
async with self._semaphore:
|
||||
async with _global_semaphore:
|
||||
async for chunk in coro:
|
||||
yield chunk
|
||||
|
||||
def reset(self) -> None:
|
||||
self._bucket.reset()
|
||||
|
||||
@property
|
||||
def waiting_count(self) -> int:
|
||||
return self._waiting
|
||||
|
||||
@property
|
||||
def config(self) -> dict:
|
||||
return {
|
||||
"rate": self._rate,
|
||||
"per_account_concurrency": self._per_account_concurrency,
|
||||
"queue_depth": self._queue_depth,
|
||||
}
|
||||
|
||||
|
||||
class RateLimitManager:
|
||||
"""限流管理器 — 按账户维度管理限流器实例。
|
||||
|
||||
每个 (channel_type, account_id) 组合共享同一个 ChannelRateLimiter。
|
||||
内置自动清理机制防止无限增长。
|
||||
"""
|
||||
|
||||
def __init__(self, max_limiters: int = DEFAULT_MAX_LIMITERS):
|
||||
self._limiters: dict[str, tuple[float, ChannelRateLimiter]] = {}
|
||||
self._max_limiters = max_limiters
|
||||
self._last_cleanup = time.monotonic()
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
def _key(self, channel_type: str, account_id: str) -> str:
|
||||
return f"{channel_type}:{account_id}"
|
||||
|
||||
async def get_limiter(
|
||||
self,
|
||||
channel_type: str,
|
||||
account_id: str,
|
||||
rate: float = DEFAULT_RATE,
|
||||
per_account_concurrency: int = DEFAULT_PER_ACCOUNT_CONCURRENCY,
|
||||
) -> ChannelRateLimiter:
|
||||
async with self._lock:
|
||||
self._maybe_cleanup()
|
||||
key = self._key(channel_type, account_id)
|
||||
if key not in self._limiters:
|
||||
self._limiters[key] = (
|
||||
time.monotonic(),
|
||||
ChannelRateLimiter(
|
||||
rate=rate,
|
||||
per_account_concurrency=per_account_concurrency,
|
||||
),
|
||||
)
|
||||
else:
|
||||
self._limiters[key] = (time.monotonic(), self._limiters[key][1])
|
||||
return self._limiters[key][1]
|
||||
|
||||
async def remove_limiter(self, channel_type: str, account_id: str) -> None:
|
||||
async with self._lock:
|
||||
key = self._key(channel_type, account_id)
|
||||
self._limiters.pop(key, None)
|
||||
|
||||
async def clear(self) -> None:
|
||||
async with self._lock:
|
||||
self._limiters.clear()
|
||||
|
||||
async def get_stats(self) -> dict[str, dict]:
|
||||
async with self._lock:
|
||||
return {
|
||||
key: {
|
||||
"waiting": limiter.waiting_count,
|
||||
**limiter.config,
|
||||
}
|
||||
for key, (_, limiter) in self._limiters.items()
|
||||
}
|
||||
|
||||
def _maybe_cleanup(self) -> None:
|
||||
if len(self._limiters) <= self._max_limiters:
|
||||
return
|
||||
now = time.monotonic()
|
||||
if now - self._last_cleanup < 300:
|
||||
return
|
||||
self._last_cleanup = now
|
||||
expired = [key for key, (last_touch, _) in self._limiters.items() if now - last_touch > 3600]
|
||||
for key in expired:
|
||||
self._limiters.pop(key, None)
|
||||
if expired:
|
||||
logger.info("Cleaned up %d stale rate limiters", len(expired))
|
||||
|
||||
|
||||
rate_limit_manager = RateLimitManager()
|
||||
269
backend/package/yuxi/channel/message/reply_dispatcher.py
Normal file
269
backend/package/yuxi/channel/message/reply_dispatcher.py
Normal file
@ -0,0 +1,269 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import random
|
||||
import time
|
||||
from collections.abc import Callable, Coroutine
|
||||
from typing import Any
|
||||
|
||||
from yuxi.channel.message.models import ReplyPayload, ReplyStage, StreamingChunk, UnifiedMessage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_DEFAULT_MIN_CHARS = 80
|
||||
_DEFAULT_DEBOUNCE_MS = 150.0
|
||||
_DEFAULT_TIMEOUT_MS = 10_000.0
|
||||
|
||||
OnIdleCallback = Callable[[], None]
|
||||
OnErrorCallback = Callable[[Exception], None]
|
||||
BeforeDeliverCallback = Callable[[str], Coroutine[Any, Any, str | None]]
|
||||
|
||||
|
||||
class ReplyDispatcher:
|
||||
def __init__(
|
||||
self,
|
||||
send_fn: Callable[..., Coroutine[Any, Any, str | None]],
|
||||
msg: UnifiedMessage,
|
||||
*,
|
||||
response_prefix: str = "",
|
||||
human_delay: tuple[float, float] | None = None,
|
||||
timeout_ms: float = _DEFAULT_TIMEOUT_MS,
|
||||
on_idle: OnIdleCallback | None = None,
|
||||
on_error: OnErrorCallback | None = None,
|
||||
before_deliver: BeforeDeliverCallback | None = None,
|
||||
):
|
||||
self._send_fn = send_fn
|
||||
self._msg = msg
|
||||
self._response_prefix = response_prefix
|
||||
self._human_delay = human_delay
|
||||
self._timeout_ms = timeout_ms
|
||||
self._on_idle = on_idle
|
||||
self._on_error = on_error
|
||||
self._before_deliver = before_deliver
|
||||
|
||||
self._send_chain: asyncio.Lock = asyncio.Lock()
|
||||
self._flush_lock: asyncio.Lock = asyncio.Lock()
|
||||
self._pending = 0
|
||||
self._complete = asyncio.Event()
|
||||
self._done = False
|
||||
|
||||
self._blocks: list[str] = []
|
||||
self._block_deadline: float = 0.0
|
||||
self._block_timer: asyncio.Task | None = None
|
||||
self._sent_payload_keys: set[str] = set()
|
||||
self._sent_content_keys: set[str] = set()
|
||||
|
||||
self._tool_count = 0
|
||||
self._block_count = 0
|
||||
self._final_count = 0
|
||||
self._failed_count = 0
|
||||
self._cancelled_count = 0
|
||||
self._started_at: float | None = None
|
||||
|
||||
self._abort_event = asyncio.Event()
|
||||
|
||||
async def enqueue(self, chunk: StreamingChunk) -> bool:
|
||||
if self._abort_event.is_set():
|
||||
self._cancelled_count += 1
|
||||
return False
|
||||
|
||||
if chunk.stage == ReplyStage.TOOL:
|
||||
return await self._dispatch_tool(chunk)
|
||||
elif chunk.stage == ReplyStage.BLOCK:
|
||||
return await self._dispatch_block(chunk)
|
||||
elif chunk.stage == ReplyStage.FINAL:
|
||||
return await self._dispatch_final(chunk)
|
||||
return False
|
||||
|
||||
async def _dispatch_tool(self, chunk: StreamingChunk) -> bool:
|
||||
self._pending += 1
|
||||
self._tool_count += 1
|
||||
try:
|
||||
prefix = self._response_prefix if self._response_prefix else ""
|
||||
content = f"{prefix}{chunk.content}" if chunk.content else ""
|
||||
if not content:
|
||||
return False
|
||||
|
||||
await self._send_ordered(content)
|
||||
self._block_count += 1
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.exception("Tool dispatch failed")
|
||||
self._failed_count += 1
|
||||
if self._on_error:
|
||||
self._on_error(e)
|
||||
raise
|
||||
finally:
|
||||
self._pending -= 1
|
||||
self._maybe_fire_idle()
|
||||
|
||||
async def _dispatch_block(self, chunk: StreamingChunk) -> bool:
|
||||
if not chunk.content:
|
||||
return False
|
||||
|
||||
async with self._flush_lock:
|
||||
self._blocks.append(chunk.content)
|
||||
now = time.monotonic()
|
||||
|
||||
if self._block_deadline == 0.0:
|
||||
self._block_deadline = now + _DEFAULT_DEBOUNCE_MS / 1000.0
|
||||
self._block_timer = asyncio.create_task(self._flush_after_debounce())
|
||||
|
||||
if sum(len(b) for b in self._blocks) >= _DEFAULT_MIN_CHARS:
|
||||
await self._flush_blocks()
|
||||
return True
|
||||
|
||||
async def _flush_after_debounce(self) -> None:
|
||||
try:
|
||||
await asyncio.sleep(_DEFAULT_DEBOUNCE_MS / 1000.0)
|
||||
except asyncio.CancelledError:
|
||||
return
|
||||
async with self._flush_lock:
|
||||
await self._flush_blocks()
|
||||
|
||||
async def _flush_blocks(self) -> None:
|
||||
if not self._blocks:
|
||||
return
|
||||
combined = "".join(self._blocks)
|
||||
self._blocks.clear()
|
||||
self._block_deadline = 0.0
|
||||
if self._block_timer and not self._block_timer.done():
|
||||
self._block_timer.cancel()
|
||||
self._block_timer = None
|
||||
|
||||
payload = ReplyPayload(
|
||||
target_id=self._target_id,
|
||||
content=combined,
|
||||
reply_to_id=self._msg.msg_id if self._block_count == 0 else None,
|
||||
)
|
||||
if payload.content_key in self._sent_content_keys:
|
||||
return
|
||||
self._sent_content_keys.add(payload.content_key)
|
||||
|
||||
self._pending += 1
|
||||
self._block_count += 1
|
||||
try:
|
||||
await self._send_ordered(combined)
|
||||
except Exception as e:
|
||||
logger.exception("Block dispatch failed")
|
||||
self._failed_count += 1
|
||||
if self._on_error:
|
||||
self._on_error(e)
|
||||
raise
|
||||
finally:
|
||||
self._pending -= 1
|
||||
self._maybe_fire_idle()
|
||||
|
||||
async def _dispatch_final(self, chunk: StreamingChunk) -> bool:
|
||||
async with self._flush_lock:
|
||||
await self._flush_blocks()
|
||||
|
||||
if chunk.content and chunk.content != "finished":
|
||||
self._pending += 1
|
||||
self._final_count += 1
|
||||
try:
|
||||
await self._send_ordered(chunk.content)
|
||||
except Exception as e:
|
||||
logger.exception("Final dispatch failed")
|
||||
self._failed_count += 1
|
||||
if self._on_error:
|
||||
self._on_error(e)
|
||||
raise
|
||||
finally:
|
||||
self._pending -= 1
|
||||
self._maybe_fire_idle()
|
||||
self._done = True
|
||||
self._complete.set()
|
||||
return True
|
||||
|
||||
self._done = True
|
||||
self._complete.set()
|
||||
self._maybe_fire_idle()
|
||||
return True
|
||||
|
||||
async def _send_ordered(self, content: str) -> None:
|
||||
async with self._send_chain:
|
||||
if self._abort_event.is_set():
|
||||
return
|
||||
|
||||
if self._human_delay and self._block_count > 0:
|
||||
delay = self._human_delay[0] + (self._human_delay[1] - self._human_delay[0]) * random.random()
|
||||
await asyncio.sleep(delay / 1000.0)
|
||||
|
||||
if self._response_prefix and self._block_count == 0:
|
||||
prefix = self._response_prefix.format(
|
||||
agent_name=self._msg.channel_type,
|
||||
model="",
|
||||
)
|
||||
content = f"{prefix}{content}"
|
||||
|
||||
if self._before_deliver:
|
||||
try:
|
||||
content = await self._before_deliver(content)
|
||||
except Exception as e:
|
||||
logger.exception("before_deliver hook failed")
|
||||
if self._on_error:
|
||||
self._on_error(e)
|
||||
return
|
||||
if content is None:
|
||||
self._cancelled_count += 1
|
||||
return
|
||||
|
||||
await self._send_fn(content)
|
||||
|
||||
@property
|
||||
def _target_id(self) -> str:
|
||||
return self._msg.group.id if self._msg.group and self._msg.group.id else self._msg.sender.id
|
||||
|
||||
async def wait_idle(self, timeout_ms: float | None = None) -> None:
|
||||
timeout = timeout_ms or self._timeout_ms
|
||||
try:
|
||||
await asyncio.wait_for(self._complete.wait(), timeout=timeout / 1000.0)
|
||||
except TimeoutError:
|
||||
self._abort_event.set()
|
||||
self._done = True
|
||||
self._complete.set()
|
||||
logger.warning("ReplyDispatcher idle timeout after %.0fms, aborting", timeout)
|
||||
|
||||
def mark_complete(self) -> None:
|
||||
self._done = True
|
||||
self._complete.set()
|
||||
self._maybe_fire_idle()
|
||||
|
||||
async def abort(self) -> None:
|
||||
async with self._flush_lock:
|
||||
if self._blocks:
|
||||
await self._flush_blocks()
|
||||
self._abort_event.set()
|
||||
self._complete.set()
|
||||
|
||||
def _maybe_fire_idle(self) -> None:
|
||||
if self._done and self._pending == 0 and self._on_idle:
|
||||
self._on_idle()
|
||||
|
||||
@property
|
||||
def tool_count(self) -> int:
|
||||
return self._tool_count
|
||||
|
||||
@property
|
||||
def block_count(self) -> int:
|
||||
return self._block_count
|
||||
|
||||
@property
|
||||
def final_count(self) -> int:
|
||||
return self._final_count
|
||||
|
||||
@property
|
||||
def failed_count(self) -> int:
|
||||
return self._failed_count
|
||||
|
||||
@property
|
||||
def cancelled_count(self) -> int:
|
||||
return self._cancelled_count
|
||||
|
||||
@property
|
||||
def queued_counts(self) -> dict[str, int]:
|
||||
return {"tool": self._tool_count, "block": self._block_count, "final": self._final_count}
|
||||
|
||||
@property
|
||||
def is_idle(self) -> bool:
|
||||
return self._done and self._pending == 0
|
||||
Loading…
Reference in New Issue
Block a user