session_persistence.py

# session_persistence.py - 会话持久化工具
import json
import pickle
from pathlib import Path
from datetime import datetime
from typing import Dict, Any

class SessionPersistence:
    """
    会话持久化管理器

    功能：
    - 导出内存中的所有会话到文件
    - 从文件恢复会话到内存
    - 支持定期自动保存
    """

    def __init__(self, save_dir='session_backups'):
        self.save_dir = Path(save_dir)
        self.save_dir.mkdir(exist_ok=True)

    def export_sessions(self, sessions: Dict, filename: str = None) -> str:
        """
        导出所有会话到pickle文件

        Args:
            sessions: sessions字典
            filename: 文件名（默认自动生成）

        Returns:
            str: 保存的文件路径
        """
        if filename is None:
            timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
            filename = f'sessions_backup_{timestamp}.pkl'

        filepath = self.save_dir / filename

        # 使用pickle保存完整对象
        with open(filepath, 'wb') as f:
            pickle.dump(sessions, f)

        print(f"[OK] 已导出 {len(sessions)} 个会话到: {filepath}")
        print(f"   文件大小: {filepath.stat().st_size / 1024:.2f} KB")

        # 同时生成一个可读的JSON摘要
        summary_file = self.save_dir / f'{filename}.summary.json'
        summary = self._generate_summary(sessions)
        with open(summary_file, 'w', encoding='utf-8') as f:
            json.dump(summary, f, indent=2, ensure_ascii=False)

        print(f"   摘要文件: {summary_file}")

        return str(filepath)

    def import_sessions(self, filename: str = None) -> Dict:
        """
        从文件导入会话

        Args:
            filename: 文件名（默认使用最新的）

        Returns:
            Dict: 恢复的sessions字典
        """
        if filename is None:
            # 找到最新的备份文件
            backup_files = sorted(self.save_dir.glob('sessions_backup_*.pkl'))
            if not backup_files:
                print("[!] 没有找到备份文件")
                return {}
            filepath = backup_files[-1]
        else:
            filepath = self.save_dir / filename

        if not filepath.exists():
            print(f"[!] 文件不存在: {filepath}")
            return {}

        # 加载pickle文件
        with open(filepath, 'rb') as f:
            sessions = pickle.load(f)

        print(f"[OK] 已导入 {len(sessions)} 个会话从: {filepath}")

        # 显示摘要
        active_count = sum(1 for s in sessions.values() if hasattr(s, 'investigator') and s.investigator)
        print(f"   其中有角色的会话: {active_count}")

        return sessions

    def _generate_summary(self, sessions: Dict) -> Dict[str, Any]:
        """生成会话摘要信息"""
        summary = {
            'export_time': datetime.now().isoformat(),
            'total_sessions': len(sessions),
            'sessions_with_characters': 0,
            'sessions': []
        }

        for session_id, session_state in sessions.items():
            session_info = {
                'session_id': session_id,
                'has_investigator': False,
                'character_name': None,
                'character_occupation': None,
                'message_count': len(getattr(session_state, 'messages', []))
            }

            # 提取调查员信息
            if hasattr(session_state, 'investigator') and session_state.investigator:
                inv = session_state.investigator
                session_info['has_investigator'] = True
                session_info['character_name'] = getattr(inv, 'name', None)
                session_info['character_occupation'] = getattr(inv, 'occupation', None)
                summary['sessions_with_characters'] += 1

            summary['sessions'].append(session_info)

        return summary

    def list_backups(self):
        """列出所有备份文件"""
        backup_files = sorted(self.save_dir.glob('sessions_backup_*.pkl'))

        if not backup_files:
            print("没有找到备份文件")
            return

        print("\n可用的备份文件:")
        print("=" * 80)

        for i, filepath in enumerate(backup_files, 1):
            size_kb = filepath.stat().st_size / 1024

            # 从文件名提取时间
            filename = filepath.stem
            timestamp_str = filename.replace('sessions_backup_', '')
            try:
                dt = datetime.strptime(timestamp_str, '%Y%m%d_%H%M%S')
                time_str = dt.strftime('%Y-%m-%d %H:%M:%S')
            except:
                time_str = timestamp_str

            print(f"{i}. {filepath.name}")
            print(f"   时间: {time_str}")
            print(f"   大小: {size_kb:.2f} KB")

            # 读取摘要
            summary_file = filepath.parent / f'{filepath.name}.summary.json'
            if summary_file.exists():
                with open(summary_file, 'r', encoding='utf-8') as f:
                    summary = json.load(f)
                print(f"   会话数: {summary['total_sessions']}")
                print(f"   有角色: {summary['sessions_with_characters']}")
            print()


def export_current_sessions():
    """
    导出当前服务器中的所有会话

    使用方法：
    1. 在服务器运行时，在Python控制台执行此函数
    2. 或者通过HTTP API调用
    """
    try:
        # 这个函数需要在服务器进程内部调用
        # 或者通过添加一个API端点来触发
        from coc_apiagent import sessions, sessions_lock

        persistence = SessionPersistence()

        with sessions_lock:
            # 复制当前sessions
            sessions_copy = dict(sessions)

        filepath = persistence.export_sessions(sessions_copy)
        return filepath
    except ImportError:
        print("[!] 无法导入sessions，请确保在服务器环境中运行")
        return None


def main():
    """命令行工具"""
    import sys

    persistence = SessionPersistence()

    if len(sys.argv) < 2:
        print("用法:")
        print("  python session_persistence.py list          - 列出所有备份")
        print("  python session_persistence.py export        - 导出当前会话")
        print("  python session_persistence.py import [file] - 导入会话")
        return

    command = sys.argv[1]

    if command == 'list':
        persistence.list_backups()

    elif command == 'export':
        filepath = export_current_sessions()
        if filepath:
            print(f"\n[OK] 导出成功: {filepath}")

    elif command == 'import':
        filename = sys.argv[2] if len(sys.argv) > 2 else None
        sessions = persistence.import_sessions(filename)
        print(f"\n[OK] 导入了 {len(sessions)} 个会话")

    else:
        print(f"未知命令: {command}")


if __name__ == '__main__':
    main()