diff --git a/runtime/python/websocket/README.md b/runtime/python/websocket/README.md index 2c130466c..bc852887c 100644 --- a/runtime/python/websocket/README.md +++ b/runtime/python/websocket/README.md @@ -68,9 +68,12 @@ python funasr_wss_client.py \ --audio_in [if set, loadding from wav.scp, else recording from mircrophone] \ --output_dir [if set, write the results to output_dir] \ --mode [`online` for streaming asr, `offline` for non-streaming, `2pass` for unifying streaming and non-streaming asr] \ ---thread_num [thread_num for send data] +--thread_num [thread_num for send data] \ +--result_timeout [seconds to wait for the final server acknowledgement] ``` +When `--audio_in` is set, the client sends `{"is_speaking": false, "is_end": true}` after the last audio frame. The server flushes pending online and offline inference before replying with `{"is_end": true, "is_final": true}`. If inference fails, the acknowledgement contains `{"is_end": true, "is_final": false, "error": "..."}`. The client waits up to `--result_timeout` seconds for this acknowledgement; the default is 300 seconds. + #### Usage examples ##### ASR offline client Recording from mircrophone diff --git a/runtime/python/websocket/funasr_wss_client.py b/runtime/python/websocket/funasr_wss_client.py index 884579d2b..b95c3867e 100644 --- a/runtime/python/websocket/funasr_wss_client.py +++ b/runtime/python/websocket/funasr_wss_client.py @@ -50,6 +50,12 @@ parser.add_argument("--ssl", type=int, default=1, help="1 for ssl connect, 0 for no ssl") parser.add_argument("--use_itn", type=int, default=1, help="1 for using itn, 0 for not itn") parser.add_argument("--mode", type=str, default="2pass", help="offline, online, 2pass") +parser.add_argument( + "--result_timeout", + type=float, + default=300.0, + help="seconds to wait for the server's end-of-input acknowledgement", +) # ✅ 验收日志输出目录（每个 meeting 单独写，避免多进程抢文件） parser.add_argument("--log_dir", type=str, default="./asr_logs", help="验收日志输出目录") @@ -185,7 +191,7 @@ async def record_microphone(): await asyncio.sleep(0.01) -async def record_from_scp(chunk_begin, chunk_size): +async def record_from_scp(chunk_begin, chunk_size, completion_queue): """从 wav/scp 文件读取音频分片发送，用于压测和延迟测试""" global voices, latency_first_audio_time, latency_last_audio_time if args.audio_in.endswith(".scp"): @@ -268,7 +274,6 @@ async def record_from_scp(chunk_begin, chunk_size): ) await websocket.send(message) - is_speaking = True # 初始化该 wav 的统计状态 latency_first_audio_time[wav_name] = None @@ -286,10 +291,6 @@ async def record_from_scp(chunk_begin, chunk_size): await websocket.send(data) - if i == chunk_num - 1: - is_speaking = False - await websocket.send(json.dumps({"is_speaking": is_speaking}, ensure_ascii=False)) - # ✅ sleep策略：默认按实时节奏；若开启 send_without_sleep 则几乎不 sleep（压测） if args.send_without_sleep: sleep_duration = 0.001 @@ -301,20 +302,30 @@ async def record_from_scp(chunk_begin, chunk_size): ) await asyncio.sleep(sleep_duration) - if not args.mode == "offline": - await asyncio.sleep(2) - - if args.mode == "offline": - global offline_msg_done - while not offline_msg_done: - await asyncio.sleep(1) - - await asyncio.sleep(10) + await websocket.send( + json.dumps({"is_speaking": False, "is_end": True}, ensure_ascii=False) + ) + try: + acknowledgement = await asyncio.wait_for( + completion_queue.get(), timeout=args.result_timeout + ) + except asyncio.TimeoutError as e: + await websocket.close() + raise TimeoutError( + f"server did not acknowledge end of input for {wav_name!r} " + f"within {args.result_timeout:g} seconds" + ) from e + if not acknowledgement.get("is_final", False): + await websocket.close() + error = acknowledgement.get("error") or "server did not finalize input" + raise RuntimeError( + f"server failed to finalize {wav_name!r}: {error}" + ) await websocket.close() -async def message(id, writer: MeetingWriter): +async def message(id, writer: MeetingWriter, completion_queue): """接收服务端识别结果 + 打印实时文本 + 打印延迟 + 写验收日志(events.jsonl)""" import websockets global websocket, voices, offline_msg_done @@ -324,6 +335,13 @@ async def message(id, writer: MeetingWriter): text_print = "" text_print_2pass_online = "" text_print_2pass_offline = "" + end_ack_received = False + + def release_sender_with_error(error): + if completion_queue is not None and not end_ack_received: + completion_queue.put_nowait( + {"is_end": True, "is_final": False, "error": error} + ) if args.output_dir is not None: ibest_writer = open( @@ -372,6 +390,10 @@ async def message(id, writer: MeetingWriter): timestamp = meg.get("timestamp", "") offline_msg_done = meg.get("is_final", False) + if meg.get("is_end"): + end_ack_received = True + if completion_queue is not None: + completion_queue.put_nowait(meg) # ✅ 验收友好：每条消息落 events.jsonl（便于后处理） event = { @@ -460,8 +482,10 @@ async def message(id, writer: MeetingWriter): except websockets.exceptions.ConnectionClosedOK: print(f"[MEETING {id}] connection closed normally") + release_sender_with_error("connection closed before end-of-input acknowledgement") except Exception as e: print(f"[MEETING {id}] Exception:", e) + release_sender_with_error(f"receiver failed before end-of-input acknowledgement: {e}") finally: try: if ibest_writer is not None: @@ -497,12 +521,15 @@ async def ws_client(id, chunk_begin, chunk_size): ) as websocket: meeting_tag = f"{id}_{i}" writer = MeetingWriter(args.log_dir, meeting_id=meeting_tag, flush_every=args.log_flush_every) + completion_queue = asyncio.Queue() if args.audio_in is not None else None try: if args.audio_in is not None: - task = asyncio.create_task(record_from_scp(i, 1)) + task = asyncio.create_task(record_from_scp(i, 1, completion_queue)) else: task = asyncio.create_task(record_microphone()) - task3 = asyncio.create_task(message(str(id) + "_" + str(i), writer)) # processid+fileid + task3 = asyncio.create_task( + message(str(id) + "_" + str(i), writer, completion_queue) + ) # processid+fileid await asyncio.gather(task, task3) finally: writer.close() diff --git a/runtime/python/websocket/funasr_wss_server.py b/runtime/python/websocket/funasr_wss_server.py index ae5cad2fb..2a449e642 100644 --- a/runtime/python/websocket/funasr_wss_server.py +++ b/runtime/python/websocket/funasr_wss_server.py @@ -346,6 +346,7 @@ async def ws_serve(websocket, path=None): frames = [] frames_asr = [] frames_asr_online = [] + pending_offline_audio = [] global websocket_users websocket_users.add(websocket) @@ -358,6 +359,8 @@ async def ws_serve(websocket, path=None): websocket.vad_pre_idx = 0 speech_start = False speech_end_i = -1 + online_needs_finalization = False + session_errors = [] websocket.wav_name = "microphone" websocket.mode = "2pass" @@ -374,6 +377,88 @@ async def ws_serve(websocket, path=None): print("new user connected", flush=True) + def record_error(message): + if message not in session_errors: + session_errors.append(message) + + async def finalize_online_segment(): + nonlocal frames_asr_online, online_needs_finalization + + if websocket.mode not in ("2pass", "online") or not online_needs_finalization: + return + + websocket.status_dict_asr_online["is_final"] = True + try: + await async_asr_online(websocket, b"".join(frames_asr_online)) + except Exception as e: + print("error in final asr streaming:", e) + record_error(f"online inference failed: {e}") + + frames_asr_online = [] + websocket.status_dict_asr_online["cache"] = {} + websocket.status_dict_asr_online["is_final"] = False + online_needs_finalization = False + + async def finish_input(send_end_ack): + nonlocal frames, frames_asr, frames_asr_online, pending_offline_audio + nonlocal speech_start, speech_end_i, online_needs_finalization + + await finalize_online_segment() + + if websocket.mode in ("2pass", "offline"): + audio_in = b"".join(frames_asr) + if not audio_in: + audio_in = b"".join(pending_offline_audio) + + if audio_in: + if websocket.save_offline_segments and audio_in: + try: + await run_blocking( + save_offline_wav_segment_sync, + websocket, + audio_in, + "not_speaking", + sem=SEM_WAV, + ) + except Exception as e: + print("[SAVE_OFFLINE_SEG] async failed:", e) + + try: + await async_asr(websocket, audio_in) + pending_offline_audio = [] + except Exception as e: + print("error in final asr offline:", e) + record_error(f"offline inference failed: {e}") + + errors = list(session_errors) + + frames = [] + frames_asr = [] + frames_asr_online = [] + pending_offline_audio = [] + speech_start = False + speech_end_i = -1 + online_needs_finalization = False + websocket.vad_pre_idx = 0 + websocket.status_dict_asr_online["cache"] = {} + websocket.status_dict_vad["cache"] = {} + + if send_end_ack: + acknowledgement = { + "mode": websocket.mode, + "wav_name": websocket.wav_name, + "is_final": not errors, + "is_end": True, + } + if errors: + acknowledgement["error"] = "; ".join(errors) + await websocket.send( + json.dumps(acknowledgement, ensure_ascii=False) + ) + session_errors.clear() + elif errors: + raise RuntimeError("; ".join(errors)) + try: async for message in websocket: # ========== 1) 先处理“文本配置消息” ========== @@ -386,9 +471,11 @@ async def ws_serve(websocket, path=None): print("=============messagejson============", messagejson) + end_of_input = False if "is_speaking" in messagejson: websocket.is_speaking = bool(messagejson["is_speaking"]) websocket.status_dict_asr_online["is_final"] = (not websocket.is_speaking) + end_of_input = not websocket.is_speaking if "chunk_interval" in messagejson: websocket.chunk_interval = _safe_int( @@ -421,16 +508,28 @@ async def ws_serve(websocket, path=None): print(f"热词已更新: {hotword_data}") if "mode" in messagejson: - websocket.mode = messagejson["mode"] or websocket.mode + requested_mode = messagejson["mode"] + if requested_mode and requested_mode not in ("online", "offline", "2pass"): + websocket.mode = requested_mode + record_error(f"unsupported mode: {requested_mode!r}") + else: + websocket.mode = requested_mode or websocket.mode if "audio_fs" in messagejson: websocket.audio_fs = _safe_int(messagejson["audio_fs"], 16000) + if end_of_input: + await finish_input(send_end_ack=bool(messagejson.get("is_end"))) + continue # ========== 2) 处理“二进制音频消息” ========== + if websocket.mode not in ("online", "offline", "2pass"): + continue + if "chunk_size" not in websocket.status_dict_asr_online: print("[WARN] chunk_size not set yet, skip audio frame (send config first).") + record_error("audio frame discarded: chunk_size is not configured") continue try: @@ -439,16 +538,21 @@ async def ws_serve(websocket, path=None): ) except Exception as e: print("[WARN] set vad chunk_size failed:", e) + record_error(f"audio frame discarded: invalid VAD chunk_size: {e}") continue pcm = message frames.append(pcm) + if websocket.mode in ("2pass", "offline"): + pending_offline_audio.append(pcm) duration_ms = _pcm_duration_ms(pcm, fs=websocket.audio_fs, ch=1, sampwidth=2) websocket.vad_pre_idx += duration_ms # online asr frames_asr_online.append(pcm) + if websocket.mode in ("2pass", "online"): + online_needs_finalization = True websocket.status_dict_asr_online["is_final"] = (speech_end_i != -1) if (len(frames_asr_online) % websocket.chunk_interval == 0) or websocket.status_dict_asr_online["is_final"]: @@ -456,8 +560,9 @@ async def ws_serve(websocket, path=None): audio_in = b"".join(frames_asr_online) try: await async_asr_online(websocket, audio_in) - except Exception: + except Exception as e: print(f"error in asr streaming, {websocket.status_dict_asr_online}") + record_error(f"online inference failed: {e}") frames_asr_online = [] if speech_start: @@ -468,6 +573,7 @@ async def ws_serve(websocket, path=None): speech_start_i, speech_end_i = await async_vad(websocket, pcm) except Exception as e: print("error in vad:", e) + record_error(f"vad inference failed: {e}") speech_start_i, speech_end_i = -1, -1 if speech_start_i != -1: @@ -482,8 +588,12 @@ async def ws_serve(websocket, path=None): # ========== 3) 2pass：离线阶段触发点 ========== if (speech_end_i != -1) or (not websocket.is_speaking): + await finalize_online_segment() + if websocket.mode in ("2pass", "offline"): audio_in = b"".join(frames_asr) + if not audio_in and speech_end_i != -1: + audio_in = b"".join(pending_offline_audio) reason = "vad_end" if speech_end_i != -1 else "not_speaking" # 保存 wav：放线程池，避免磁盘 IO 卡 loop @@ -499,21 +609,26 @@ async def ws_serve(websocket, path=None): except Exception as e: print("[SAVE_OFFLINE_SEG] async failed:", e) - try: - await async_asr(websocket, audio_in) - except Exception as e: - print("error in asr offline:", e) + if audio_in: + try: + await async_asr(websocket, audio_in) + pending_offline_audio = [] + except Exception as e: + print("error in asr offline:", e) + record_error(f"offline inference failed: {e}") frames_asr = [] speech_start = False frames_asr_online = [] websocket.status_dict_asr_online["cache"] = {} + websocket.status_dict_asr_online["is_final"] = False + online_needs_finalization = False + speech_end_i = -1 if not websocket.is_speaking: websocket.vad_pre_idx = 0 frames = [] websocket.status_dict_vad["cache"] = {} - speech_end_i = -1 else: frames = frames[-20:] @@ -524,6 +639,11 @@ async def ws_serve(websocket, path=None): websocket_users.remove(websocket) except websockets.InvalidState: print("InvalidState...") + try: + await ws_reset(websocket) + except Exception: + pass + websocket_users.discard(websocket) except Exception as e: print("Exception:", e) try: @@ -663,11 +783,7 @@ async def async_asr(websocket, audio_in: bytes): if punc_array is not None: message["punc_array"] = to_python(punc_array) - try: - await websocket.send(json.dumps(message, ensure_ascii=False)) - except Exception as e: - print("send json failed:", e) - print("message types:", {k: type(v) for k, v in message.items()}) + await websocket.send(json.dumps(message, ensure_ascii=False)) else: message = { "mode": mode, @@ -681,7 +797,7 @@ async def async_asr(websocket, audio_in: bytes): async def async_asr_online(websocket, audio_in: bytes): - if len(audio_in) <= 0: + if len(audio_in) <= 0 and not websocket.status_dict_asr_online.get("is_final", False): return # streaming generate 也是阻塞：线程池执行 diff --git a/tests/test_websocket_file_finalization.py b/tests/test_websocket_file_finalization.py new file mode 100644 index 000000000..031a4adbd --- /dev/null +++ b/tests/test_websocket_file_finalization.py @@ -0,0 +1,760 @@ +import ast +import asyncio +import json +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + + +ROOT = Path(__file__).resolve().parents[1] +SERVER_PATH = ROOT / "runtime/python/websocket/funasr_wss_server.py" +CLIENT_PATH = ROOT / "runtime/python/websocket/funasr_wss_client.py" +README_PATH = ROOT / "runtime/python/websocket/README.md" + + +def _load_async_function(path, name, namespace): + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + function = next( + node + for node in tree.body + if isinstance(node, ast.AsyncFunctionDef) and node.name == name + ) + exec( + compile(ast.Module(body=[function], type_ignores=[]), str(path), "exec"), + namespace, + ) + return namespace[name] + + +class FakeConnectionClosed(Exception): + pass + + +class FakeInvalidState(Exception): + pass + + +class FakeWebSocket: + def __init__(self, incoming=()): + self.incoming = iter(incoming) + self.sent = [] + self.closed = False + self.path = None + + def __aiter__(self): + return self + + async def __anext__(self): + try: + return next(self.incoming) + except StopIteration: + raise StopAsyncIteration + + async def send(self, message): + self.sent.append(message) + + async def close(self): + self.closed = True + + +def _server_namespace( + asr_calls, + online_calls, + vad_results=None, + asr_error=None, + online_error=None, + asr_error_once=None, + online_error_once=None, + online_cache_states=None, + reset_calls=None, + vad_calls=None, +): + vad_results = iter(vad_results or []) + online_cache_states = online_cache_states if online_cache_states is not None else [] + reset_calls = reset_calls if reset_calls is not None else [] + vad_calls = vad_calls if vad_calls is not None else [] + + async def async_vad(websocket, pcm): + vad_calls.append(pcm) + return next(vad_results, (0, -1)) + + async def async_asr(websocket, audio): + asr_calls.append(audio) + if asr_error is not None: + raise asr_error + if asr_error_once is not None and len(asr_calls) == 1: + raise asr_error_once + await websocket.send( + json.dumps( + { + "mode": ( + "2pass-offline" if websocket.mode == "2pass" else websocket.mode + ), + "text": "offline final", + "wav_name": websocket.wav_name, + "is_final": True, + } + ) + ) + + async def async_asr_online(websocket, audio): + online_calls.append(audio) + online_cache_states.append(dict(websocket.status_dict_asr_online["cache"])) + if online_error is not None: + raise online_error + if online_error_once is not None and len(online_calls) == 1: + raise online_error_once + if not websocket.status_dict_asr_online.get("is_final", False): + websocket.status_dict_asr_online["cache"]["partial"] = True + await websocket.send( + json.dumps( + { + "mode": ( + "2pass-online" if websocket.mode == "2pass" else websocket.mode + ), + "text": "online final", + "wav_name": websocket.wav_name, + "is_final": True, + } + ) + ) + + async def ws_reset(websocket): + reset_calls.append(websocket) + await websocket.close() + + return { + "asyncio": asyncio, + "json": json, + "websockets": SimpleNamespace( + ConnectionClosed=FakeConnectionClosed, + InvalidState=FakeInvalidState, + ), + "websocket_users": set(), + "args": SimpleNamespace( + save_offline_segments=False, save_offline_segments_dir="unused" + ), + "_safe_int": lambda value, default: ( + int(value) if value is not None else default + ), + "_pcm_duration_ms": lambda pcm, fs, ch, sampwidth: 10, + "async_vad": async_vad, + "async_asr": async_asr, + "async_asr_online": async_asr_online, + "run_blocking": None, + "save_offline_wav_segment_sync": None, + "SEM_WAV": None, + "ws_reset": ws_reset, + } + + +def _run_server( + mode, + frame_count=1, + chunk_interval=10, + chunk_size_value=(5, 10, 5), + send_end_ack=True, + vad_results=None, + asr_error=None, + online_error=None, + asr_error_once=None, + online_error_once=None, + online_cache_states=None, + vad_calls=None, +): + audio = b"\x01\x00" * 160 + config = { + "mode": mode, + "chunk_interval": chunk_interval, + "audio_fs": 16000, + "wav_name": "sample", + "is_speaking": True, + } + if chunk_size_value is not None: + config["chunk_size"] = list(chunk_size_value) + end_message = {"is_speaking": False} + if send_end_ack: + end_message["is_end"] = True + incoming = [ + json.dumps(config), + *([audio] * frame_count), + json.dumps(end_message), + ] + websocket = FakeWebSocket(incoming) + asr_calls = [] + online_calls = [] + namespace = _server_namespace( + asr_calls, + online_calls, + vad_results=vad_results, + asr_error=asr_error, + online_error=online_error, + asr_error_once=asr_error_once, + online_error_once=online_error_once, + online_cache_states=online_cache_states, + vad_calls=vad_calls, + ) + ws_serve = _load_async_function(SERVER_PATH, "ws_serve", namespace) + asyncio.run(ws_serve(websocket)) + sent = [json.loads(message) for message in websocket.sent] + return audio, asr_calls, online_calls, sent + + +def test_offline_end_control_flushes_buffer_before_acknowledgement(): + audio, asr_calls, online_calls, sent = _run_server("offline") + + assert asr_calls == [audio] + assert online_calls == [] + assert sent[-1] == { + "mode": "offline", + "wav_name": "sample", + "is_final": True, + "is_end": True, + } + + +def test_online_end_control_flushes_residual_frames_before_acknowledgement(): + audio, asr_calls, online_calls, sent = _run_server("online") + + assert asr_calls == [] + assert online_calls == [audio] + assert sent[-1] == { + "mode": "online", + "wav_name": "sample", + "is_final": True, + "is_end": True, + } + + +def test_2pass_end_control_flushes_online_and_offline_before_acknowledgement(): + audio, asr_calls, online_calls, sent = _run_server("2pass") + + assert online_calls == [audio] + assert asr_calls == [audio] + assert sent[-1] == { + "mode": "2pass", + "wav_name": "sample", + "is_final": True, + "is_end": True, + } + + +def test_online_finalizes_cache_after_exact_chunk_boundary(): + audio, asr_calls, online_calls, sent = _run_server( + "online", frame_count=2, chunk_interval=2 + ) + + assert asr_calls == [] + assert online_calls == [audio * 2, b""] + assert sent[-1]["is_end"] is True + + +def test_online_vad_endpoint_flushes_residual_audio_before_cache_reset(): + audio, asr_calls, online_calls, sent = _run_server("online", vad_results=[(0, 10)]) + + assert asr_calls == [] + assert online_calls == [audio] + assert sent[-1]["is_end"] is True + + +def test_online_vad_endpoint_flushes_exact_chunk_cache_before_reset(): + cache_states = [] + audio, asr_calls, online_calls, sent = _run_server( + "online", + chunk_interval=1, + vad_results=[(0, 10)], + online_cache_states=cache_states, + ) + + assert asr_calls == [] + assert online_calls == [audio, b""] + assert cache_states == [{}, {"partial": True}] + assert sent[-1]["is_end"] is True + + +def test_offline_end_uses_buffered_audio_when_vad_has_not_started(): + audio, asr_calls, online_calls, sent = _run_server( + "offline", vad_results=[(-1, -1)] + ) + + assert online_calls == [] + assert asr_calls == [audio] + assert sent[-1]["is_end"] is True + + +def test_offline_end_does_not_emit_empty_final_after_vad_completion(): + audio, asr_calls, online_calls, sent = _run_server("offline", vad_results=[(0, 10)]) + + assert online_calls == [] + assert asr_calls == [audio] + assert sent[-1]["is_end"] is True + + +def test_offline_end_flushes_tail_after_an_earlier_vad_completion(): + audio, asr_calls, online_calls, sent = _run_server( + "offline", frame_count=2, vad_results=[(0, 10), (-1, -1)] + ) + + assert online_calls == [] + assert asr_calls == [audio, audio] + assert sent[-1]["is_end"] is True + + +def test_offline_end_falls_back_when_vad_reports_only_an_endpoint(): + audio, asr_calls, online_calls, sent = _run_server( + "offline", vad_results=[(-1, 10)] + ) + + assert online_calls == [] + assert asr_calls == [audio] + assert sent[-1]["is_end"] is True + + +def test_offline_endpoint_only_audio_is_not_dropped_by_a_later_segment(): + audio, asr_calls, online_calls, sent = _run_server( + "offline", + frame_count=3, + vad_results=[(-1, 10), (10, -1), (-1, 30)], + ) + + assert online_calls == [] + assert asr_calls == [audio, audio * 2] + assert sent[-1]["is_end"] is True + + +def test_end_ack_reports_audio_discarded_without_chunk_size(): + _, asr_calls, online_calls, sent = _run_server("offline", chunk_size_value=None) + + assert asr_calls == [] + assert online_calls == [] + assert sent[-1]["is_end"] is True + assert sent[-1]["is_final"] is False + assert "chunk_size" in sent[-1]["error"] + + +def test_end_ack_reports_audio_discarded_for_invalid_vad_chunk_size(): + _, asr_calls, online_calls, sent = _run_server("offline", chunk_size_value=(5,)) + + assert asr_calls == [] + assert online_calls == [] + assert sent[-1]["is_end"] is True + assert sent[-1]["is_final"] is False + assert "VAD chunk_size" in sent[-1]["error"] + + +def test_end_ack_rejects_unsupported_mode(): + vad_calls = [] + _, asr_calls, online_calls, sent = _run_server("bogus", vad_calls=vad_calls) + + assert vad_calls == [] + assert asr_calls == [] + assert online_calls == [] + assert sent[-1]["is_end"] is True + assert sent[-1]["is_final"] is False + assert "unsupported mode" in sent[-1]["error"] + + +def test_invalid_state_resets_and_unregisters_connection(): + audio = b"\x01\x00" * 160 + incoming = [ + json.dumps( + { + "mode": "offline", + "chunk_interval": 10, + "audio_fs": 16000, + "wav_name": "sample", + "is_speaking": True, + "chunk_size": [5, 10, 5], + } + ), + audio, + json.dumps({"is_speaking": False, "is_end": True}), + ] + + class InvalidStateWebSocket(FakeWebSocket): + async def send(self, message): + raise FakeInvalidState + + websocket = InvalidStateWebSocket(incoming) + reset_calls = [] + namespace = _server_namespace([], [], reset_calls=reset_calls) + ws_serve = _load_async_function(SERVER_PATH, "ws_serve", namespace) + + asyncio.run(ws_serve(websocket)) + + assert reset_calls == [websocket] + assert websocket.closed is True + assert websocket not in namespace["websocket_users"] + + +def test_old_end_control_without_is_end_remains_supported(): + audio, asr_calls, online_calls, sent = _run_server("offline", send_end_ack=False) + + assert asr_calls == [audio] + assert online_calls == [] + assert not any(message.get("is_end") for message in sent) + + +def test_end_ack_reports_inference_failure(): + _, asr_calls, _, sent = _run_server( + "offline", asr_error=RuntimeError("offline inference failed") + ) + + assert len(asr_calls) == 1 + assert sent[-1]["is_end"] is True + assert sent[-1]["is_final"] is False + assert "offline inference failed" in sent[-1]["error"] + + +def test_end_ack_preserves_earlier_online_inference_failure(): + _, _, online_calls, sent = _run_server( + "online", + chunk_interval=1, + online_error_once=RuntimeError("streaming chunk failed"), + ) + + assert len(online_calls) == 2 + assert sent[-1]["is_end"] is True + assert sent[-1]["is_final"] is False + assert "streaming chunk failed" in sent[-1]["error"] + + +def test_end_ack_preserves_earlier_offline_inference_failure(): + _, asr_calls, _, sent = _run_server( + "offline", + vad_results=[(0, 10)], + asr_error_once=RuntimeError("offline segment failed"), + ) + + assert len(asr_calls) == 2 + assert sent[-1]["is_end"] is True + assert sent[-1]["is_final"] is False + assert "offline segment failed" in sent[-1]["error"] + + +def test_online_final_with_empty_audio_still_flushes_model_cache(): + calls = [] + + async def run_blocking(function, model, audio, status, sem): + calls.append((function, model, audio, status, sem)) + return [{"text": ""}] + + namespace = { + "json": json, + "run_blocking": run_blocking, + "_generate_sync": object(), + "model_asr_streaming": object(), + "SEM_ASR_ONLINE": object(), + } + async_asr_online = _load_async_function(SERVER_PATH, "async_asr_online", namespace) + websocket = FakeWebSocket() + websocket.mode = "online" + websocket.wav_name = "sample" + websocket.is_speaking = False + websocket.status_dict_asr_online = {"cache": {"partial": True}, "is_final": True} + + asyncio.run(async_asr_online(websocket, b"")) + + assert len(calls) == 1 + assert calls[0][2] == b"" + + +def test_offline_result_send_failure_propagates_to_the_caller(): + generate_sync = object() + speaker_sync = object() + + async def run_blocking(function, *args, sem): + if function is generate_sync: + return [{"text": "final text"}] + if function is speaker_sync: + return "unknown", 0.0 + raise AssertionError("unexpected blocking function") + + namespace = { + "json": json, + "run_blocking": run_blocking, + "_generate_sync": generate_sync, + "_sv_and_match_sync": speaker_sync, + "model_asr": object(), + "model_punc": None, + "SEM_ASR_OFFLINE": object(), + "SEM_SV": object(), + "args": SimpleNamespace(speaker_db_reload_sec=60), + "to_python": lambda value: value, + } + async_asr = _load_async_function(SERVER_PATH, "async_asr", namespace) + + class FailingResultWebSocket(FakeWebSocket): + async def send(self, message): + raise RuntimeError("result send failed") + + websocket = FailingResultWebSocket() + websocket.mode = "offline" + websocket.wav_name = "sample" + websocket.status_dict_asr = {} + websocket.status_dict_punc = {"cache": {}} + + with pytest.raises(RuntimeError, match="result send failed"): + asyncio.run(async_asr(websocket, b"\x01\x00")) + + +def _client_namespace(args, websocket): + return { + "asyncio": asyncio, + "json": json, + "os": __import__("os"), + "time": __import__("time"), + "args": args, + "websocket": websocket, + "voices": None, + "offline_msg_done": False, + "latency_first_audio_time": {}, + "latency_last_audio_time": {}, + "latency_first_text_printed": {}, + } + + +def test_file_sender_waits_for_protocol_ack_instead_of_fixed_sleep(tmp_path): + pcm_path = tmp_path / "sample.pcm" + pcm_path.write_bytes(b"\x01\x00" * 160) + websocket = FakeWebSocket() + args = SimpleNamespace( + audio_in=str(pcm_path), + hotword="", + audio_fs=16000, + use_itn=1, + chunk_size=[5, 10, 5], + chunk_interval=10, + encoder_chunk_look_back=4, + decoder_chunk_look_back=0, + mode="offline", + send_without_sleep=True, + result_timeout=1.0, + ) + namespace = _client_namespace(args, websocket) + record_from_scp = _load_async_function(CLIENT_PATH, "record_from_scp", namespace) + + async def exercise(): + completed = asyncio.Queue() + task = asyncio.create_task(record_from_scp(0, 1, completed)) + try: + for _ in range(100): + if any( + isinstance(message, str) and json.loads(message).get("is_end") + for message in websocket.sent + ): + break + await asyncio.sleep(0) + + assert not task.done() + completed.put_nowait({"is_end": True, "is_final": True}) + await asyncio.wait_for(task, timeout=0.2) + finally: + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + asyncio.run(exercise()) + + end_message = json.loads(websocket.sent[-1]) + assert end_message == {"is_speaking": False, "is_end": True} + assert websocket.closed is True + + +def test_file_sender_surfaces_server_error_ack(tmp_path): + pcm_path = tmp_path / "sample.pcm" + pcm_path.write_bytes(b"\x01\x00" * 160) + websocket = FakeWebSocket() + args = SimpleNamespace( + audio_in=str(pcm_path), + hotword="", + audio_fs=16000, + use_itn=1, + chunk_size=[5, 10, 5], + chunk_interval=10, + encoder_chunk_look_back=4, + decoder_chunk_look_back=0, + mode="offline", + send_without_sleep=True, + result_timeout=1.0, + ) + namespace = _client_namespace(args, websocket) + record_from_scp = _load_async_function(CLIENT_PATH, "record_from_scp", namespace) + + async def exercise(): + completed = asyncio.Queue() + task = asyncio.create_task(record_from_scp(0, 1, completed)) + try: + for _ in range(100): + if any( + isinstance(message, str) and json.loads(message).get("is_end") + for message in websocket.sent + ): + break + await asyncio.sleep(0) + + completed.put_nowait( + {"is_end": True, "is_final": False, "error": "offline inference failed"} + ) + with pytest.raises(RuntimeError, match="offline inference failed"): + await asyncio.wait_for(task, timeout=0.2) + finally: + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + asyncio.run(exercise()) + + +def test_file_sender_rejects_nonfinal_ack_without_error(tmp_path): + pcm_path = tmp_path / "sample.pcm" + pcm_path.write_bytes(b"\x01\x00" * 160) + websocket = FakeWebSocket() + args = SimpleNamespace( + audio_in=str(pcm_path), + hotword="", + audio_fs=16000, + use_itn=1, + chunk_size=[5, 10, 5], + chunk_interval=10, + encoder_chunk_look_back=4, + decoder_chunk_look_back=0, + mode="offline", + send_without_sleep=True, + result_timeout=1.0, + ) + namespace = _client_namespace(args, websocket) + record_from_scp = _load_async_function(CLIENT_PATH, "record_from_scp", namespace) + + async def exercise(): + completed = asyncio.Queue() + task = asyncio.create_task(record_from_scp(0, 1, completed)) + try: + for _ in range(100): + if any( + isinstance(message, str) and json.loads(message).get("is_end") + for message in websocket.sent + ): + break + await asyncio.sleep(0) + + completed.put_nowait({"is_end": True, "is_final": False}) + with pytest.raises(RuntimeError, match="did not finalize"): + await asyncio.wait_for(task, timeout=0.2) + finally: + if not task.done(): + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + asyncio.run(exercise()) + + +def test_file_sender_closes_connection_after_result_timeout(tmp_path): + pcm_path = tmp_path / "sample.pcm" + pcm_path.write_bytes(b"\x01\x00" * 160) + websocket = FakeWebSocket() + args = SimpleNamespace( + audio_in=str(pcm_path), + hotword="", + audio_fs=16000, + use_itn=1, + chunk_size=[5, 10, 5], + chunk_interval=10, + encoder_chunk_look_back=4, + decoder_chunk_look_back=0, + mode="offline", + send_without_sleep=True, + result_timeout=0.01, + ) + namespace = _client_namespace(args, websocket) + record_from_scp = _load_async_function(CLIENT_PATH, "record_from_scp", namespace) + + async def exercise(): + completed = asyncio.Queue() + with pytest.raises(TimeoutError, match="did not acknowledge"): + await record_from_scp(0, 1, completed) + + asyncio.run(exercise()) + + assert websocket.closed is True + + +def test_receiver_releases_file_sender_only_on_end_acknowledgement(monkeypatch): + class ReceivingWebSocket(FakeWebSocket): + async def recv(self): + try: + return next(self.incoming) + except StopIteration: + raise FakeConnectionClosed + + websocket = ReceivingWebSocket( + [ + json.dumps( + {"mode": "offline", "text": "done", "is_final": True, "is_end": True} + ) + ] + ) + args = SimpleNamespace(thread_num=1, output_dir=None, words_max_print=10000) + namespace = _client_namespace(args, websocket) + namespace.update( + { + "websockets": SimpleNamespace( + exceptions=SimpleNamespace(ConnectionClosedOK=FakeConnectionClosed) + ), + "MeetingWriter": object, + "_iso": lambda ts: str(ts), + } + ) + message = _load_async_function(CLIENT_PATH, "message", namespace) + monkeypatch.setitem(sys.modules, "websockets", namespace["websockets"]) + + async def exercise(): + completed = asyncio.Queue() + await message("0", None, completed) + acknowledgement = completed.get_nowait() + assert acknowledgement["is_end"] is True + + asyncio.run(exercise()) + + +def test_receiver_failure_releases_waiting_file_sender(monkeypatch): + class ReceivingWebSocket(FakeWebSocket): + async def recv(self): + try: + return next(self.incoming) + except StopIteration: + raise FakeConnectionClosed + + websocket = ReceivingWebSocket(["not-json"]) + args = SimpleNamespace(thread_num=1, output_dir=None, words_max_print=10000) + namespace = _client_namespace(args, websocket) + namespace.update( + { + "websockets": SimpleNamespace( + exceptions=SimpleNamespace(ConnectionClosedOK=FakeConnectionClosed) + ), + "MeetingWriter": object, + "_iso": lambda ts: str(ts), + } + ) + message = _load_async_function(CLIENT_PATH, "message", namespace) + monkeypatch.setitem(sys.modules, "websockets", namespace["websockets"]) + + async def exercise(): + completed = asyncio.Queue() + await message("0", None, completed) + failure = completed.get_nowait() + assert "receiver failed" in failure["error"] + + asyncio.run(exercise()) + + +def test_readme_documents_file_completion_acknowledgement(): + readme = README_PATH.read_text(encoding="utf-8") + + assert "--result_timeout" in readme + assert '"is_speaking": false, "is_end": true' in readme + assert '"is_end": true, "is_final": true' in readme + assert '"is_final": false, "error"' in readme