"""统一业务会话与请求池。 业务端点通过本模块创建一次下游会话；会话内部为每个需要转发到 Vcore AI Studio 匿名上游的请求创建滚动并行请求池。管理端点不走本模块。 """ import time import uuid import asyncio from contextlib import asynccontextmanager from dataclasses import dataclass from enum import Enum from typing import Any, AsyncGenerator, Coroutine from src.core.config import load_config from src.utils.logger import get_logger from .vcore_client import VcoreAIClient logger = get_logger(__name__) class SessionState(str, Enum): WAITING = "waiting" ACTIVE = "active" CLOSED = "closed" @dataclass class _SessionPoolState: active_count: int = 0 waiting_count: int = 0 class BusinessSessionPool: """进程内业务会话池状态管理。""" def __init__(self) -> None: self._condition = asyncio.Condition() self._state = _SessionPoolState() def _limit_from_config(self) -> int: cfg = load_config() limit = int(cfg.get("business_session_concurrency_limit", 0) or 0) return max(0, limit) async def activate(self, session: "GatewaySession") -> None: async with self._condition: self._state.waiting_count += 1 session.state = SessionState.WAITING logger.debug( f"业务会话进入等待: {session.session_id}, " f"active={self._state.active_count}, waiting={self._state.waiting_count}" ) try: while True: limit = self._limit_from_config() if limit <= 0 or self._state.active_count < limit: self._state.active_count += 1 self._state.waiting_count -= 1 session.state = SessionState.ACTIVE logger.debug( f"业务会话激活: {session.session_id}, " f"active={self._state.active_count}, waiting={self._state.waiting_count}, limit={'不限' if limit <= 0 else limit}" ) return await self._condition.wait() except BaseException: self._state.waiting_count = max(0, self._state.waiting_count - 1) session.state = SessionState.CLOSED self._condition.notify(1) raise async def close(self, session: "GatewaySession") -> None: async with self._condition: if session.state == SessionState.ACTIVE: self._state.active_count = max(0, self._state.active_count - 1) elif session.state == SessionState.WAITING: self._state.waiting_count = max(0, self._state.waiting_count - 1) session.state = SessionState.CLOSED self._condition.notify(1) business_session_pool = BusinessSessionPool() class GatewaySession: """单个下游请求会话。""" def __init__(self, vcore_client: VcoreAIClient, session_type: str = "business", source: str = "unknown") -> None: self.vcore_client = vcore_client self.session_type = session_type self.source = source self.session_id = f"{session_type}-{uuid.uuid4().hex[:12]}" self.started_at = time.perf_counter() self.state = SessionState.WAITING self._tasks: set[asyncio.Task[Any]] = set() self._closed = False logger.debug(f"创建{self.session_type}会话: {self.session_id}, source={self.source}") def create_task(self, coro: Coroutine[Any, Any, Any], name: str | None = None) -> asyncio.Task[Any]: """创建受当前会话托管的后台任务；会话关闭时会立即取消。""" if self._closed: raise asyncio.CancelledError(f"会话已关闭: {self.session_id}") task = asyncio.create_task(coro, name=name) self.track_task(task) return task def track_task(self, task: asyncio.Task[Any]) -> asyncio.Task[Any]: """把外部已创建任务纳入当前会话生命周期。""" if self._closed: task.cancel() return task self._tasks.add(task) task.add_done_callback(self._tasks.discard) return task def untrack_task(self, task: asyncio.Task[Any]) -> None: """解除任务与当前会话的等待关系；任务自身仍会继续响应已发出的取消。""" self._tasks.discard(task) @staticmethod def _consume_detached_task_result(task: asyncio.Task[Any]) -> None: try: task.result() except asyncio.CancelledError: pass except Exception as e: logger.debug(f"后台清理任务结束时出现异常: {e}") async def cancel_running_resources(self) -> None: """取消并回收当前会话占用的所有运行中任务。""" self._closed = True tasks = [task for task in self._tasks if not task.done()] if not tasks: logger.debug(f"取消{self.session_type}会话运行资源: {self.session_id}, 无运行中任务") return logger.info(f"取消{self.session_type}会话运行资源: {self.session_id}, tasks={len(tasks)}") for task in tasks: task.cancel() done, pending = await asyncio.wait(tasks, timeout=1.0) if done: await asyncio.gather(*done, return_exceptions=True) self._tasks.difference_update(done) if pending: for task in pending: self.untrack_task(task) task.add_done_callback(self._consume_detached_task_result) logger.warning( f"取消{self.session_type}会话运行资源超时，已转后台清理: " f"{self.session_id}, pending={len(pending)}" ) else: logger.info(f"取消{self.session_type}会话运行资源完成: {self.session_id}") async def close(self) -> None: await self.cancel_running_resources() await business_session_pool.close(self) elapsed_ms = (time.perf_counter() - self.started_at) * 1000 logger.debug(f"清理{self.session_type}会话: {self.session_id}, state={self.state}, elapsed={elapsed_ms:.0f}ms") async def stream_gemini(self, model: str, gemini_payload: dict[str, Any], **kwargs: Any) -> AsyncGenerator[dict[str, Any], None]: """通过请求池流式返回 Gemini chunk。""" async for chunk in self.vcore_client.stream_chat_realtime( model=model, gemini_payload=gemini_payload, business_session_id=self.session_id, gateway_session=self, **kwargs, ): yield chunk async def complete_gemini(self, model: str, gemini_payload: dict[str, Any], **kwargs: Any) -> dict[str, Any]: """通过请求池获取完整 Gemini 响应。""" return await self.vcore_client.complete_chat( model=model, gemini_payload=gemini_payload, business_session_id=self.session_id, gateway_session=self, **kwargs, ) async def count_tokens(self, model: str, contents: list[dict[str, Any]]) -> int: """业务会话内执行 countTokens。""" return await self.vcore_client.count_tokens( model=model, contents=contents, business_session_id=self.session_id, gateway_session=self, ) @asynccontextmanager async def business_session(vcore_client: VcoreAIClient, source: str = "unknown") -> AsyncGenerator[GatewaySession, None]: """创建、纳入会话池并自动清理业务会话。""" session = GatewaySession(vcore_client, session_type="business", source=source) try: await business_session_pool.activate(session) yield session finally: await session.close() @asynccontextmanager async def management_session(vcore_client: VcoreAIClient, source: str = "admin") -> AsyncGenerator[GatewaySession, None]: """创建并自动清理管理会话。当前管理端点多数本地处理，保留统一抽象。""" session = GatewaySession(vcore_client, session_type="management", source=source) try: yield session finally: await session.close()