src/api/session_pool.py

"""统一业务会话与请求池。

业务端点通过本模块创建一次下游会话；会话内部为每个需要转发到
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()