agent-Specialization/core/main_terminal_parts/context.py

import asyncio
import json
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional, Set

try:
    from config import (
        OUTPUT_FORMATS, DATA_DIR, PROMPTS_DIR, NEED_CONFIRMATION,
        MAX_TERMINALS, TERMINAL_BUFFER_SIZE, TERMINAL_DISPLAY_SIZE,
        MAX_READ_FILE_CHARS, READ_TOOL_DEFAULT_MAX_CHARS,
        READ_TOOL_DEFAULT_CONTEXT_BEFORE, READ_TOOL_DEFAULT_CONTEXT_AFTER,
        READ_TOOL_MAX_CONTEXT_BEFORE, READ_TOOL_MAX_CONTEXT_AFTER,
        READ_TOOL_DEFAULT_MAX_MATCHES, READ_TOOL_MAX_MATCHES,
        READ_TOOL_MAX_FILE_SIZE,
        TERMINAL_SANDBOX_MOUNT_PATH,
        TERMINAL_SANDBOX_MODE,
        TERMINAL_SANDBOX_CPUS,
        TERMINAL_SANDBOX_MEMORY,
        PROJECT_MAX_STORAGE_MB,
        CUSTOM_TOOLS_ENABLED,
    )
except ImportError:
    import sys
    project_root = Path(__file__).resolve().parents[2]
    if str(project_root) not in sys.path:
        sys.path.insert(0, str(project_root))
    from config import (
        OUTPUT_FORMATS, DATA_DIR, PROMPTS_DIR, NEED_CONFIRMATION,
        MAX_TERMINALS, TERMINAL_BUFFER_SIZE, TERMINAL_DISPLAY_SIZE,
        MAX_READ_FILE_CHARS, READ_TOOL_DEFAULT_MAX_CHARS,
        READ_TOOL_DEFAULT_CONTEXT_BEFORE, READ_TOOL_DEFAULT_CONTEXT_AFTER,
        READ_TOOL_MAX_CONTEXT_BEFORE, READ_TOOL_MAX_CONTEXT_AFTER,
        READ_TOOL_DEFAULT_MAX_MATCHES, READ_TOOL_MAX_MATCHES,
        READ_TOOL_MAX_FILE_SIZE,
        TERMINAL_SANDBOX_MOUNT_PATH,
        TERMINAL_SANDBOX_MODE,
        TERMINAL_SANDBOX_CPUS,
        TERMINAL_SANDBOX_MEMORY,
        PROJECT_MAX_STORAGE_MB,
        CUSTOM_TOOLS_ENABLED,
    )

from modules.file_manager import FileManager
from modules.search_engine import SearchEngine
from modules.terminal_ops import TerminalOperator
from modules.memory_manager import MemoryManager
from modules.terminal_manager import TerminalManager
from modules.todo_manager import TodoManager
from modules.sub_agent_manager import SubAgentManager
from modules.webpage_extractor import extract_webpage_content, tavily_extract
from modules.ocr_client import OCRClient
from modules.easter_egg_manager import EasterEggManager
from modules.personalization_manager import (
    load_personalization_config,
    build_personalization_prompt,
)
from modules.skills_manager import (
    get_skills_catalog,
    build_skills_list,
    merge_enabled_skills,
    build_skills_prompt,
)
from modules.custom_tool_registry import CustomToolRegistry, build_default_tool_category
from modules.custom_tool_executor import CustomToolExecutor

try:
    from config.limits import THINKING_FAST_INTERVAL
except ImportError:
    THINKING_FAST_INTERVAL = 10

from modules.container_monitor import collect_stats, inspect_state
from core.tool_config import TOOL_CATEGORIES
from utils.api_client import DeepSeekClient
from utils.context_manager import ContextManager, AUTO_SHALLOW_PLACEHOLDER
from utils.host_workspace_debug import write_host_workspace_debug
from utils.tool_result_formatter import format_tool_result_for_context
from utils.logger import setup_logger
from config.model_profiles import (
    get_model_profile,
    get_model_prompt_replacements,
    get_model_context_window,
    model_supports_image,
    model_supports_video,
)

logger = setup_logger(__name__)
DISABLE_LENGTH_CHECK = True

class MainTerminalContextMixin:
    # 权限模式说明（代码中存储的简要版本）
    _PERMISSION_MODE_DESC = {
        "unrestricted": "当前处于无限制模式，所有工具均可直接使用，无需额外批准。",
        "readonly": "当前处于只读模式，仅能执行不会修改工作区的读取类操作。",
        "approval": "当前处于批准模式，修改工作区的操作需要用户批准后方可执行。",
    }

    _PERMISSION_MODE_LABEL = {
        "unrestricted": "无限制",
        "readonly": "只读",
        "approval": "批准",
    }

    _EXECUTION_MODE_LABEL = {
        "sandbox": "沙箱",
        "direct": "完全访问权限",
    }

    def _build_permission_mode_message(self) -> Optional[str]:
        """根据当前权限模式构建权限说明消息（从模板读取并替换占位符）"""
        template = self.load_prompt("permission_mode")
        if not template:
            return None

        mode = self.get_permission_mode()

        # 根据模式生成详细规则
        if mode == "unrestricted":
            detailed_rules = (
                "- 所有命令和工具均可直接使用，无需用户批准\n"
                "- run_command：支持任意命令，包括管道、重定向、子shell等\n"
                "- 文件操作：write_file、edit_file、create_file、delete_file、rename_file、create_folder 均可直接使用"
            )
        elif mode == "readonly":
            detailed_rules = (
                "- run_command：仅允许纯只读命令（ls, cat, grep, find, head, tail, wc, stat, file, sort, uniq 等）\n"
                "- 允许只读管道：如 `find . -name '*.py' | head -20`、`cat file | sort | uniq`\n"
                "- 禁止：子shell（$() ``）、逻辑运算符（&& || ;）、重定向（> <）\n"
                "- 禁止危险管道：管道指向 shell（bash/sh）、tee、xargs、chmod 等写入或执行操作\n"
                "- 文件操作：仅允许 read_file、view_image、view_video 等读取类工具\n"
                "- 禁止：write_file、edit_file、create_file、delete_file、rename_file、create_folder 等修改类工具"
            )
        elif mode == "approval":
            detailed_rules = (
                "- 免批准直通：纯只读命令（ls, cat, grep, find 等）及只读管道可直接执行\n"
                "- 免批准直通：如 `find . -name '*.py' | head -20`、`grep 'error' log | wc -l`\n"
                "- 需要用户批准：非只读命令（带重定向、子shell、逻辑运算符）\n"
                "- 需要用户批准：危险管道（指向 bash、tee、xargs、chmod 等）\n"
                "- 需要用户批准：terminal_input、run_python、write_file、edit_file、create_file、delete_file、rename_file、create_folder、save_webpage\n"
                "- 系统会在执行前弹出批准请求，请等待用户确认后再继续"
            )
        else:
            detailed_rules = ""

        return template.format(
            permission_mode=mode,
            permission_mode_label=self._PERMISSION_MODE_LABEL.get(mode, mode),
            mode_description=self._PERMISSION_MODE_DESC.get(mode, ""),
            detailed_rules=detailed_rules,
        )

    def _build_execution_mode_message(self) -> Optional[str]:
        """根据当前执行环境模式构建提示消息。"""
        # 仅宿主机模式注入；Docker 模式不需要该提示。
        try:
            if not getattr(self, "_is_host_mode", lambda: False)():
                return None
        except Exception:
            return None
        template = self.load_prompt("execution_mode")
        if not template:
            return None
        state = {}
        if hasattr(self, "get_execution_mode_state"):
            try:
                state = self.get_execution_mode_state() or {}
            except Exception:
                state = {}
        mode = str(state.get("mode") or "sandbox").strip().lower()
        mode_label = self._EXECUTION_MODE_LABEL.get(mode, mode)
        if mode == "sandbox":
            rules = (
                "- 所有命令默认在系统 OS 沙箱中执行\n"
                "- 若命令因系统权限或目录访问限制失败，请先给出替代方案（只读检查、分步验证）；若该操作确属必须，再提示用户切换为“完全访问权限”或让用户在本地授权白名单后重试"
            )
        else:
            rules = (
                "- 当前为宿主机直接执行模式（完全访问权限）\n"
                "- 仅在必须时执行高权限操作，保持最小化命令范围\n"
                "- 涉及删除/覆盖/系统级变更前，先说明风险再执行"
            )
        return template.format(
            execution_mode=mode,
            execution_mode_label=mode_label,
            rules=rules,
        )

    def build_context(self) -> Dict:
            """构建主终端上下文"""
            # 读取记忆
            memory = self.memory_manager.read_main_memory()

            # 构建上下文
            return self.context_manager.build_main_context(memory)

    def _tool_calls_followed_by_tools(self, conversation: List[Dict], start_idx: int, tool_calls: List[Dict]) -> bool:
            """判断指定助手消息的工具调用是否拥有后续的工具响应。"""
            if not tool_calls:
                return False
            expected_ids = [tc.get("id") for tc in tool_calls if tc.get("id")]
            if not expected_ids:
                return False
            matched: Set[str] = set()
            idx = start_idx + 1
            total = len(conversation)
            while idx < total and len(matched) < len(expected_ids):
                next_conv = conversation[idx]
                role = next_conv.get("role")
                if role == "tool":
                    call_id = next_conv.get("tool_call_id")
                    if call_id in expected_ids:
                        matched.add(call_id)
                    else:
                        break
                elif role in ("assistant", "user"):
                    break
                idx += 1
            return len(matched) == len(expected_ids)

    def build_messages(self, context: Dict, user_input: str) -> List[Dict]:
            """构建消息列表（添加终端内容注入）"""
            try:
                file_tree_preview = (context.get("project_info", {}).get("file_tree") or "").splitlines()
                write_host_workspace_debug(
                    "main_terminal.build_messages.context_snapshot",
                    terminal_id=id(self),
                    terminal_project_path=str(getattr(self, "project_path", "")),
                    context_project_path=str(getattr(getattr(self, "context_manager", None), "project_path", "")),
                    project_info_path=str(context.get("project_info", {}).get("path", "")),
                    file_tree_first_line=file_tree_preview[0] if file_tree_preview else "",
                    current_conversation_id=getattr(getattr(self, "context_manager", None), "current_conversation_id", None),
                )
            except Exception:
                pass
            # 加载系统提示（Qwen3.5 使用专用提示）
            current_model = getattr(self, "model_key", "kimi")
            prompt_name = "main_system_qwenvl" if (model_supports_image(current_model) or model_supports_video(current_model)) else "main_system"
            system_prompt = self.load_prompt(prompt_name)

            # 格式化系统提示
            container_path = self.container_mount_path or "/workspace"
            container_cpus = self.container_cpu_limit
            container_memory = self.container_memory_limit
            project_storage = self.project_storage_limit
            model_key = getattr(self, "model_key", "kimi")
            prompt_replacements = get_model_prompt_replacements(model_key)
            system_prompt = system_prompt.format(
                project_path=container_path,
                container_path=container_path,
                container_cpus=container_cpus,
                container_memory=container_memory,
                project_storage=project_storage,
                file_tree=context["project_info"]["file_tree"],
                memory=context["memory"],
                current_time=datetime.now().strftime("%Y-%m-%d %H"),
                model_description=prompt_replacements.get("model_description", "")
            )

            messages = [
                {"role": "system", "content": system_prompt}
            ]

            personalization_config = getattr(self.context_manager, "custom_personalization_config", None) or load_personalization_config(self.data_dir)
            shallow_replace_enabled = bool(personalization_config.get("auto_shallow_compress_enabled", False)) if isinstance(personalization_config, dict) else False
            skills_catalog = get_skills_catalog()
            enabled_skills = merge_enabled_skills(
                personalization_config.get("enabled_skills") if isinstance(personalization_config, dict) else None,
                skills_catalog,
                personalization_config.get("skills_catalog_snapshot") if isinstance(personalization_config, dict) else None,
            )
            skills_template = self.load_prompt("skills_system").strip()
            skills_list = build_skills_list(skills_catalog, enabled_skills)
            skills_prompt = build_skills_prompt(skills_template, skills_list)
            if skills_prompt:
                messages.append({"role": "system", "content": skills_prompt})

            workspace_system = self.context_manager._build_workspace_system_message(context)
            if workspace_system:
                messages.append({"role": "system", "content": workspace_system})

            # 注入权限模式说明（根据当前模式动态生成）
            permission_mode_message = self._build_permission_mode_message()
            if permission_mode_message:
                messages.append({"role": "system", "content": permission_mode_message})
            execution_mode_message = self._build_execution_mode_message()
            if execution_mode_message:
                messages.append({"role": "system", "content": execution_mode_message})

            if self.tool_category_states.get("sub_agent", True):
                sub_agent_prompt = self.load_prompt("sub_agent_guidelines").strip()
                if sub_agent_prompt:
                    messages.append({"role": "system", "content": sub_agent_prompt})

            if self.deep_thinking_mode:
                deep_prompt = self.load_prompt("deep_thinking_mode_guidelines").strip()
                if deep_prompt:
                    deep_prompt = deep_prompt.format(
                        deep_thinking_line=prompt_replacements.get("deep_thinking_line", "")
                    )
                    messages.append({"role": "system", "content": deep_prompt})
            elif self.thinking_mode:
                thinking_prompt = self.load_prompt("thinking_mode_guidelines").strip()
                if thinking_prompt:
                    thinking_prompt = thinking_prompt.format(
                        thinking_model_line=prompt_replacements.get("thinking_model_line", "")
                    )
                    messages.append({"role": "system", "content": thinking_prompt})

            # 支持按对话覆盖的个性化配置
            personalization_block = build_personalization_prompt(personalization_config, include_header=False)
            if personalization_block:
                personalization_template = self.load_prompt("personalization").strip()
                if personalization_template and "{personalization_block}" in personalization_template:
                    personalization_text = personalization_template.format(personalization_block=personalization_block)
                elif personalization_template:
                    personalization_text = f"{personalization_template}\n{personalization_block}"
                else:
                    personalization_text = personalization_block
                messages.append({"role": "system", "content": personalization_text})

            # AGENTS.md 自动注入
            agents_md_inject_enabled = bool(personalization_config.get("agents_md_auto_inject", False)) if isinstance(personalization_config, dict) else False
            if agents_md_inject_enabled:
                agents_md_content = self._load_agents_md_content()
                if agents_md_content:
                    agents_md_template = self.load_prompt("agents_md_inject").strip()
                    if agents_md_template and "{{AGENTS_MD_CONTENT}}" in agents_md_template:
                        agents_md_text = agents_md_template.replace("{{AGENTS_MD_CONTENT}}", agents_md_content)
                    else:
                        agents_md_text = f"【AGENTS.md 项目规范】\n\n{agents_md_content}\n\n---\n请注意：以上规范来自工作区根目录的 AGENTS.md 文件，若有冲突请以 AGENTS.md 为准。"
                    messages.append({"role": "system", "content": agents_md_text})

            # 支持按对话覆盖的自定义 system prompt（API 用途）。
            # 放在最后一个 system 消息位置，确保优先级最高，便于业务场景强约束。
            custom_system_prompt = getattr(self.context_manager, "custom_system_prompt", None)
            if isinstance(custom_system_prompt, str) and custom_system_prompt.strip():
                messages.append({"role": "system", "content": custom_system_prompt.strip()})

            # 禁用工具提示也作为“开头 system prompt”注入，便于后续统一合并。
            disabled_notice = self._format_disabled_tool_notice()
            if disabled_notice:
                messages.append({
                    "role": "system",
                    "content": disabled_notice
                })

            # 添加对话历史（保留完整结构，包括tool_calls和tool消息）
            conversation = context["conversation"]
            replaced_tool_count = 0
            for idx, conv in enumerate(conversation):
                metadata = conv.get("metadata") or {}
                if conv["role"] == "assistant":
                    # Assistant消息可能包含工具调用
                    message = {
                        "role": conv["role"],
                        "content": conv["content"]
                    }
                    # 对于思考模式（如 DeepSeek thinking），assistant 历史消息中的
                    # reasoning_content 即使为空字符串也需要原样回传，否则下一轮可能被
                    # API 判定为“未回传 reasoning_content”并返回 400。
                    if "reasoning_content" in conv:
                        message["reasoning_content"] = conv.get("reasoning_content", "")
                    # 如果有工具调用信息，添加到消息中
                    tool_calls = conv.get("tool_calls") or []
                    if tool_calls and self._tool_calls_followed_by_tools(conversation, idx, tool_calls):
                        message["tool_calls"] = tool_calls
                    messages.append(message)

                elif conv["role"] == "tool":
                    if shallow_replace_enabled and metadata.get("auto_shallow_compacted"):
                        messages.append({
                            "role": "tool",
                            "content": AUTO_SHALLOW_PLACEHOLDER,
                            "tool_call_id": conv.get("tool_call_id", ""),
                            "name": conv.get("name", "")
                        })
                        replaced_tool_count += 1
                        continue
                    # Tool消息需要保留完整结构
                    images = conv.get("images") or metadata.get("images") or []
                    videos = conv.get("videos") or metadata.get("videos") or []
                    media_refs = conv.get("media_refs") or metadata.get("media_refs") or []
                    content_value = conv.get("content")
                    if isinstance(content_value, list):
                        content_payload = content_value
                    elif images or videos or media_refs:
                        content_payload = self.context_manager._build_content_with_images(
                            content_value,
                            images,
                            videos,
                            media_refs=media_refs,
                        )
                    else:
                        content_payload = content_value
                    message = {
                        "role": "tool",
                        "content": content_payload,
                        "tool_call_id": conv.get("tool_call_id", ""),
                        "name": conv.get("name", "")
                    }
                    messages.append(message)

                elif conv["role"] == "system" and metadata.get("sub_agent_notice"):
                    # 转换为用户消息，让模型能及时响应
                    messages.append({
                        "role": "user",
                        "content": conv["content"]
                    })
                else:
                    # User 或普通 System 消息
                    images = conv.get("images") or metadata.get("images") or []
                    videos = conv.get("videos") or metadata.get("videos") or []
                    media_refs = conv.get("media_refs") or metadata.get("media_refs") or []
                    content_payload = (
                        self.context_manager._build_content_with_images(
                            conv["content"],
                            images,
                            videos,
                            media_refs=media_refs,
                        )
                        if (images or videos or media_refs) else conv["content"]
                    )
                    # 调试：记录所有 system 消息
                    if conv["role"] == "system":
                        logger.info(f"[DEBUG build_messages] 添加 system 消息: content前50字={conv['content'][:50]}")
                    messages.append({
                        "role": conv["role"],
                        "content": content_payload
                    })

            # 当前用户输入已经在conversation中了，不需要重复添加
            if shallow_replace_enabled:
                print(f"[ContextCompression] build_messages 替换tool占位符: {replaced_tool_count} 条")
            return messages

    def _load_agents_md_content(self) -> Optional[str]:
        """加载工作区根目录的 AGENTS.md 文件内容。
        
        如果存在多个 AGENTS.md 文件，返回最新更新的那一个。
        """
        try:
            project_path = Path(self.project_path)
            if not project_path.exists():
                return None
            
            # 查找所有 AGENTS.md 文件
            agents_md_files = list(project_path.rglob("AGENTS.md"))
            if not agents_md_files:
                return None
            
            # 找到最新更新的文件
            latest_file = max(agents_md_files, key=lambda p: p.stat().st_mtime)
            
            # 读取文件内容
            content = latest_file.read_text(encoding='utf-8')
            return content.strip() if content else None
        except Exception as exc:
            logger.warning(f"[AGENTS.md] 读取失败: {exc}")
            return None

    def load_prompt(self, name: str) -> str:
            """加载提示模板"""
            prompt_file = Path(PROMPTS_DIR) / f"{name}.txt"
            if prompt_file.exists():
                with open(prompt_file, 'r', encoding='utf-8') as f:
                    return f.read()
            return "你是一个AI助手。"