185 lines
6.4 KiB
Python
185 lines
6.4 KiB
Python
"""默认入站中间件实现（聚合文件）。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from typing import TYPE_CHECKING
|
||
|
||
from yuxi.channel.constants import InboundRejectionReason
|
||
from yuxi.channel.middlewares.protocols import InboundResult
|
||
from yuxi.storage.postgres.manager import pg_manager
|
||
from yuxi.utils.logging_config import logger
|
||
|
||
if TYPE_CHECKING:
|
||
from collections.abc import Awaitable, Callable
|
||
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from yuxi.channel.message.dedupe import MessageDeduper
|
||
from yuxi.channel.middlewares.protocols import InboundContext
|
||
from yuxi.channel.plugins.protocol import BindingRoute, InboundMessage
|
||
from yuxi.channel.routing.router import BindingRouter
|
||
from yuxi.channel.security.policy import SecurityPolicy
|
||
from yuxi.channel.session.manager import SessionManager
|
||
from yuxi.storage.postgres.model_channel import ChannelSession
|
||
|
||
|
||
class DedupeMiddleware:
|
||
"""基于 Redis 的入站消息幂等去重。"""
|
||
|
||
name = "dedupe"
|
||
default_order = 100
|
||
|
||
def __init__(self, deduper: MessageDeduper) -> None:
|
||
self._deduper = deduper
|
||
|
||
async def process(
|
||
self,
|
||
ctx: InboundContext,
|
||
next_mw: Callable[[], Awaitable[InboundResult]],
|
||
) -> InboundResult:
|
||
if await self._deduper.is_processed(ctx.inbound.channel_message_id):
|
||
return InboundResult(accepted=False, reason=InboundRejectionReason.DUPLICATE)
|
||
return await next_mw()
|
||
|
||
|
||
class SecurityMiddleware:
|
||
"""调用安全策略链，按结果决定是否继续后续中间件。"""
|
||
|
||
name = "security"
|
||
default_order = 200
|
||
|
||
def __init__(self, security_policy: SecurityPolicy, deduper: MessageDeduper) -> None:
|
||
self._security = security_policy
|
||
self._deduper = deduper
|
||
|
||
async def process(
|
||
self,
|
||
ctx: InboundContext,
|
||
next_mw: Callable[[], Awaitable[InboundResult]],
|
||
) -> InboundResult:
|
||
from yuxi.channel.metrics import channel_rate_limited_total
|
||
|
||
security_result = await self._security.check(ctx.config, ctx.plugin, ctx.inbound)
|
||
if security_result.allowed:
|
||
return await next_mw()
|
||
|
||
if security_result.reason == InboundRejectionReason.RATE_LIMITED:
|
||
channel_rate_limited_total.inc({"channel_type": ctx.channel_type, "account_id": ctx.account_id})
|
||
|
||
logger.bind(
|
||
event="channel_message_rejected",
|
||
channel_type=ctx.channel_type,
|
||
account_id=ctx.account_id,
|
||
session_key=ctx.inbound.session_key,
|
||
message_id=ctx.inbound.channel_message_id,
|
||
sender_id=ctx.inbound.sender_id,
|
||
reason=security_result.reason,
|
||
).warning("Channel inbound message rejected by security policy")
|
||
|
||
if security_result.reason == InboundRejectionReason.DM_PAIRING_REQUIRED and security_result.pairing_code:
|
||
return InboundResult(
|
||
accepted=False,
|
||
reason=security_result.reason,
|
||
pairing_code=security_result.pairing_code,
|
||
qr_content=security_result.qr_content,
|
||
qr_reply=security_result.qr_reply,
|
||
)
|
||
|
||
# 限流后外部平台会重试，清除去重标记避免重试被误判为 DUPLICATE。
|
||
if security_result.reason == InboundRejectionReason.RATE_LIMITED:
|
||
await self._deduper.safe_clear_processed(ctx.inbound.channel_message_id)
|
||
|
||
return InboundResult(
|
||
accepted=False,
|
||
reason=security_result.reason or InboundRejectionReason.SECURITY_REJECTED,
|
||
)
|
||
|
||
|
||
class TransactionMiddleware:
|
||
"""开启数据库事务并将 AsyncSession 注入上下文。"""
|
||
|
||
name = "transaction"
|
||
default_order = 300
|
||
|
||
def __init__(self, deduper: MessageDeduper) -> None:
|
||
self._deduper = deduper
|
||
|
||
async def process(
|
||
self,
|
||
ctx: InboundContext,
|
||
next_mw: Callable[[], Awaitable[InboundResult]],
|
||
) -> InboundResult:
|
||
try:
|
||
async with pg_manager.get_async_session_context() as db:
|
||
ctx.db = db
|
||
return await next_mw()
|
||
except Exception:
|
||
# 事务回滚在 async context manager 异常退出时自动完成。
|
||
# 清除去重标记，使外部平台收到非 2xx 后重试不会被误判为 DUPLICATE。
|
||
await self._deduper.safe_clear_processed(ctx.inbound.channel_message_id)
|
||
raise
|
||
|
||
|
||
class SessionMiddleware:
|
||
"""解析或创建 ChannelSession，并将结果注入上下文。"""
|
||
|
||
name = "session"
|
||
default_order = 400
|
||
|
||
def __init__(self, session_manager: SessionManager) -> None:
|
||
self._session_manager = session_manager
|
||
|
||
async def process(
|
||
self,
|
||
ctx: InboundContext,
|
||
next_mw: Callable[[], Awaitable[InboundResult]],
|
||
) -> InboundResult:
|
||
session, route = await self._session_manager.resolve(ctx.config, ctx.plugin, ctx.inbound, ctx.db)
|
||
ctx.session = session
|
||
ctx.route = route
|
||
return await next_mw()
|
||
|
||
|
||
class RouteMiddleware:
|
||
"""在 session 未携带路由时解析 BindingRoute。"""
|
||
|
||
name = "route"
|
||
default_order = 500
|
||
|
||
def __init__(self, binding_router: BindingRouter) -> None:
|
||
self._binding_router = binding_router
|
||
|
||
async def process(
|
||
self,
|
||
ctx: InboundContext,
|
||
next_mw: Callable[[], Awaitable[InboundResult]],
|
||
) -> InboundResult:
|
||
if ctx.route is None:
|
||
ctx.route = await self._binding_router.resolve(ctx.session, ctx.config, ctx.plugin, ctx.inbound, ctx.db)
|
||
return await next_mw()
|
||
|
||
|
||
class CreateRunMiddleware:
|
||
"""创建 AgentRun 并记录消息；作为默认入站链的终点，不再调用 next_mw。"""
|
||
|
||
name = "create_run"
|
||
default_order = 600
|
||
|
||
def __init__(
|
||
self,
|
||
create_run: Callable[[ChannelSession, BindingRoute, InboundMessage, AsyncSession], Awaitable[dict]],
|
||
record_message: Callable[[ChannelSession, InboundMessage, AsyncSession], Awaitable[None]],
|
||
) -> None:
|
||
self._create_run = create_run
|
||
self._record_message = record_message
|
||
|
||
async def process(
|
||
self,
|
||
ctx: InboundContext,
|
||
next_mw: Callable[[], Awaitable[InboundResult]],
|
||
) -> InboundResult:
|
||
run_result = await self._create_run(ctx.session, ctx.route, ctx.inbound, ctx.db)
|
||
ctx.run_result = run_result
|
||
await self._record_message(ctx.session, ctx.inbound, ctx.db)
|
||
return InboundResult(accepted=True, run_id=run_result.get("run_id"))
|