fix: manage WebRTC pipeline shutdown

This commit is contained in:
Xin Wang
2026-08-05 13:48:29 +08:00
parent 10bea48d72
commit 16156028c9
6 changed files with 253 additions and 9 deletions
+3
View File
@@ -41,7 +41,10 @@ async def lifespan(_app: FastAPI):
await sync_interface_definitions() await sync_interface_definitions()
await sync_default_tools() await sync_default_tools()
await recover_interrupted_documents() await recover_interrupted_documents()
try:
yield yield
finally:
await voice_webrtc.shutdown_active_sessions()
app = FastAPI(title="AI Video Assistant 平台 - 后端", lifespan=lifespan) app = FastAPI(title="AI Video Assistant 平台 - 后端", lifespan=lifespan)
+128 -5
View File
@@ -11,6 +11,8 @@
import asyncio import asyncio
import base64 import base64
import json import json
from collections.abc import Coroutine
from typing import Any
from db.session import SessionLocal from db.session import SessionLocal
from fastapi import APIRouter, Body, Depends, Request, WebSocket 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"]) router = APIRouter(tags=["voice"])
_http_peers: dict[str, object] = {} _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)]) @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")) restart_pc = bool(payload.get("restart_pc"))
if restart_pc and pc_id and pc_id in peers: if restart_pc and pc_id and pc_id in peers:
old_pc = peers.pop(pc_id) old_pc = peers.pop(pc_id)
old_task = _connection_tasks.get(old_pc)
await old_pc.disconnect() 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: if pc_id and pc_id in peers:
pc = peers[pc_id] pc = peers[pc_id]
@@ -187,23 +304,29 @@ async def _handle_offer_payload(payload, peers):
await pc.initialize(sdp=offer.sdp, type=offer.type) await pc.initialize(sdp=offer.sdp, type=offer.type)
peers[pc.pc_id] = pc peers[pc.pc_id] = pc
@pc.event_handler("closed")
async def _on_closed(conn: SmallWebRTCConnection):
peers.pop(conn.pc_id, None)
# 后台跑管线:WebRTC transport + 解析出的运行时配置 # 后台跑管线:WebRTC transport + 解析出的运行时配置
transport = build_webrtc_transport( transport = build_webrtc_transport(
pc, pc,
video_in_enabled=vision_enabled, video_in_enabled=vision_enabled,
) )
asyncio.create_task( pipeline_task = _start_pipeline_task(
pc,
run_pipeline( run_pipeline(
transport, transport,
cfg, cfg,
vision_enabled=vision_enabled, vision_enabled=vision_enabled,
assistant_id=offer.assistant_id, assistant_id=offer.assistant_id,
channel="webrtc", 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() answer = pc.get_answer()
+15
View File
@@ -169,6 +169,21 @@ class ConversationRecorder:
logger.error(f"保存对话文本失败,不影响本次通话: {exc}") logger.error(f"保存对话文本失败,不影响本次通话: {exc}")
async def finish(self, *, status: str = "completed") -> None: 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: async with self._lock:
try: try:
async with SessionLocal() as db: async with SessionLocal() as db:
+15 -2
View File
@@ -810,14 +810,22 @@ async def run_pipeline(
) )
runner = WorkerRunner(handle_sigint=False) runner = WorkerRunner(handle_sigint=False)
run_status = "completed" run_status = "completed"
can_finalize_recorder = True
try: try:
await runner.add_workers(worker) await runner.add_workers(worker)
await runner.run() 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: except Exception:
run_status = "failed" run_status = "failed"
raise raise
finally: finally:
if recorder: if recorder and can_finalize_recorder:
await recorder.finish(status=run_status) await recorder.finish(status=run_status)
logger.info("管线已结束") logger.info("管线已结束")
@@ -971,13 +979,18 @@ async def run_realtime_pipeline(
) )
runner = WorkerRunner(handle_sigint=False) runner = WorkerRunner(handle_sigint=False)
run_status = "completed" run_status = "completed"
can_finalize_recorder = True
try: try:
await runner.add_workers(worker) await runner.add_workers(worker)
await runner.run() await runner.run()
except GeneratorExit:
run_status = "failed"
can_finalize_recorder = False
raise
except Exception: except Exception:
run_status = "failed" run_status = "failed"
raise raise
finally: finally:
if recorder: if recorder and can_finalize_recorder:
await recorder.finish(status=run_status) await recorder.finish(status=run_status)
logger.info("Realtime 管线已结束") logger.info("Realtime 管线已结束")
@@ -1,5 +1,6 @@
from __future__ import annotations from __future__ import annotations
import asyncio
import unittest import unittest
from types import SimpleNamespace from types import SimpleNamespace
from unittest.mock import AsyncMock, patch from unittest.mock import AsyncMock, patch
@@ -8,6 +9,30 @@ from services.conversation_history import ConversationRecorder
class ConversationRecorderTest(unittest.IsolatedAsyncioTestCase): 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): async def test_fixed_reply_transcript_keeps_workflow_metadata(self):
recorder = ConversationRecorder("conv_test") recorder = ConversationRecorder("conv_test")
recorder._append = AsyncMock() recorder._append = AsyncMock()
@@ -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()