From 16156028c9444adef126d24e9f6d9be99cde1ad0 Mon Sep 17 00:00:00 2001 From: Xin Wang Date: Wed, 5 Aug 2026 13:48:29 +0800 Subject: [PATCH] fix: manage WebRTC pipeline shutdown --- backend/app.py | 5 +- backend/routes/voice_webrtc.py | 135 ++++++++++++++++++- backend/services/conversation_history.py | 15 +++ backend/services/pipecat/pipeline.py | 17 ++- backend/tests/test_conversation_history.py | 25 ++++ backend/tests/test_voice_webrtc_lifecycle.py | 65 +++++++++ 6 files changed, 253 insertions(+), 9 deletions(-) create mode 100644 backend/tests/test_voice_webrtc_lifecycle.py diff --git a/backend/app.py b/backend/app.py index 9b41798..35f3041 100644 --- a/backend/app.py +++ b/backend/app.py @@ -41,7 +41,10 @@ async def lifespan(_app: FastAPI): await sync_interface_definitions() await sync_default_tools() await recover_interrupted_documents() - yield + try: + yield + finally: + await voice_webrtc.shutdown_active_sessions() app = FastAPI(title="AI Video Assistant 平台 - 后端", lifespan=lifespan) diff --git a/backend/routes/voice_webrtc.py b/backend/routes/voice_webrtc.py index 092a447..729353c 100644 --- a/backend/routes/voice_webrtc.py +++ b/backend/routes/voice_webrtc.py @@ -11,6 +11,8 @@ import asyncio import base64 import json +from collections.abc import Coroutine +from typing import Any from db.session import SessionLocal from fastapi import APIRouter, Body, Depends, Request, WebSocket @@ -27,6 +29,116 @@ from services.webrtc_ice import aiortc_ice_servers, client_ice_servers router = APIRouter(tags=["voice"]) _http_peers: dict[str, object] = {} +_active_connections: set[object] = set() +_pipeline_tasks: set[asyncio.Task[None]] = set() +_connection_tasks: dict[object, asyncio.Task[None]] = {} +PIPELINE_CLOSE_GRACE_SECONDS = 10.0 + + +def _consume_pipeline_result(task: asyncio.Task[None], connection: object) -> None: + """Retrieve background exceptions and release the task's strong reference.""" + _pipeline_tasks.discard(task) + if _connection_tasks.get(connection) is task: + _connection_tasks.pop(connection, None) + try: + error = task.exception() + except asyncio.CancelledError: + logger.info(f"WebRTC pipeline 已取消: task={task.get_name()}") + return + if error is not None: + logger.opt(exception=error).error( + f"WebRTC pipeline 异常结束: task={task.get_name()}" + ) + + +def _start_pipeline_task( + connection: object, + coroutine: Coroutine[Any, Any, None], +) -> asyncio.Task[None]: + """Start one strongly referenced pipeline task for a WebRTC connection.""" + connection_id = str(getattr(connection, "pc_id", "unknown")) + task = asyncio.create_task( + coroutine, + name=f"webrtc-pipeline:{connection_id}", + ) + _active_connections.add(connection) + _pipeline_tasks.add(task) + _connection_tasks[connection] = task + task.add_done_callback( + lambda completed, connection=connection: _consume_pipeline_result( + completed, + connection, + ) + ) + return task + + +async def _wait_for_pipeline_close( + task: asyncio.Task[None] | None, + *, + connection_id: str, +) -> None: + """Let transport disconnect finish normally, then cancel a stuck pipeline.""" + if task is None: + return + try: + await asyncio.wait_for( + asyncio.shield(task), + timeout=PIPELINE_CLOSE_GRACE_SECONDS, + ) + return + except TimeoutError: + logger.warning( + f"WebRTC pipeline 关闭超过 {PIPELINE_CLOSE_GRACE_SECONDS:g} 秒，" + f"执行取消: pc_id={connection_id}" + ) + except asyncio.CancelledError: + raise + except Exception: + # The done callback owns exception reporting and retrieval. + return + + task.cancel() + done, _pending = await asyncio.wait( + {task}, + timeout=PIPELINE_CLOSE_GRACE_SECONDS, + ) + if not done: + logger.error(f"WebRTC pipeline 取消后仍未退出: pc_id={connection_id}") + + +async def shutdown_active_sessions() -> None: + """Close active peers and drain every managed pipeline during app shutdown.""" + connections = list(_active_connections) + if connections: + await asyncio.gather( + *(connection.disconnect() for connection in connections), + return_exceptions=True, + ) + _active_connections.difference_update(connections) + _http_peers.clear() + + tasks = list(_pipeline_tasks) + if not tasks: + return + done, pending = await asyncio.wait( + tasks, + timeout=PIPELINE_CLOSE_GRACE_SECONDS, + ) + for task in pending: + task.cancel() + if pending: + cancelled, stuck = await asyncio.wait( + pending, + timeout=PIPELINE_CLOSE_GRACE_SECONDS, + ) + if stuck: + logger.error(f"应用关闭时仍有 {len(stuck)} 个 WebRTC pipeline 未退出") + else: + cancelled = set() + logger.info( + f"WebRTC 会话清理完成: normal={len(done)} cancelled={len(cancelled)}" + ) @router.get("/api/webrtc/ice-servers", dependencies=[Depends(require_admin)]) @@ -155,7 +267,12 @@ async def _handle_offer_payload(payload, peers): 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] @@ -187,25 +304,31 @@ async def _handle_offer_payload(payload, peers): await pc.initialize(sdp=offer.sdp, type=offer.type) peers[pc.pc_id] = pc - @pc.event_handler("closed") - async def _on_closed(conn: SmallWebRTCConnection): - peers.pop(conn.pc_id, None) - # 后台跑管线:WebRTC transport + 解析出的运行时配置 transport = build_webrtc_transport( pc, video_in_enabled=vision_enabled, ) - asyncio.create_task( + pipeline_task = _start_pipeline_task( + pc, run_pipeline( transport, cfg, vision_enabled=vision_enabled, assistant_id=offer.assistant_id, channel="webrtc", - ) + ), ) + @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"], diff --git a/backend/services/conversation_history.py b/backend/services/conversation_history.py index 96fb755..5787070 100644 --- a/backend/services/conversation_history.py +++ b/backend/services/conversation_history.py @@ -169,6 +169,21 @@ class ConversationRecorder: logger.error(f"保存对话文本失败，不影响本次通话: {exc}") async def finish(self, *, status: str = "completed") -> None: + """Finish the session even if the owning pipeline is being cancelled.""" + finish_task = asyncio.create_task( + self._finish(status=status), + name=f"conversation-recorder-finish:{self.session_id}", + ) + try: + await asyncio.shield(finish_task) + except asyncio.CancelledError: + # Shield prevents an interrupted pipeline cleanup from abandoning a + # checked-out asyncpg connection. Preserve cancellation only after + # the short database cleanup has returned the connection. + await asyncio.shield(finish_task) + raise + + async def _finish(self, *, status: str) -> None: async with self._lock: try: async with SessionLocal() as db: diff --git a/backend/services/pipecat/pipeline.py b/backend/services/pipecat/pipeline.py index 5597043..db60e86 100644 --- a/backend/services/pipecat/pipeline.py +++ b/backend/services/pipecat/pipeline.py @@ -810,14 +810,22 @@ async def run_pipeline( ) runner = WorkerRunner(handle_sigint=False) run_status = "completed" + can_finalize_recorder = True try: await runner.add_workers(worker) await runner.run() + except GeneratorExit: + # A coroutine being synchronously closed cannot perform async cleanup. + # The WebRTC route now prevents this by retaining every pipeline task; + # keep this guard so a future unmanaged caller cannot corrupt asyncpg. + run_status = "failed" + can_finalize_recorder = False + raise except Exception: run_status = "failed" raise finally: - if recorder: + if recorder and can_finalize_recorder: await recorder.finish(status=run_status) logger.info("管线已结束") @@ -971,13 +979,18 @@ async def run_realtime_pipeline( ) runner = WorkerRunner(handle_sigint=False) run_status = "completed" + can_finalize_recorder = True try: await runner.add_workers(worker) await runner.run() + except GeneratorExit: + run_status = "failed" + can_finalize_recorder = False + raise except Exception: run_status = "failed" raise finally: - if recorder: + if recorder and can_finalize_recorder: await recorder.finish(status=run_status) logger.info("Realtime 管线已结束") diff --git a/backend/tests/test_conversation_history.py b/backend/tests/test_conversation_history.py index ada7ccd..a60dcde 100644 --- a/backend/tests/test_conversation_history.py +++ b/backend/tests/test_conversation_history.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import unittest from types import SimpleNamespace from unittest.mock import AsyncMock, patch @@ -8,6 +9,30 @@ from services.conversation_history import ConversationRecorder class ConversationRecorderTest(unittest.IsolatedAsyncioTestCase): + async def test_finish_waits_for_database_cleanup_when_cancelled(self): + recorder = ConversationRecorder("conv_test") + started = asyncio.Event() + release = asyncio.Event() + completed = asyncio.Event() + + async def finish_database_write(*, status: str): + self.assertEqual(status, "failed") + started.set() + await release.wait() + completed.set() + + recorder._finish = finish_database_write + task = asyncio.create_task(recorder.finish(status="failed")) + await started.wait() + task.cancel() + await asyncio.sleep(0) + + self.assertFalse(task.done()) + release.set() + with self.assertRaises(asyncio.CancelledError): + await task + self.assertTrue(completed.is_set()) + async def test_fixed_reply_transcript_keeps_workflow_metadata(self): recorder = ConversationRecorder("conv_test") recorder._append = AsyncMock() diff --git a/backend/tests/test_voice_webrtc_lifecycle.py b/backend/tests/test_voice_webrtc_lifecycle.py new file mode 100644 index 0000000..ca76a67 --- /dev/null +++ b/backend/tests/test_voice_webrtc_lifecycle.py @@ -0,0 +1,65 @@ +from __future__ import annotations + +import asyncio +import unittest +from unittest.mock import patch + +from routes import voice_webrtc + + +class FakeConnection: + pc_id = "pc_test" + + +class WebRTCPipelineLifecycleTest(unittest.IsolatedAsyncioTestCase): + async def asyncTearDown(self): + tasks = list(voice_webrtc._pipeline_tasks) + for task in tasks: + task.cancel() + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + voice_webrtc._pipeline_tasks.clear() + voice_webrtc._connection_tasks.clear() + voice_webrtc._active_connections.clear() + + async def test_pipeline_task_is_retained_until_completion(self): + connection = FakeConnection() + release = asyncio.Event() + + async def pipeline(): + await release.wait() + + task = voice_webrtc._start_pipeline_task(connection, pipeline()) + self.assertIn(task, voice_webrtc._pipeline_tasks) + self.assertIs(voice_webrtc._connection_tasks[connection], task) + + release.set() + await task + await asyncio.sleep(0) + + self.assertNotIn(task, voice_webrtc._pipeline_tasks) + self.assertNotIn(connection, voice_webrtc._connection_tasks) + + async def test_stuck_pipeline_is_cancelled_after_close_grace_period(self): + connection = FakeConnection() + cancelled = asyncio.Event() + + async def pipeline(): + try: + await asyncio.Event().wait() + finally: + cancelled.set() + + task = voice_webrtc._start_pipeline_task(connection, pipeline()) + with patch.object(voice_webrtc, "PIPELINE_CLOSE_GRACE_SECONDS", 0.01): + await voice_webrtc._wait_for_pipeline_close( + task, + connection_id=connection.pc_id, + ) + + self.assertTrue(cancelled.is_set()) + self.assertTrue(task.cancelled()) + + +if __name__ == "__main__": + unittest.main()