MoviePilot/app/agent/llm/session.py

"""
LLM Provider 临时授权会话状态的唯一 owner。
"""

from __future__ import annotations

import time
from dataclasses import dataclass, field
from typing import Any, Dict, Optional

from app.agent.llm.runtime import LLMProviderAuthError


@dataclass
class PendingAuthSession:
    """
    保存临时鉴权会话，避免把 PKCE/device code 等状态写回配置。

    :param session_id: 临时授权会话的随机标识
    :param provider_id: 需要授权的提供商标识
    :param method_id: 本次采用的授权方式标识
    :param flow_type: 浏览器回调或设备码等授权流程类型
    :param status: 当前授权状态，初始为 pending
    :param message: 面向用户的进度或失败说明
    :param authorize_url: 浏览器授权入口；不适用时为 None
    :param instructions: 完成授权所需的操作提示；无提示时为 None
    :param verification_url: 设备码授权使用的验证页面；不适用时为 None
    :param user_code: 需要用户输入的设备验证码；不适用时为 None
    :param interval_seconds: 设备授权状态的轮询间隔，单位为秒，默认 5
    :param expires_at: 授权到期的 Unix 时间戳；0 表示按创建时间加默认时限计算
    :param created_at: 授权会话创建的 Unix 时间戳，单位为秒
    :param context: 仅在内存中保存的 PKCE、设备码和回调关联状态，不写入配置
    """

    session_id: str
    provider_id: str
    method_id: str
    flow_type: str
    status: str = "pending"
    message: str = ""
    authorize_url: Optional[str] = None
    instructions: Optional[str] = None
    verification_url: Optional[str] = None
    user_code: Optional[str] = None
    interval_seconds: int = 5
    expires_at: float = 0
    created_at: float = field(default_factory=time.time)
    context: Dict[str, Any] = field(default_factory=dict)


class _ProviderSession:
    """
    LLM Provider 临时授权会话状态的唯一 owner。
    """

    _lock: Any
    _pending_sessions: dict[str, PendingAuthSession]
    _oauth_state_index: dict[str, str]

    def __getattr__(self, name: str) -> Any:
        """
        将跨 owner 调用交给最终 Facade 的 MRO 解析。
        """
        raise AttributeError(name)

    _AUTH_SESSION_DONE_RETENTION = 300

    def _cleanup_auth_sessions_locked(self, now: Optional[float] = None) -> None:
        """
        清理过期或已完成一段时间的临时授权会话。

        调用方必须已经持有 `_lock`，这样 `_pending_sessions` 与
        `_oauth_state_index` 能保持一致，避免 state 残留。
        """
        now = time.time() if now is None else now
        expired_session_ids = []
        for session_id, session in self._pending_sessions.items():
            expires_at = session.expires_at or session.created_at + 600
            if session.status == "pending":
                if expires_at <= now:
                    expired_session_ids.append(session_id)
            elif expires_at + self._AUTH_SESSION_DONE_RETENTION <= now:
                expired_session_ids.append(session_id)

        if not expired_session_ids:
            return

        expired_session_ids_set = set(expired_session_ids)
        for session_id in expired_session_ids:
            self._pending_sessions.pop(session_id, None)
        for state, session_id in list(self._oauth_state_index.items()):
            if session_id in expired_session_ids_set:
                self._oauth_state_index.pop(state, None)

    def get_session_status(self, session_id: str) -> dict[str, Any]:
        """
        读取临时授权会话状态。
        """
        with self._lock:
            self._cleanup_auth_sessions_locked()
            session = self._pending_sessions.get(session_id)
            if not session:
                raise LLMProviderAuthError("授权会话不存在或已过期")
            return {
                "session_id": session.session_id,
                "provider_id": session.provider_id,
                "status": session.status,
                "message": session.message,
                "user_code": session.user_code,
                "verification_url": session.verification_url,
                "authorize_url": session.authorize_url,
                "instructions": session.instructions,
                "interval_seconds": session.interval_seconds,
                "expires_at": session.expires_at,
            }

    async def _mark_session_success(self, session: PendingAuthSession, auth_data: dict[str, Any]) -> None:
        """
        标记授权会话为成功，并保存认证信息。
        """
        auth_data["updated_at"] = int(time.time())
        await self.save_auth(session.provider_id, auth_data)
        session.status = "authorized"
        session.message = "授权成功"

    @staticmethod
    def _mark_session_error(session: PendingAuthSession, message: str) -> None:
        """
        标记授权会话为失败，并记录错误信息。
        """
        session.status = "failed"
        session.message = message

    async def poll_auth_session(self, session_id: str) -> dict[str, Any]:
        """
        执行一次 device code 轮询，并返回最新状态。

        前端可按 interval_seconds 轮询，直到状态变为 authorized / failed。
        """
        with self._lock:
            self._cleanup_auth_sessions_locked()
            session = self._pending_sessions.get(session_id)
        if not session:
            raise LLMProviderAuthError("授权会话不存在或已过期")
        if session.status != "pending":
            return self.get_session_status(session_id)

        try:
            if session.provider_id == "chatgpt" and session.method_id == "device_code":
                await self._poll_chatgpt_device_auth(session)
            elif session.provider_id == "github-copilot" and session.method_id == "device_code":
                await self._poll_copilot_device_auth(session)
            else:
                raise LLMProviderAuthError("当前授权会话不支持轮询")
        except Exception as err:
            self._mark_session_error(session, str(err))
        return self.get_session_status(session_id)