ai-video-fullstack/backend/routes/voice_webrtc.py

"""WebRTC 输出:SmallWebRTC 信令握手。

参考 dograh 的 webrtc_signaling.py,砍掉鉴权/配额/DB/org/ICE 过滤策略/TURN。
握手消息:
    client → {type:"offer",         payload:{pc_id, sdp, type, assistant_id}}
    server → {type:"answer",        payload:{pc_id, sdp, type}}
    both   → {type:"ice-candidate", payload:{pc_id, candidate:{...}}}
    server → {type:"error",         payload:{message}}
"""

import base64
import json
from collections.abc import Coroutine
from typing import Any

from fastapi import APIRouter, Body, Depends, HTTPException, Request, WebSocket
from loguru import logger
from models import AssistantConfig, SignalingOffer
from services.auth import require_admin, require_admin_websocket
from services.realtime import lifecycle as realtime_lifecycle
from services.realtime.launcher import (
    resolve_assistant_config,
    validate_visual_runtime,
)
from services.runtime_variables import DynamicVariableError, prepare_dynamic_config
from starlette.websockets import WebSocketDisconnect, WebSocketState

from services.webrtc_ice import aiortc_ice_servers, client_ice_servers

# 注意:pipecat 是重依赖(语音才用),在 _handle_offer 等处惰性导入。

router = APIRouter(tags=["voice"])
_http_peers: dict[str, object] = {}
_active_connections = realtime_lifecycle.active_connections
_pipeline_tasks = realtime_lifecycle.pipeline_tasks
_connection_tasks = realtime_lifecycle.connection_tasks
PIPELINE_CLOSE_GRACE_SECONDS = 10.0
MAX_CLIENT_TOOL_NAME_LENGTH = 64


def _start_pipeline_task(
    connection: object,
    coroutine: Coroutine[Any, Any, None],
) -> object:
    """Compatibility proxy for existing RTVI callers and tests."""
    return realtime_lifecycle.start_pipeline_task(
        connection,
        coroutine,
        protocol="webrtc",
    )


async def _wait_for_pipeline_close(
    task,
    *,
    connection_id: str,
) -> None:
    await realtime_lifecycle.wait_for_pipeline_close(
        task,
        connection_id=connection_id,
        timeout=PIPELINE_CLOSE_GRACE_SECONDS,
    )


async def shutdown_active_sessions() -> None:
    """Drain RTVI and OpenAI Realtime sessions."""
    _http_peers.clear()
    await realtime_lifecycle.shutdown_active_sessions(
        timeout=PIPELINE_CLOSE_GRACE_SECONDS,
    )


def _decode_client_tool_names(encoded: str | None) -> list[str]:
    """Decode client capabilities carried outside the WebRTC SDP body.

    Pipecat's browser transport replaces the body of a ``Request`` endpoint
    with its SDP offer and drops ``APIRequest.requestData``. Request headers are
    preserved, so they are the single metadata channel for this path.
    """

    if encoded is None:
        return []
    try:
        padding = "=" * (-len(encoded) % 4)
        value = json.loads(
            base64.urlsafe_b64decode(encoded + padding).decode("utf-8")
        )
    except (ValueError, UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise HTTPException(
            status_code=400,
            detail="客户端工具能力请求头格式无效",
        ) from exc
    if not isinstance(value, list):
        raise HTTPException(
            status_code=400,
            detail="客户端工具能力必须是函数名列表",
        )

    names: list[str] = []
    seen: set[str] = set()
    for item in value:
        if not isinstance(item, str):
            raise HTTPException(
                status_code=400,
                detail="客户端工具函数名必须是字符串",
            )
        name = item.strip()
        if not name or len(name) > MAX_CLIENT_TOOL_NAME_LENGTH:
            raise HTTPException(
                status_code=400,
                detail="客户端工具函数名为空或过长",
            )
        if name not in seen:
            names.append(name)
            seen.add(name)
    return names


@router.get("/api/webrtc/ice-servers", dependencies=[Depends(require_admin)])
async def ice_servers():
    """Browser fetches STUN/TURN config (with ephemeral TURN creds when configured)."""
    return {"iceServers": client_ice_servers()}


@router.post("/api/webrtc/offer", dependencies=[Depends(require_admin)])
async def http_offer(request: Request, payload: dict = Body(...)):
    """Official SmallWebRTC JS transport offer endpoint."""
    request_data = payload.get("requestData") or {}
    encoded_variables = request.headers.get("x-pipecat-dynamic-variables", "")
    header_variables = {}
    if encoded_variables:
        padding = "=" * (-len(encoded_variables) % 4)
        header_variables = json.loads(
            base64.urlsafe_b64decode(encoded_variables + padding).decode("utf-8")
        )
    client_tool_names = _decode_client_tool_names(
        request.headers.get("x-pipecat-client-tools")
    )
    logger.debug(f"WebRTC 客户端工具上报: {client_tool_names or '<none>'}")
    offer_payload = {
        **payload,
        "assistant_id": request.headers.get("x-pipecat-assistant-id"),
        "dynamic_variables": header_variables,
        **request_data,
        # Client capabilities have one canonical transport: the request header.
        "client_tool_names": client_tool_names,
    }
    answer = await _handle_offer_payload(offer_payload, _http_peers)
    return answer


@router.patch("/api/webrtc/offer", dependencies=[Depends(require_admin)])
async def http_ice_candidates(payload: dict = Body(...)):
    """Accept the batched trickle ICE format used by the official JS client."""
    pc_id = payload.get("pc_id")
    pc = _http_peers.get(pc_id) if pc_id else None
    if not pc:
        return {"ok": False}

    from aiortc.sdp import candidate_from_sdp

    for item in payload.get("candidates") or []:
        candidate = candidate_from_sdp(item.get("candidate", ""))
        candidate.sdpMid = item.get("sdp_mid")
        candidate.sdpMLineIndex = item.get("sdp_mline_index")
        await pc.add_ice_candidate(candidate)
    return {"ok": True}


@router.websocket("/ws/voice")
async def voice_signaling(websocket: WebSocket):
    if not await require_admin_websocket(websocket):
        return
    await websocket.accept()
    peers: dict = {}
    try:
        while True:
            message = await websocket.receive_json()
            try:
                if message.get("type") == "offer":
                    await _handle_offer(websocket, message.get("payload", {}), peers)
                elif message.get("type") == "ice-candidate":
                    await _handle_ice(message.get("payload", {}), peers)
            except Exception as e:
                logger.exception(f"处理 WebRTC 信令消息失败: {e}")
                if websocket.application_state == WebSocketState.CONNECTED:
                    await websocket.send_json(
                        {
                            "type": "error",
                            "payload": {
                                "message": (
                                    str(e)
                                    if isinstance(e, DynamicVariableError)
                                    else f"语音会话启动失败: {type(e).__name__}"
                                )
                            },
                        }
                    )
    except WebSocketDisconnect:
        logger.info("WebRTC 信令断开")
    except Exception as e:
        logger.error(f"WebRTC 信令出错: {e}")
    finally:
        # disconnect() triggers the registered closed callback, which removes
        # the peer from this dict. Iterate over a snapshot to avoid mutation.
        for pc in list(peers.values()):
            await pc.disconnect()


async def _resolve_config(offer: SignalingOffer) -> AssistantConfig:
    """优先用 assistant_id 从 DB 解析(含真 key);否则用调试内联配置。"""
    if offer.assistant_id:
        return await resolve_assistant_config(
            offer.assistant_id,
            dynamic_variables=offer.dynamic_variables,
        )
    if offer.inline_config:
        return prepare_dynamic_config(
            offer.inline_config,
            offer.dynamic_variables,
            assistant_id=None,
        )
    raise ValueError("offer 缺少 assistant_id 或 inline_config")


async def _handle_offer(websocket, payload, peers):
    answer = await _handle_offer_payload(payload, peers)
    if websocket.application_state == WebSocketState.CONNECTED:
        await websocket.send_json(
            {
                "type": "answer",
                "payload": answer,
            }
        )


async def _handle_offer_payload(payload, peers):
    from pipecat.transports.smallwebrtc.connection import SmallWebRTCConnection
    from services.pipecat.pipeline import run_pipeline
    from services.pipecat.transports import build_webrtc_transport

    offer = SignalingOffer(**payload)
    pc_id = offer.pc_id

    restart_pc = bool(payload.get("restart_pc"))
    if restart_pc and pc_id and pc_id in peers:
        old_pc = peers.pop(pc_id)
        old_task = _connection_tasks.get(old_pc)
        await old_pc.disconnect()
        await _wait_for_pipeline_close(
            old_task,
            connection_id=str(getattr(old_pc, "pc_id", pc_id)),
        )

    if pc_id and pc_id in peers:
        pc = peers[pc_id]
        await pc.renegotiate(sdp=offer.sdp, type=offer.type, restart_pc=False)
    else:
        cfg = await _resolve_config(offer)  # 解析放在建连前,配置错就别建连
        # 服务端助手配置是视觉理解的唯一授权来源；客户端 offer 只负责携带媒体轨。
        vision_enabled = validate_visual_runtime(cfg)
        pc = SmallWebRTCConnection(ice_servers=aiortc_ice_servers())
        if pc_id:
            pc._pc_id = pc_id
        await pc.initialize(sdp=offer.sdp, type=offer.type)
        peers[pc.pc_id] = pc

        # 后台跑管线:WebRTC transport + 解析出的运行时配置
        transport = build_webrtc_transport(
            pc,
            video_in_enabled=vision_enabled,
        )
        pipeline_task = _start_pipeline_task(
            pc,
            run_pipeline(
                transport,
                cfg,
                vision_enabled=vision_enabled,
                assistant_id=offer.assistant_id,
                channel="webrtc",
                client_tool_names=set(offer.client_tool_names),
            ),
        )

        @pc.event_handler("closed")
        async def _on_closed(conn: SmallWebRTCConnection):
            peers.pop(conn.pc_id, None)
            _active_connections.discard(conn)
            await _wait_for_pipeline_close(
                pipeline_task,
                connection_id=conn.pc_id,
            )

    answer = pc.get_answer()
    return {
        "pc_id": answer["pc_id"],
        "sdp": answer["sdp"],
        "type": answer["type"],
    }


async def _handle_ice(payload, peers):
    from aiortc.sdp import candidate_from_sdp

    pc_id = payload.get("pc_id")
    candidate_data = payload.get("candidate")
    pc = peers.get(pc_id) if pc_id else None
    if not pc or not candidate_data:
        return
    try:
        candidate = candidate_from_sdp(candidate_data.get("candidate", ""))
        candidate.sdpMid = candidate_data.get("sdpMid")
        candidate.sdpMLineIndex = candidate_data.get("sdpMLineIndex")
        await pc.add_ice_candidate(candidate)
    except Exception as e:
        logger.error(f"添加 ICE candidate 失败: {e}")