ForcePilot/backend/package/yuxi/channel/message/handler.py

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