# utils/context_manager.py - 上下文管理器（集成对话持久化和Token统计） import os import json import base64 import mimetypes import io import uuid import platform import shutil import subprocess from copy import deepcopy from typing import Dict, List, Optional, Any from pathlib import Path from datetime import datetime try: from config import ( MAX_CONTEXT_SIZE, DATA_DIR, PROMPTS_DIR, TERMINAL_SANDBOX_MOUNT_PATH, TERMINAL_SANDBOX_CPUS, TERMINAL_SANDBOX_MEMORY, PROJECT_MAX_STORAGE_MB, TERMINAL_SANDBOX_MODE, LINUX_SAFETY, ) from config.model_profiles import ( get_model_prompt_replacements, get_registered_model_keys, model_supports_image, model_supports_video, ) 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 ( MAX_CONTEXT_SIZE, DATA_DIR, PROMPTS_DIR, TERMINAL_SANDBOX_MOUNT_PATH, TERMINAL_SANDBOX_CPUS, TERMINAL_SANDBOX_MEMORY, PROJECT_MAX_STORAGE_MB, TERMINAL_SANDBOX_MODE, LINUX_SAFETY, ) from config.model_profiles import ( get_model_prompt_replacements, get_registered_model_keys, model_supports_image, model_supports_video, ) from utils.conversation_manager import ConversationManager from utils.host_workspace_debug import write_host_workspace_debug from utils.media_store import MediaStore from utils.token_usage import normalize_usage_payload AUTO_SHALLOW_PLACEHOLDER = "过早的工具结果已经被自动压缩" AUTO_SHALLOW_TOOL_WHITELIST = { "write_file", "read_file", "edit_file", "terminal_input", "terminal_snapshot", "web_search", "extract_webpage", "run_command", "view_image", "view_video", } class ProjectMixin: """ContextManager project mixin 能力 mixin。""" def _ensure_project_snapshot(self) -> Dict[str, Any]: """ 确保当前对话拥有项目文件树快照： - 若已缓存/存档则直接返回； - 若不存在，则第一次扫描目录并存入对话文件，后续复用。 """ if self.project_snapshot: return self.project_snapshot meta = self.conversation_metadata or {} stored_tree = meta.get("project_file_tree") if stored_tree and stored_tree != "宿主机模式下文件树不可用": use_stored_snapshot = True stored_path = meta.get("project_path") if isinstance(stored_path, str) and stored_path.strip(): try: stored_path_obj = Path(stored_path.strip()).expanduser().resolve() if stored_path_obj != Path(self.project_path).resolve(): use_stored_snapshot = False write_host_workspace_debug( "context_manager.project_snapshot.discard_stale_snapshot", conversation_id=self.current_conversation_id, metadata_project_path=str(stored_path_obj), current_project_path=str(Path(self.project_path).resolve()), ) except Exception: pass if use_stored_snapshot: self.project_snapshot = { "file_tree": stored_tree, "statistics": meta.get("project_statistics"), "snapshot_at": meta.get("project_snapshot_at") } write_host_workspace_debug( "context_manager.project_snapshot.use_metadata_snapshot", conversation_id=self.current_conversation_id, project_path=str(self.project_path), ) return self.project_snapshot # 首次生成并缓存 structure = self._get_project_structure_for_prompt() snapshot = { "file_tree": self._build_file_tree(structure), "statistics": { "total_files": structure["total_files"], "total_size": structure["total_size"] }, "snapshot_at": datetime.now().isoformat() } self.project_snapshot = snapshot write_host_workspace_debug( "context_manager.project_snapshot.generated_new", conversation_id=self.current_conversation_id, project_path=str(self.project_path), total_files=snapshot.get("statistics", {}).get("total_files"), ) if self.current_conversation_id: self.conversation_manager.update_project_snapshot( self.current_conversation_id, project_file_tree=snapshot["file_tree"], project_statistics=snapshot["statistics"], project_snapshot_at=snapshot["snapshot_at"] ) # 同步内存元数据 self.conversation_metadata["project_file_tree"] = snapshot["file_tree"] self.conversation_metadata["project_statistics"] = snapshot["statistics"] self.conversation_metadata["project_snapshot_at"] = snapshot["snapshot_at"] return snapshot def _get_project_structure_for_prompt(self, limit: int = 20) -> Dict: """获取用于 prompt 的浅层文件结构（仅根目录，优先文件夹）。""" structure = { "path": str(self.project_path), "files": [], "folders": [], "total_files": 0, "total_size": 0, "tree": {} } if not self.project_path.exists(): return structure try: entries = [p for p in self.project_path.iterdir() if not p.name.startswith('.')] except PermissionError: return structure folders = [p for p in entries if p.is_dir()] files = [p for p in entries if p.is_file()] folders.sort(key=lambda p: p.name.lower()) files.sort(key=lambda p: p.name.lower()) selected = (folders + files)[:max(0, limit)] def add_entry(entry: Path, *, display_path: Optional[str] = None): relative_path = str(entry.relative_to(self.project_path)) if display_path: relative_path = display_path if entry.is_dir(): structure["folders"].append({ "name": entry.name, "path": relative_path }) parts = relative_path.split("/") parent = structure["tree"] for idx, part in enumerate(parts): is_leaf = idx == len(parts) - 1 parent.setdefault(part, { "type": "folder", "path": "/".join(parts[:idx + 1]), "children": {} }) if is_leaf: parent[part]["type"] = "folder" parent[part]["path"] = relative_path parent[part].setdefault("children", {}) parent = parent[part].setdefault("children", {}) else: try: size = entry.stat().st_size except OSError: size = 0 file_info = { "name": entry.name, "path": relative_path, "size": size, "annotation": self.file_annotations.get(relative_path, "") } structure["files"].append(file_info) structure["total_files"] += 1 structure["total_size"] += size structure["tree"][entry.name] = { "type": "file", "path": relative_path, "size": size, "annotation": file_info["annotation"] } for entry in selected: add_entry(entry) # 根目录浅扫会跳过点目录；内部文件迁移到 .agents 后，单独展示这几个对模型有意义的目录。 agents_dir = self.project_path / ".agents" for internal_name in ("skills", "user_upload", "compact_result", "review"): internal_dir = agents_dir / internal_name if internal_dir.exists() and internal_dir.is_dir(): add_entry(internal_dir, display_path=f".astrion/{internal_name}") return structure def get_project_structure(self) -> Dict: """获取项目文件结构""" structure = { "path": str(self.project_path), "files": [], "folders": [], "total_files": 0, "total_size": 0, "tree": {} # 新增：树形结构数据 } if self._is_host_mode_without_safety(): structure["unavailable"] = True structure["message"] = "宿主机模式下文件树不可用" return structure # 记录实际存在的文件 existing_files = set() def scan_directory(path: Path, level: int = 0, max_level: int = 5, parent_tree: Dict = None): if level > max_level: return if parent_tree is None: parent_tree = structure["tree"] try: # 获取目录内容并排序（文件夹在前，文件在后） items = sorted(path.iterdir(), key=lambda x: (not x.is_dir(), x.name.lower())) for item in items: if item.name.startswith('.'): continue relative_path = str(item.relative_to(self.project_path)) if item.is_file(): existing_files.add(relative_path) # 记录存在的文件 file_info = { "name": item.name, "path": relative_path, "size": item.stat().st_size, "modified": datetime.fromtimestamp(item.stat().st_mtime).isoformat(), "annotation": self.file_annotations.get(relative_path, "") } structure["files"].append(file_info) structure["total_files"] += 1 structure["total_size"] += file_info["size"] # 添加到树形结构 parent_tree[item.name] = { "type": "file", "path": relative_path, "size": file_info["size"], "annotation": file_info["annotation"] } elif item.is_dir(): folder_info = { "name": item.name, "path": relative_path } structure["folders"].append(folder_info) # 创建文件夹节点 parent_tree[item.name] = { "type": "folder", "path": relative_path, "children": {} } # 递归扫描子目录 scan_directory(item, level + 1, max_level, parent_tree[item.name]["children"]) except PermissionError: pass scan_directory(self.project_path) # 清理不存在文件的备注 invalid_annotations = [] for annotation_path in self.file_annotations.keys(): if annotation_path not in existing_files: invalid_annotations.append(annotation_path) if invalid_annotations: for path in invalid_annotations: del self.file_annotations[path] print(f"🧹 清理无效备注: {path}") self.save_annotations() return structure def _build_file_tree(self, structure: Dict) -> str: """构建文件树字符串（修复版：正确显示树形结构）""" if not structure.get("tree"): return f"📁 {structure['path']}/\n(空项目)" lines = [] project_name = Path(structure['path']).name if self._is_host_mode_without_safety(): root_label = f"{structure['path']} (项目根)" else: container_root = (self.container_mount_path or "").strip() or "/workspace" container_root = container_root.rstrip("/") or "/" if container_root == "/": root_label = f"/ (项目根)" elif container_root.endswith(project_name): root_label = f"{container_root} (项目根)" else: root_label = f"{container_root} (映射自 {project_name})" lines.append(f"{root_label}/") ROOT_FOLDER_CHILD_LIMIT = 20 def count_descendants(item: Dict) -> int: """计算某个文件夹下（含多层）所有子项数量。""" if item.get("type") != "folder": return 0 children = item.get("children") or {} total = len(children) for child in children.values(): if child.get("type") == "folder": total += count_descendants(child) return total def build_tree_recursive(tree_dict: Dict, prefix: str = "", depth: int = 0): """递归构建树形结构""" if not tree_dict: return # 将项目按类型和名称排序：文件夹在前，文件在后，同类型按名称排序 items = list(tree_dict.items()) folders = [(name, info) for name, info in items if info["type"] == "folder"] files = [(name, info) for name, info in items if info["type"] == "file"] # 排序 folders.sort(key=lambda x: x[0].lower()) files.sort(key=lambda x: x[0].lower()) # 合并列表 sorted_items = folders + files for i, (name, info) in enumerate(sorted_items): is_last = (i == len(sorted_items) - 1) # 选择连接符 if is_last: current_connector = "└── " next_prefix = prefix + " " else: current_connector = "├── " next_prefix = prefix + "│ " if info["type"] == "folder": # 文件夹 lines.append(f"{prefix}{current_connector}{name}/") # 递归处理子项目 children = info.get("children") or {} if depth == 0: total_entries = count_descendants(info) else: total_entries = None if depth == 0 and total_entries is not None and total_entries > ROOT_FOLDER_CHILD_LIMIT: lines.append( f"{next_prefix}… (该目录包含 {total_entries} 项，已省略以控制 prompt 体积)" ) elif children: build_tree_recursive(children, next_prefix, depth + 1) else: # 文件 size_info = self._format_file_size(info['size']) # 构建文件行 file_line = f"{prefix}{current_connector}{name}" # 添加大小信息（简化版） if info['size'] > 1024: # 只显示大于1KB的文件大小 file_line += f" {size_info}" # 添加备注 if info.get('annotation'): file_line += f" # {info['annotation']}" lines.append(file_line) # 构建树形结构 build_tree_recursive(structure["tree"]) # 添加统计信息 lines.append("") lines.append(f"统计: {structure['total_files']} 个文件, {structure['total_size']/1024/1024:.2f}MB") return "\n".join(lines) def _format_file_size(self, size_bytes: int) -> str: """格式化文件大小""" if size_bytes < 1024: return f"({size_bytes}B)" elif size_bytes < 1024 * 1024: return f"({size_bytes/1024:.1f}KB)" else: return f"({size_bytes/1024/1024:.1f}MB)" def _get_file_icon(self, filename: str) -> str: """根据文件类型返回合适的图标（当前已禁用 emoji，返回空字符串）。""" return ""