beaver_project/app-instance/backend/beaver/interfaces/gateway/main.py

"""Gateway entrypoint for Beaver.

当前阶段先不扩 bus / channels adapter，只做最小消息桥接：
1. 启动时托管 `AgentService.start()`
2. 常驻消费 `MessageBus.inbound`
3. 调 `service.submit_direct(...)`
4. 将结果写回 `MessageBus.outbound`
5. 退出时走 `AgentService.shutdown()`
"""

from __future__ import annotations

import asyncio
from pathlib import Path

from beaver.foundation.events import InboundMessage, MessageBus, OutboundMessage
from beaver.services.agent_service import AgentService


async def _publish_bridge_error(
    bus: MessageBus,
    inbound: InboundMessage,
    *,
    detail: str,
    finish_reason: str = "error",
) -> None:
    """把 bridge 处理失败转换成结构化 outbound 错误消息。"""

    await bus.publish_outbound(
        OutboundMessage(
            message_id=inbound.message_id,
            channel=inbound.channel,
            session_id=inbound.session_id,
            content=detail,
            finish_reason=finish_reason,
            metadata={"error": detail, "inbound_metadata": dict(inbound.metadata)},
        )
    )


async def _flush_pending_inbound(bus: MessageBus, *, reason: str) -> None:
    """把尚未处理的 inbound 明确冲刷成 outbound 错误，而不是静默丢弃。"""

    while True:
        try:
            pending = bus.inbound.get_nowait()
        except asyncio.QueueEmpty:
            break
        await _publish_bridge_error(bus, pending, detail=reason, finish_reason="stopped")


async def _await_bridge_shutdown(task: asyncio.Task[None], *, timeout_seconds: float = 1.0) -> None:
    """等待 bridge 退出；超时则取消，避免 shutdown 被桥接层反向卡死。"""

    try:
        await asyncio.wait_for(task, timeout=timeout_seconds)
    except asyncio.CancelledError:
        pass
    except asyncio.TimeoutError:
        task.cancel()
        try:
            await task
        except asyncio.CancelledError:
            pass


async def _bridge_inbound_to_runtime(
    service: AgentService,
    bus: MessageBus,
    stop_event: asyncio.Event,
) -> None:
    """Consume inbound messages, run the agent, and publish outbound results."""

    while True:
        if stop_event.is_set():
            await _flush_pending_inbound(
                bus,
                reason="Gateway stopped before processing the inbound message",
            )
            break

        try:
            inbound = await asyncio.wait_for(bus.consume_inbound(), timeout=0.25)
        except asyncio.TimeoutError:
            continue

        try:
            result = await service.submit_direct(
                inbound.content,
                session_id=inbound.session_id,
                source=f"gateway:{inbound.channel}",
                user_id=inbound.user_id,
                title=inbound.title,
                execution_context=inbound.execution_context,
                model=inbound.model,
                provider_name=inbound.provider_name,
                embedding_model=inbound.embedding_model,
            )
        except asyncio.CancelledError:
            await _publish_bridge_error(
                bus,
                inbound,
                detail="Gateway stopped before completing the inbound message",
                finish_reason="cancelled",
            )
            raise
        except Exception as exc:  # pragma: no cover - defensive bridge path
            await _publish_bridge_error(
                bus,
                inbound,
                detail=str(exc),
            )
        else:
            await bus.publish_outbound(
                OutboundMessage(
                    message_id=inbound.message_id,
                    channel=inbound.channel,
                    session_id=result.session_id,
                    run_id=result.run_id,
                    content=result.output_text,
                    finish_reason=result.finish_reason,
                    provider_name=result.provider_name,
                    model=result.model,
                    usage=dict(result.usage),
                    metadata={"inbound_metadata": dict(inbound.metadata)},
                )
            )


async def run_gateway(
    *,
    workspace: str | Path | None = None,
    service: AgentService | None = None,
    bus: MessageBus | None = None,
    manage_service_lifecycle: bool | None = None,
    stop_event: asyncio.Event | None = None,
    shutdown_timeout_seconds: float | None = 5.0,
    shutdown_force: bool = True,
) -> None:
    """运行最小 gateway 宿主层与消息桥接。

    默认 ownership 语义：
    - 未传 `service`：gateway 自己创建并接管其 lifecycle
    - 传入外部 `service`：默认只使用，不自动 start/shutdown
    """

    attached_service = service or AgentService(workspace=workspace)
    attached_bus = bus or MessageBus()
    owns_service = manage_service_lifecycle if manage_service_lifecycle is not None else service is None
    owned_stop_event = stop_event or asyncio.Event()
    started = False
    if owns_service:
        try:
            await attached_service.start()
            started = True
        except Exception:
            attached_service.close()
            raise

    if not attached_service.is_running:
        raise RuntimeError(
            "Gateway requires AgentService running mode; start the injected service first "
            "or allow the gateway to manage its lifecycle."
        )

    bridge_task = asyncio.create_task(_bridge_inbound_to_runtime(attached_service, attached_bus, owned_stop_event))
    try:
        await owned_stop_event.wait()
    finally:
        owned_stop_event.set()
        if owns_service and started:
            try:
                await attached_service.shutdown(
                    timeout_seconds=shutdown_timeout_seconds,
                    force=shutdown_force,
                )
            finally:
                await _await_bridge_shutdown(bridge_task)
        else:
            await _await_bridge_shutdown(bridge_task)


def main() -> None:
    """同步 gateway 入口。"""

    try:
        asyncio.run(run_gateway())
    except KeyboardInterrupt:
        pass