agent-Specialization/utils/conversation_manager/index_mixin.py

# utils/conversation_manager.py - 对话持久化管理器（集成Token统计）

import json
import os
import time
import tempfile
import threading
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional, Any
from dataclasses import dataclass
try:
    from config import DATA_DIR, HOST_WORKSPACES_FILE
except ImportError:
    import sys
    from pathlib import Path
    project_root = Path(__file__).resolve().parents[1]
    if str(project_root) not in sys.path:
        sys.path.insert(0, str(project_root))
    from config import DATA_DIR, HOST_WORKSPACES_FILE

try:
    from utils.perf_log import perf_log
except Exception:
    perf_log = None

@dataclass
class ConversationMetadata:
    """对话元数据"""
    id: str
    title: str
    created_at: str
    updated_at: str
    project_path: Optional[str]
    project_relative_path: Optional[str]
    thinking_mode: bool
    total_messages: int
    total_tools: int
    run_mode: str = "fast"
    model_key: Optional[str] = None
    has_images: bool = False
    has_videos: bool = False
    status: str = "active"  # active, archived, error


class IndexMixin:
    """ConversationManager index mixin 能力 mixin。"""

    def _migrate_legacy_conversation_files(self):
        """把 data/conversations/conv_*.json 按 metadata.project_path 移入对应 workspace 子目录。"""
        try:
            legacy_files = [
                p for p in self.conversations_root.glob("conv_*.json")
                if p.is_file() and p.parent == self.conversations_root
            ]
        except Exception:
            return

        moved_count = 0
        for file_path in legacy_files:
            try:
                data = json.loads(file_path.read_text(encoding="utf-8"))
                metadata = data.get("metadata", {}) if isinstance(data, dict) else {}
                project_path = self._normalize_path(metadata.get("project_path"))
                workspace_id = self._resolve_workspace_id(project_path)
                if not workspace_id:
                    continue
                target_dir = self._conversation_dir_for_workspace(workspace_id)
                target_dir.mkdir(parents=True, exist_ok=True)
                target = target_dir / file_path.name
                if target.exists():
                    print(f"⚠️ legacy 对话迁移跳过，目标已存在: {target}")
                    continue
                file_path.replace(target)
                moved_count += 1
            except Exception as exc:
                print(f"⚠️ legacy 对话迁移失败 {file_path.name}: {exc}")
        if moved_count:
            print(f"🔄 已迁移 {moved_count} 个 legacy 对话到工作区目录")

    def _iter_conversation_files(self, sort_by_mtime: bool = True):
        """遍历对话文件（排除索引文件），可按修改时间降序排序。"""
        files = []
        for p in self.conversations_dir.glob("*.json"):
            if p == self.index_file:
                continue
            stem = p.stem or ""
            # 跳过索引备份/损坏文件，避免被误当作对话文件参与重建
            if stem == "index" or stem.startswith("index_corrupt_"):
                continue
            # 仅纳入标准对话文件，避免其他 json 干扰索引
            if not stem.startswith("conv_"):
                continue
            files.append(p)
        if sort_by_mtime:
            files.sort(key=lambda p: p.stat().st_mtime, reverse=True)
        return files

    def _atomic_write_json(self, target_file: Path, payload: Dict[str, Any]):
        """原子写入 JSON：使用唯一临时文件，避免并发写同名 .tmp 产生竞态。"""
        target_file.parent.mkdir(parents=True, exist_ok=True)
        temp_path: Optional[Path] = None
        try:
            with tempfile.NamedTemporaryFile(
                mode="w",
                encoding="utf-8",
                dir=str(target_file.parent),
                prefix=f".{target_file.name}.",
                suffix=".tmp",
                delete=False,
            ) as fh:
                temp_path = Path(fh.name)
                json.dump(payload, fh, ensure_ascii=False, indent=2)
                fh.flush()
                os.fsync(fh.fileno())
            os.replace(str(temp_path), str(target_file))
        finally:
            if temp_path and temp_path.exists():
                try:
                    temp_path.unlink()
                except Exception:
                    pass

    def _rebuild_index_from_files(self, max_count: Optional[int] = None) -> Dict:
        """
        从现有对话文件重建索引。

        Args:
            max_count: 限制重建的条目数（按文件修改时间倒序）；None 表示全量重建。
        """
        t0 = time.perf_counter()
        rebuilt_index: Dict[str, Dict] = {}
        files = self._iter_conversation_files(sort_by_mtime=True)
        if max_count is not None:
            files = files[:max(0, int(max_count))]
        if perf_log:
            perf_log("_rebuild_index_from_files start", extra={"file_count": len(files), "max_count": max_count})
        for file_path in files:
            try:
                with open(file_path, "r", encoding="utf-8") as f:
                    raw = f.read().strip()
                    if not raw:
                        continue
                    data = json.loads(raw)
            except Exception as exc:
                print(f"⚠️ 重建索引时跳过 {file_path.name}: {exc}")
                continue

            conv_id = data.get("id") or file_path.stem
            metadata = data.get("metadata", {}) or {}

            rebuilt_index[conv_id] = {
                "title": data.get("title") or "未命名对话",
                "created_at": data.get("created_at"),
                "updated_at": data.get("updated_at"),
                "project_path": metadata.get("project_path"),
                "project_relative_path": metadata.get("project_relative_path"),
                "thinking_mode": metadata.get("thinking_mode", False),
                "run_mode": metadata.get("run_mode") or ("thinking" if metadata.get("thinking_mode") else "fast"),
                "model_key": metadata.get("model_key"),
                "has_images": metadata.get("has_images", False),
                "has_videos": metadata.get("has_videos", False),
                "total_messages": metadata.get("total_messages", 0),
                "total_tools": metadata.get("total_tools", 0),
                "status": metadata.get("status", "active"),
            }
        elapsed_ms = (time.perf_counter() - t0) * 1000
        if perf_log:
            perf_log("_rebuild_index_from_files done", elapsed_ms=elapsed_ms, extra={"rebuilt_size": len(rebuilt_index), "max_count": max_count})
        if rebuilt_index:
            print(f"🔄 已从对话文件重建索引，共 {len(rebuilt_index)} 条记录")
        return rebuilt_index

    def _index_missing_conversations(self, index: Dict) -> bool:
        """检查索引是否缺失本地对话文件"""
        t0 = time.perf_counter()
        index_ids = set(index.keys())
        files = list(self._iter_conversation_files())
        missing = False
        for file_path in files:
            conv_id = file_path.stem
            if conv_id and conv_id not in index_ids:
                print(f"🔍 对话 {conv_id} 未出现在索引中，将重建索引。")
                missing = True
                break
        elapsed_ms = (time.perf_counter() - t0) * 1000
        if perf_log:
            perf_log("_index_missing_conversations", elapsed_ms=elapsed_ms, extra={"file_count": len(files), "index_size": len(index), "missing": missing})
        return missing

    def _load_index(self, ensure_integrity: bool = False, max_rebuild: Optional[int] = None) -> Dict:
        """加载对话索引，可选地在缺失时自动重建（可限制重建条数）"""
        try:
            index: Dict = {}
            if self.index_file.exists():
                with open(self.index_file, 'r', encoding='utf-8') as f:
                    content = f.read().strip()
                    if content:
                        index = json.loads(content)
                        if index:
                            if ensure_integrity and not self._index_verified:
                                if self._index_missing_conversations(index):
                                    rebuilt = self._rebuild_index_from_files(max_count=max_rebuild)
                                    if rebuilt:
                                        self._save_index(rebuilt)
                                        index = rebuilt
                                self._index_verified = True
                            return index
                        # 索引为空但对话文件仍然存在时尝试重建
                        rebuilt = self._rebuild_index_from_files(max_count=max_rebuild)
                        if rebuilt:
                            self._save_index(rebuilt)
                            if ensure_integrity:
                                self._index_verified = True
                            return rebuilt
                        return {}
            # 索引缺失但存在对话文件时重建
            rebuilt = self._rebuild_index_from_files(max_count=max_rebuild)
            if rebuilt:
                self._save_index(rebuilt)
                if ensure_integrity:
                    self._index_verified = True
                return rebuilt
            return index
        except (json.JSONDecodeError, Exception) as e:
            print(f"⚠️ 加载对话索引失败，将尝试重建: {e}")
            backup_path = self.index_file.with_name(
                f"{self.index_file.stem}_corrupt_{int(time.time() * 1000)}{self.index_file.suffix}"
            )
            try:
                if self.index_file.exists():
                    self.index_file.replace(backup_path)
                    print(f"🗄️ 已备份损坏的索引文件到: {backup_path.name}")
            except Exception as backup_exc:
                print(f"⚠️ 备份损坏索引文件失败: {backup_exc}")
            # 索引损坏场景优先全量重建，避免仅重建部分导致“对话丢失”错觉
            rebuilt = self._rebuild_index_from_files(max_count=None)
            if rebuilt:
                self._save_index(rebuilt)
                if ensure_integrity:
                    self._index_verified = True
                return rebuilt
            return {}

    def _save_index(self, index: Dict):
        """保存对话索引"""
        try:
            with self._io_lock:
                self._atomic_write_json(self.index_file, index)
        except Exception as e:
            print(f"⌘ 保存对话索引失败: {e}")

    def _ensure_index_covering(self, limit: int, offset: int) -> Dict:
        """
        确保索引涵盖到 offset+limit 条记录，不足时按需扩展重建（仍按 mtime 倒序，增量加载批量）。
        每次先检查是否有新增对话文件未被纳入索引，避免"创建后列表不显示"的问题。
        """
        t0 = time.perf_counter()
        needed = max(0, int(offset) + int(limit))
        index = self._load_index()
        load_ms = (time.perf_counter() - t0) * 1000
        if perf_log:
            perf_log("_ensure_index_covering load_index", elapsed_ms=load_ms, extra={"index_size": len(index), "needed": needed})

        # 主动检查新增文件：索引可能滞后于实际对话文件，导致列表不同步。
        t1 = time.perf_counter()
        missing = self._index_missing_conversations(index)
        missing_ms = (time.perf_counter() - t1) * 1000
        if perf_log:
            perf_log("_ensure_index_covering check_missing", elapsed_ms=missing_ms, extra={"missing": missing, "index_size": len(index)})
        if missing:
            t2 = time.perf_counter()
            rebuilt = self._rebuild_index_from_files(max_count=max(needed, len(index) + 10))
            rebuild_ms = (time.perf_counter() - t2) * 1000
            if perf_log:
                perf_log("_ensure_index_covering rebuild_missing", elapsed_ms=rebuild_ms, extra={"rebuilt_size": len(rebuilt) if rebuilt else 0})
            if rebuilt:
                self._save_index(rebuilt)
                index = rebuilt

        if len(index) >= needed:
            total_ms = (time.perf_counter() - t0) * 1000
            if perf_log:
                perf_log("_ensure_index_covering return", elapsed_ms=total_ms, extra={"index_size": len(index), "needed": needed})
            return index

        # 第一次尝试：扩展到需要的数量（按更新时间倒序）
        t3 = time.perf_counter()
        rebuilt = self._rebuild_index_from_files(max_count=needed)
        rebuild_ms = (time.perf_counter() - t3) * 1000
        if perf_log:
            perf_log("_ensure_index_covering rebuild_needed", elapsed_ms=rebuild_ms, extra={"rebuilt_size": len(rebuilt) if rebuilt else 0})
        if rebuilt:
            self._save_index(rebuilt)
            index = rebuilt

        # 如果仍不足且存在更多文件可能未被纳入（例如首批限定过小），进行一次全量重建兜底
        if len(index) < needed:
            t4 = time.perf_counter()
            rebuilt_full = self._rebuild_index_from_files(max_count=None)
            rebuild_full_ms = (time.perf_counter() - t4) * 1000
            if perf_log:
                perf_log("_ensure_index_covering rebuild_full", elapsed_ms=rebuild_full_ms, extra={"rebuilt_size": len(rebuilt_full) if rebuilt_full else 0})
            if rebuilt_full:
                self._save_index(rebuilt_full)
                index = rebuilt_full

        total_ms = (time.perf_counter() - t0) * 1000
        if perf_log:
            perf_log("_ensure_index_covering return", elapsed_ms=total_ms, extra={"index_size": len(index), "needed": needed})
        return index