WebSocket 流式语音合成（TTS）协议

{
  "NotificationType": "Handshake",
  "TaskId": "125xxxxxx9-tts-xxxxxxxxxxxxxxxx",
  "HandshakeResult": {
    "Code": 0,
    "Message": "success",
    "Format": "mp3",
    "SampleRate": 24000
  }
}

{ "Text": "你好，世界", "Final": false }

{
  "NotificationType": "Handshake | ProcessEof",
  "TaskId": "...",
  "HandshakeResult": { ... },
  "ProcessEofInfo":  { ... }
}

{
  "NotificationType": "ProcessEof",
  "TaskId": "...",
  "ProcessEofInfo": {
    "Code": 0,
    "Message": "finished normally"
  }
}

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""WebSocket TTS demo client.

Sign the request with TC3-HMAC-SHA256, open a WSS connection, send text,
collect audio frames (binary ws messages) and save the merged audio to a
local file. Control signals (Handshake / ProcessEof) still come as JSON text
frames; audio payload is delivered as raw binary frames without JSON wrapping.

Dependency:
    pip install websockets

Usage:
    export TTS_SECRET_ID=AKIDxxxx
    export TTS_SECRET_KEY=xxxxxxxx
    python3 wss-tts.py \\
        --host mps.cloud.tencent.com \\
        --appid your_appid \\
        --voice s1_iX4D1zb9hyem/Bp2GpyJ7cD0miMrFVWDbtKap71N8cF8M7raq9RAupL+bRWn \\
        --format mp3 \\
        --sample-rate 22050 \\
        --language zh \\
        --text "你好，世界。这是一条测试合成的文本。" \\
        --output out.mp3
"""

import argparse
import asyncio
import hashlib
import hmac
import json
import logging
import os
import random
import ssl
import struct
import sys
import time
from datetime import datetime, timezone
from typing import Dict, List, Tuple
from urllib.parse import quote

import websockets

LOG = logging.getLogger("wss-tts")

# CAM 签名链路所识别的服务名（需与服务端配置保持一致）
SERVICE_NAME = "mps"

# 规范请求中使用的 HTTP 请求方法（服务端要求小写 post）
HTTP_METHOD = "post"

# 构造规范请求头时默认使用的内容类型
CONTENT_TYPE = "application/json; charset=utf-8"


def build_signed_url(
    host: str,
    appid: str,
    secret_id: str,
    secret_key: str,
    biz_params: Dict[str, str],
    valid_seconds: int = 3600,
    use_ssl: bool = True,
) -> str:
    """Return a signed wss/ws URL ready to be passed to websockets.connect.

    biz_params holds voiceId / format / sampleRate / language / timeoutSec /
    speed / vol / resId.  The function appends secretId / nonce / timeStamp /
    expired / signature with TC3-HMAC-SHA256 signing.
    """
    timestamp = int(time.time())
    expired = timestamp + valid_seconds
    nonce = str(random.randint(10 ** 9, 10 ** 10 - 1))

    params = {k: str(v) for k, v in biz_params.items() if v not in (None, "")}
    params["secretId"] = secret_id
    params["nonce"] = nonce
    params["timeStamp"] = str(timestamp)
    params["expired"] = str(expired)

    # 规范查询字符串：键名升序排列，值进行 URL 编码处理。
    keys_sorted = sorted(params.keys())
    canonical_qs = "&".join(f"{k}={quote(params[k], safe='')}" for k in keys_sorted)

    path = f"/tts/v1/{appid}"

    # 规范请求格式必须与服务端（CAM api_sha256_v4）及可用网页客户端完全一致:
    #   post\\n
    #   /tts/v1/{appid}\\n
    #   {canonical_qs}\\n
    #   content-type:application/json; charset=utf-8\\n
    #   host:{host}\\n
    #   \\n                     <- blank line that separates headers and signed_headers
    #   content-type;host\\n    <- signed_headers, followed by an empty trailing line
    canonical_request = "\\n".join([
        HTTP_METHOD,
        path,
        canonical_qs,
        f"content-type:{CONTENT_TYPE}",
        f"host:{host}",
        "",
        "content-type;host",
        "",
    ])

    date = datetime.fromtimestamp(timestamp, tz=timezone.utc).strftime("%Y-%m-%d")
    credential_scope = f"{date}/{SERVICE_NAME}/tc3_request"
    string_to_sign = "\\n".join([
        "TC3-HMAC-SHA256",
        str(timestamp),
        credential_scope,
        hashlib.sha256(canonical_request.encode("utf-8")).hexdigest(),
    ])

    k_date = hmac.new(("TC3" + secret_key).encode("utf-8"), date.encode("utf-8"), hashlib.sha256).digest()
    k_service = hmac.new(k_date, SERVICE_NAME.encode("utf-8"), hashlib.sha256).digest()
    k_signing = hmac.new(k_service, b"tc3_request", hashlib.sha256).digest()
    signature = hmac.new(k_signing, string_to_sign.encode("utf-8"), hashlib.sha256).hexdigest()

    # 拼接签名，生成最终查询字符串
    final_qs = "&".join(f"{k}={quote(params[k], safe='')}" for k in keys_sorted)
    final_qs += f"&signature={signature}"

    scheme = "wss" if use_ssl else "ws"
    return f"{scheme}://{host}{path}?{final_qs}"


def wrap_pcm_to_wav(pcm: bytes, sample_rate: int, sample_width: int = 2, channels: int = 1) -> bytes:
    """Wrap raw PCM s16le bytes with a minimal RIFF/WAVE header for playback."""
    byte_rate = sample_rate * channels * sample_width
    block_align = channels * sample_width
    data_size = len(pcm)
    riff_size = 36 + data_size
    header = b"RIFF" + struct.pack("<I", riff_size) + b"WAVE"
    fmt_chunk = b"fmt " + struct.pack("<IHHIIHH", 16, 1, channels, sample_rate, byte_rate, block_align, sample_width * 8)
    data_chunk = b"data" + struct.pack("<I", data_size)
    return header + fmt_chunk + data_chunk + pcm


def redact_for_log(payload: dict) -> str:
    """Render a control-signal JSON payload for logging.

    Audio bytes are delivered as binary ws frames now, so this only needs to
    pretty-print Handshake / ProcessEof JSON dicts.
    """
    if not isinstance(payload, dict):
        return json.dumps(payload, ensure_ascii=False)
    return json.dumps(payload, ensure_ascii=False)


async def run_session(url: str, texts: List[str], output: str) -> int:
    """Run one TTS session and persist the merged audio."""
    audio_chunks: List[bytes] = []
    handshake_format = ""
    handshake_sample_rate = 0
    final_code = -1
    final_message = ""

    ssl_ctx = ssl.create_default_context() if url.startswith("wss://") else None
    LOG.info("connect %s", url)
    ws = await websockets.connect(url, ssl=ssl_ctx, max_size=None)
    try:
        # 步骤 1：需等待握手应答（Handshake ack）返回后，再发送业务文本数据
        first = await ws.recv()
        first_obj = json.loads(first) if isinstance(first, str) else {}
        LOG.info("recv handshake: %s", redact_for_log(first_obj))
        if first_obj.get("NotificationType") != "Handshake":
            LOG.error("first frame is not Handshake: %s", first_obj)
            return 1
        hs = first_obj.get("HandshakeResult") or {}
        if hs.get("Code") != 0:
            LOG.error("handshake failed code=%s message=%s", hs.get("Code"), hs.get("Message"))
            return 1
        handshake_format = (hs.get("Format") or "").lower()
        handshake_sample_rate = int(hs.get("SampleRate") or 0)
        LOG.info("handshake ok format=%s sampleRate=%d", handshake_format, handshake_sample_rate)

        # 步骤 2：发送全部文本分片，最后发送空数据包标识数据流结束（EOS）
        for idx, seg in enumerate(texts):
            pkt = {"Text": seg, "Final": False}
            await ws.send(json.dumps(pkt, ensure_ascii=False))
            LOG.info("send #%d: %s", idx + 1, pkt)
        eos = {"Text": "", "Final": True}
        await ws.send(json.dumps(eos))
        LOG.info("send EOS: %s", eos)

        # 步骤 3：持续接收数据帧，直至收到结束帧（ProcessEof）
        # - 二进制 WebSocket 帧  -> 原始音频字节流，直接追加拼接
        # - 文本 WebSocket 帧    -> 控制类 JSON 报文（握手消息 / 结束指令 ProcessEof）
        while True:
            try:
                msg = await ws.recv()
            except websockets.ConnectionClosed as exc:
                LOG.info("ws closed: code=%s reason=%s", exc.code, exc.reason)
                break
            if isinstance(msg, (bytes, bytearray)):
                chunk = bytes(msg)
                if not chunk:
                    continue
                audio_chunks.append(chunk)
                LOG.info("audio chunk %d bytes (total=%d)",
                         len(chunk), sum(len(c) for c in audio_chunks))
                continue
            try:
                obj = json.loads(msg)
            except json.JSONDecodeError as exc:
                LOG.warning("ignore non-json text frame: %s err=%s", msg, exc)
                continue
            LOG.info("recv: %s", redact_for_log(obj))
            ntype = obj.get("NotificationType")
            if ntype == "ProcessEof":
                info = obj.get("ProcessEofInfo") or {}
                final_code = int(info.get("Code", -1))
                final_message = info.get("Message", "")
                LOG.info("process eof code=%d message=%s, exit", final_code, final_message)
                break
            else:
                LOG.warning("unknown NotificationType: %s", ntype)
    finally:
        # 强制关闭连接，不再等待服务端关闭帧；
        # 收到 ProcessEof 后进程立即退出，
        # 避免在优雅关闭握手上阻塞，防止卡死 asyncio 事件循环
        try:
            await asyncio.wait_for(ws.close(), timeout=1.0)
        except Exception:
            try:
                ws.transport.abort()
            except Exception:
                pass

    if not audio_chunks:
        LOG.error("no audio received, code=%d message=%s", final_code, final_message)
        return 2

    merged = b"".join(audio_chunks)
    out_path = output
    # 原始 PCM 响应数据需封装 WAV 文件头，以便直接播放
    if handshake_format == "pcm":
        merged = wrap_pcm_to_wav(merged, handshake_sample_rate or 16000)
        if not out_path.lower().endswith(".wav"):
            out_path = out_path + ".wav"
    with open(out_path, "wb") as fp:
        fp.write(merged)
    LOG.info("saved %d bytes to %s (format=%s sampleRate=%d code=%d)",
             len(merged), out_path, handshake_format, handshake_sample_rate, final_code)
    return 0 if final_code == 0 else 3


def parse_args() -> Tuple[argparse.Namespace, Dict[str, str]]:
    """Parse CLI arguments and assemble the business URL query params."""
    parser = argparse.ArgumentParser(description="WebSocket TTS demo client")
    parser.add_argument("--host", default="mps.cloud.tencent.com", help="server host")
    parser.add_argument("--appid", required=True, help="tencent cloud appid")
    parser.add_argument("--secret-id", default=os.environ.get("TTS_SECRET_ID", ""),
                        help="cam secret id (env TTS_SECRET_ID)")
    parser.add_argument("--secret-key", default=os.environ.get("TTS_SECRET_KEY", ""),
                        help="cam secret key (env TTS_SECRET_KEY)")
    parser.add_argument("--voice", required=True, help="voiceId, e.g. ws_voice_101")
    parser.add_argument("--format", default="mp3", choices=["pcm", "mp3", "wav", "flac", "opus", "ulaw", "alaw"],
                        help="audio format (server may downgrade, see Handshake ack)")
    parser.add_argument("--sample-rate", type=int, default=24000, help="audio sample rate (Hz)")
    parser.add_argument("--language", default="zh", help="language hint, e.g. zh / en / auto")
    parser.add_argument("--timeout-sec", type=int, default=60, help="session idle timeout (<=120)")
    parser.add_argument("--speed", type=float, default=0.0, help="speech speed, 0 = use upstream default")
    parser.add_argument("--vol", type=float, default=0.0, help="speech volume, 0 = use upstream default")
    parser.add_argument("--res-id", default="", help="optional resource tag for billing")
    parser.add_argument("--text", action="append", default=None,
                        help="text segment, can be repeated to send multiple segments")
    parser.add_argument("--output", default="out.bin", help="output file path")
    parser.add_argument("--no-ssl", action="store_true", help="use ws:// instead of wss://")
    parser.add_argument("--verbose", action="store_true", help="enable debug logging")
    args = parser.parse_args()

    if not args.secret_id or not args.secret_key:
        parser.error("secret-id/secret-key required (or via env TTS_SECRET_ID/TTS_SECRET_KEY)")
    if not args.text:
        args.text = ["真正的危险不是计算机开始像人一样思考，而是人开始像计算机一样思考。"]

    biz: Dict[str, str] = {
        "voiceId": args.voice,
        "format": args.format,
        "sampleRate": str(args.sample_rate),
        "timeoutSec": str(args.timeout_sec),
    }
    if args.language:
        biz["language"] = args.language
    if args.res_id:
        biz["resId"] = args.res_id
    if args.speed > 0:
        biz["speed"] = str(args.speed)
    if args.vol > 0:
        biz["vol"] = str(args.vol)
    return args, biz


def main() -> int:
    args, biz = parse_args()
    logging.basicConfig(
        level=logging.DEBUG if args.verbose else logging.INFO,
        format="%(asctime)s %(levelname)s %(message)s",
    )
    url = build_signed_url(
        host=args.host,
        appid=args.appid,
        secret_id=args.secret_id,
        secret_key=args.secret_key,
        biz_params=biz,
        use_ssl=not args.no_ssl,
    )
    return run_async(run_session(url, args.text, args.output))


def run_async(coro) -> int:
    """Run an async coroutine on Python >=3.6 (asyncio.run only exists in 3.7+)."""
    if hasattr(asyncio, "run"):
        return asyncio.run(coro)
    loop = asyncio.new_event_loop()
    try:
        asyncio.set_event_loop(loop)
        return loop.run_until_complete(coro)
    finally:
        # 关闭事件循环前取消所有未完成任务，否则
        # 后台任务（如 WebSocket 关闭连接）可能在已关闭的循环上执行回调，
        # 导致程序退出时产生大量冗余报错堆栈
        try:
            pending = [t for t in asyncio.Task.all_tasks(loop) if not t.done()]
            for t in pending:
                t.cancel()
            if pending:
                loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
        except Exception:
            pass
        try:
            loop.run_until_complete(loop.shutdown_asyncgens())
        except Exception:
            pass
        loop.close()
        asyncio.set_event_loop(None)


if __name__ == "__main__":
    sys.exit(main())