- 工作区内部路径 .agents/ -> .astrion/ - 默认运行态数据根 ~/.agents/agents/ -> ~/.astrion/astrion/ - 环境变量 AGENTS_DATA_ROOT -> ASTRION_DATA_ROOT(无兼容) - 同步更新代码、测试、文档与脚本中的路径引用 - 新增 .gitignore 忽略 .astrion/ 与 .agents/ 运行态目录
435 lines
17 KiB
Python
435 lines
17 KiB
Python
# 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 ""
|