# 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 CompressionMixin: """ContextManager compression mixin 能力 mixin。""" def is_compression_in_progress(self) -> bool: return bool(self._get_meta_flag("compression_in_progress", False)) def set_compression_state( self, *, in_progress: bool, mode: Optional[str] = None, stage: Optional[str] = None, error: Optional[str] = None, resume_payload: Optional[Dict[str, Any]] = None, job_id: Optional[str] = None, ): updates: Dict[str, Any] = { "compression_in_progress": bool(in_progress), "compression_mode": mode if in_progress else None, "compression_stage": stage if in_progress else None, "compression_error": error if in_progress else None, "compression_resume_payload": resume_payload if in_progress else None, "compression_job_id": job_id if in_progress else None, } for k, v in updates.items(): self.conversation_metadata[k] = v if self.current_conversation_id: try: target_manager = ( self._get_conversation_manager_for_id(self.current_conversation_id) if hasattr(self, "_get_conversation_manager_for_id") else self.conversation_manager ) target_manager.update_conversation_metadata(self.current_conversation_id, updates) except Exception as exc: print(f"[ContextCompression] 更新压缩状态失败: {exc}") def on_tool_call_finished( self, tool_name: Optional[str] = None, *, enable_shallow: bool = True, shallow_trigger_tokens: int = 80_000, deep_trigger_tokens: int = 150_000, shallow_batch_size: int = 10, shallow_keep_recent_tools: int = 15, shallow_trigger_tool_calls_interval: int = 10, shallow_keep_user_turn_tools: int = 3, ) -> int: """每次工具调用完成后触发：更新计数并执行自动浅压缩。返回本轮浅压缩条数。""" if not self.current_conversation_id: return 0 # 兼容历史对话：tool_call_count 可能从 0 开始，但历史里已存在大量 tool 结果。 # 这里优先对齐到“当前历史中的 tool 消息数量”，避免长期无法触发浅压缩。 meta_tool_count = int(self._get_meta_flag("tool_call_count", 0) or 0) history_tool_count = 0 try: history_tool_count = sum( 1 for msg in (self.conversation_history or []) if isinstance(msg, dict) and msg.get("role") == "tool" ) except Exception: history_tool_count = 0 tool_count = max(meta_tool_count + 1, history_tool_count) self._set_meta_flag("tool_call_count", tool_count, save=False) current_tokens = self.get_current_context_tokens(self.current_conversation_id) shallow_trigger_tokens = max(1, int(80_000 if shallow_trigger_tokens is None else shallow_trigger_tokens)) deep_trigger_tokens = max( shallow_trigger_tokens + 1, int(150_000 if deep_trigger_tokens is None else deep_trigger_tokens), ) shallow_batch_size = max(1, int(10 if shallow_batch_size is None else shallow_batch_size)) shallow_keep_recent_tools = max(0, int(15 if shallow_keep_recent_tools is None else shallow_keep_recent_tools)) shallow_keep_user_turn_tools = max(0, int(3 if shallow_keep_user_turn_tools is None else shallow_keep_user_turn_tools)) shallow_trigger_tool_calls_interval = max( 1, int(10 if shallow_trigger_tool_calls_interval is None else shallow_trigger_tool_calls_interval), ) was_long_conversation = bool(self._get_meta_flag("is_long_conversation", False)) if current_tokens > shallow_trigger_tokens and not was_long_conversation: self._set_meta_flag("is_long_conversation", True, save=False) just_marked_long = (not was_long_conversation) and bool(self._get_meta_flag("is_long_conversation", False)) if current_tokens > deep_trigger_tokens and not self._get_meta_flag("is_ultra_long_conversation", False): self._set_meta_flag("is_ultra_long_conversation", True, save=False) # 每 10 次工具调用触发一次浅压缩（当 long 已标记） last_checkpoint = int(self._get_meta_flag("last_shallow_compress_tool_count", 0) or 0) should_try_shallow = ( bool(enable_shallow) and bool(self._get_meta_flag("is_long_conversation", False)) and ( just_marked_long or (last_checkpoint <= 0 and tool_count > 0) or (tool_count - last_checkpoint >= shallow_trigger_tool_calls_interval) ) ) changed = False compressed_count = 0 if should_try_shallow: compressed = self._run_auto_shallow_compression( batch_size=shallow_batch_size, keep_recent_tools=shallow_keep_recent_tools, keep_user_turn_tools=shallow_keep_user_turn_tools, ) self._set_meta_flag("last_shallow_compress_tool_count", tool_count, save=False) compressed_count = max(0, int(compressed or 0)) changed = compressed > 0 if self.current_conversation_id: updates = { "tool_call_count": tool_count, "is_long_conversation": bool(self._get_meta_flag("is_long_conversation", False)), "is_ultra_long_conversation": bool(self._get_meta_flag("is_ultra_long_conversation", False)), "last_shallow_compress_tool_count": int(self._get_meta_flag("last_shallow_compress_tool_count", 0) or 0), } try: self.conversation_manager.update_conversation_metadata(self.current_conversation_id, updates) except Exception as exc: print(f"[ContextCompression] 写入计数失败: {exc}") if changed: self.auto_save_conversation(force=True) return compressed_count def _run_auto_shallow_compression(self, batch_size: int = 10, keep_recent_tools: int = 15, keep_user_turn_tools: int = 3) -> int: """浅压缩：仅打标记，不修改原文。返回本轮标记条数。 Args: batch_size: 每轮最大压缩数量 keep_recent_tools: 保留最近的N条工具消息不压缩 keep_user_turn_tools: 保留最近N次用户输入后的工具消息不压缩 """ history = self.conversation_history or [] if not history: return 0 tool_indices = [idx for idx, msg in enumerate(history) if msg.get("role") == "tool"] protected_indices = set(tool_indices[-max(0, keep_recent_tools):]) if keep_recent_tools > 0 else set() # 新增机制：保留最近keep_user_turn_tools次用户输入后的工具消息 if keep_user_turn_tools > 0: # 从后往前找真正的用户输入（排除系统自动发送的） user_turn_indices = [] for idx in range(len(history) - 1, -1, -1): msg = history[idx] if msg.get("role") == "user": # 检查是否是系统自动发送的消息（向后兼容：检查metadata或前缀文本） metadata = msg.get("metadata") or {} content = msg.get("content") or "" is_auto = metadata.get("is_auto_generated") or \ content.startswith("这是一句系统自动发送的user消息，用于通知你") if not is_auto: user_turn_indices.append(idx) if len(user_turn_indices) >= keep_user_turn_tools: break # 保护这些用户输入之后的所有工具消息（直到下一个user消息） for user_idx in user_turn_indices: for idx in range(user_idx + 1, len(history)): msg = history[idx] if msg.get("role") == "user": break # 遇到下一个user消息，停止 if msg.get("role") == "tool": protected_indices.add(idx) candidates: List[int] = [] for idx, msg in enumerate(history): if msg.get("role") != "tool": continue if idx in protected_indices: continue tool_name = (msg.get("name") or "").strip() if tool_name not in AUTO_SHALLOW_TOOL_WHITELIST: continue metadata = msg.get("metadata") or {} if metadata.get("auto_shallow_compacted"): continue candidates.append(idx) if len(candidates) >= max(1, batch_size): break if not candidates: return 0 self._shallow_compact_round += 1 now = datetime.now().isoformat() for idx in candidates: msg = history[idx] metadata = msg.get("metadata") or {} metadata["auto_shallow_compacted"] = True metadata["auto_shallow_compacted_at"] = now metadata["auto_shallow_compact_round"] = self._shallow_compact_round msg["metadata"] = metadata history[idx] = msg self.conversation_history = history return len(candidates) def compress_conversation(self, conversation_id: str) -> Dict: """ 压缩指定对话：保留用户/助手原文（不含 reasoning），提取工具意图/名称， 生成一条 system 消息作为新对话的压缩版历史。 """ conversation_data = self.conversation_manager.load_conversation(conversation_id) if not conversation_data: return { "success": False, "error": f"对话不存在: {conversation_id}" } original_messages = conversation_data.get("messages", []) or [] if not original_messages: return { "success": False, "error": "当前对话没有可压缩的内容" } header_text = ( f"系统提示：根据压缩后的工作记录继续这个任务。" f"如果信息不足，提示用户使用对话回顾功能。源对话：{conversation_id}" ) lines: List[str] = [] tool_buffer: List[str] = [] seen_tool_call_ids = set() def add_spacing(): if lines and lines[-1] != "": lines.append("") def flush_tools(): if not tool_buffer: return add_spacing() lines.append("tool：") lines.extend(f"- {entry}" for entry in tool_buffer) tool_buffer.clear() for message in original_messages: role = message.get("role") if role == "user": flush_tools() content = message.get("content") or "" add_spacing() lines.append(f"user：{content}") continue if role == "assistant": content = message.get("content") or "" has_visible_content = bool(str(content).strip()) if has_visible_content: flush_tools() add_spacing() lines.append(f"assistant：{content}") tool_calls = message.get("tool_calls") or [] for tc in tool_calls: tc_id = tc.get("id") or tc.get("tool_call_id") if tc_id: seen_tool_call_ids.add(tc_id) func = tc.get("function") or {} arguments = func.get("arguments") args_obj = {} if isinstance(arguments, str): try: args_obj = json.loads(arguments) except Exception: args_obj = {} elif isinstance(arguments, dict): args_obj = arguments intent = args_obj.get("intent") if isinstance(args_obj, dict) else None name = func.get("name") or tc.get("name") or "unknown_tool" entry = intent.strip() if isinstance(intent, str) and intent.strip() else name tool_buffer.append(entry) continue if role == "tool": tc_id = message.get("tool_call_id") or message.get("id") if tc_id and tc_id in seen_tool_call_ids: # 已经通过 intent 记录，无需重复 continue name = message.get("name") or "unknown_tool" tool_buffer.append(name) continue # 其他角色（如 system）原样保留 flush_tools() content = message.get("content") or "" add_spacing() lines.append(f"{role}：{content}" if role else content) flush_tools() summary_text = header_text + "\n\n" + "\n".join(lines) system_message = { "role": "system", "content": summary_text, "timestamp": datetime.now().isoformat(), "metadata": { "compression": { "source_conversation_id": conversation_id, "types": ["intent_summary"], "created_at": datetime.now().isoformat() } } } metadata = conversation_data.get("metadata", {}) resolved_project_path = self._resolve_project_path_from_metadata(metadata) project_path = str(resolved_project_path) thinking_mode = metadata.get("thinking_mode", False) run_mode = metadata.get("run_mode") or ("thinking" if thinking_mode else "fast") model_key = metadata.get("model_key") has_images = metadata.get("has_images", False) original_title = conversation_data.get("title") compressed_conversation_id = self.conversation_manager.create_conversation( project_path=project_path, thinking_mode=thinking_mode, run_mode=run_mode, initial_messages=[system_message], model_key=model_key, has_images=has_images, metadata_overrides={ "permission_mode": metadata.get("permission_mode", "unrestricted"), }, ) # 设置压缩后的对话标题 if original_title: try: new_title = f"{original_title} 压缩后" self.conversation_manager.update_conversation_title(compressed_conversation_id, new_title) except Exception: pass return { "success": True, "compressed_conversation_id": compressed_conversation_id, "compressed_types": ["intent_summary"], "system_message": summary_text }