86639692ba
- Added support for public Realtime API, including new routes for managing API keys and handling WebRTC connections. - Introduced RealtimeApiKey model and associated CRUD operations for admin management of API keys. - Implemented authentication mechanisms for API keys and client secrets. - Enhanced environment configuration with new secrets for Realtime API. - Created OpenAIRealtime session management and event processing for real-time interactions. - Updated schemas and settings to accommodate new features and ensure compatibility with existing systems.
110 lines
3.1 KiB
Python
110 lines
3.1 KiB
Python
"""Shared ownership of peer connections and their pipeline tasks."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from collections.abc import Coroutine
|
|
from typing import Any
|
|
|
|
from loguru import logger
|
|
|
|
|
|
active_connections: set[object] = set()
|
|
pipeline_tasks: set[asyncio.Task[None]] = set()
|
|
connection_tasks: dict[object, asyncio.Task[None]] = {}
|
|
DEFAULT_CLOSE_GRACE_SECONDS = 10.0
|
|
|
|
|
|
def _consume_pipeline_result(task: asyncio.Task[None], connection: object) -> None:
|
|
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"Realtime pipeline 已取消: task={task.get_name()}")
|
|
return
|
|
if error is not None:
|
|
logger.opt(exception=error).error(
|
|
f"Realtime pipeline 异常结束: task={task.get_name()}"
|
|
)
|
|
|
|
|
|
def start_pipeline_task(
|
|
connection: object,
|
|
coroutine: Coroutine[Any, Any, None],
|
|
*,
|
|
protocol: str,
|
|
) -> asyncio.Task[None]:
|
|
connection_id = str(getattr(connection, "pc_id", id(connection)))
|
|
task = asyncio.create_task(
|
|
coroutine,
|
|
name=f"{protocol}-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,
|
|
timeout: float = DEFAULT_CLOSE_GRACE_SECONDS,
|
|
) -> None:
|
|
if task is None:
|
|
return
|
|
try:
|
|
await asyncio.wait_for(asyncio.shield(task), timeout=timeout)
|
|
return
|
|
except TimeoutError:
|
|
logger.warning(
|
|
f"Realtime pipeline 关闭超过 {timeout:g} 秒，执行取消: "
|
|
f"connection_id={connection_id}"
|
|
)
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception:
|
|
return
|
|
|
|
task.cancel()
|
|
done, _pending = await asyncio.wait({task}, timeout=timeout)
|
|
if not done:
|
|
logger.error(f"Realtime pipeline 取消后仍未退出: connection_id={connection_id}")
|
|
|
|
|
|
async def shutdown_active_sessions(
|
|
*,
|
|
timeout: float = DEFAULT_CLOSE_GRACE_SECONDS,
|
|
) -> None:
|
|
connections = list(active_connections)
|
|
if connections:
|
|
await asyncio.gather(
|
|
*(connection.disconnect() for connection in connections),
|
|
return_exceptions=True,
|
|
)
|
|
active_connections.difference_update(connections)
|
|
|
|
tasks = list(pipeline_tasks)
|
|
if not tasks:
|
|
return
|
|
done, pending = await asyncio.wait(tasks, timeout=timeout)
|
|
for task in pending:
|
|
task.cancel()
|
|
cancelled: set[asyncio.Task[None]] = set()
|
|
if pending:
|
|
cancelled, stuck = await asyncio.wait(pending, timeout=timeout)
|
|
if stuck:
|
|
logger.error(f"应用关闭时仍有 {len(stuck)} 个 Realtime pipeline 未退出")
|
|
logger.info(
|
|
f"Realtime 会话清理完成: normal={len(done)} cancelled={len(cancelled)}"
|
|
)
|
|
|