agent_util.cpp

// agent_util.cpp —— Windows 基础设施（进程执行/文件读写/编码/崩溃诊断/重启自更新）
#include "agent_internal.h"

// 唯一实例：SSE 回调（net）与 chat()（core）共享
bool g_assistantStartEmitted = false;

// 取系统临时目录；失败时返回 false，调用方必须检查（否则会用未初始化缓冲构造 std::string）。
bool getTempDir(char *buf, size_t bufSize)
{
    if (bufSize == 0) return false;
    DWORD n = GetTempPathA(static_cast<DWORD>(bufSize), buf);
    if (n == 0 || n >= bufSize) { buf[0] = '\0'; return false; }
    return true;
}

// cmd 内置命令（dir/type 等）输出重定向到文件时按 OEM 代码页写字节（中文 Windows 为 936/GBK），
// 而 agent 侧一律按 UTF-8 处理 → 中文全乱码（实测 `dir /b` 返回的是 GBK 字节 b2e2cad4）。
// 归一化：已是合法 UTF-8 的原样保留（git/python 等自身输出 UTF-8），否则按 OEM 代码页转 UTF-8。
std::string normalizeCommandOutputEncoding(const std::string &bytes)
{
    if (bytes.empty() || isValidUtf8(bytes)) return bytes;
    const UINT cp = GetOEMCP();
    const int wlen = MultiByteToWideChar(cp, 0, bytes.c_str(), static_cast<int>(bytes.size()), nullptr, 0);
    if (wlen <= 0) return bytes;
    std::wstring wide(static_cast<size_t>(wlen), L'\0');
    MultiByteToWideChar(cp, 0, bytes.c_str(), static_cast<int>(bytes.size()), &wide[0], wlen);
    return wstringToUtf8(wide);
}

std::string executeShellCommand(const std::string &command, DWORD timeoutMs)
{
    if (command.empty())
        return "(无输出)";

    char tmpDir[MAX_PATH];
    if (!getTempDir(tmpDir, MAX_PATH))
        return "Tool Error: 无法获取系统临时目录";
    char outPath[MAX_PATH];
    if (GetTempFileNameA(tmpDir, "ycd", 0, outPath) == 0)
        return "Tool Error: 无法创建临时输出文件";

    SECURITY_ATTRIBUTES sa;
    sa.nLength = sizeof(sa);
    sa.bInheritHandle = TRUE;
    sa.lpSecurityDescriptor = NULL;

    // 命令里可能含中文路径（模型给的路径是 UTF-8），而 A 版 API 按 ANSI 解释命令行 →
    // 中文路径到 cmd 手里就成了别的字符。统一走 W 版：命令、输出文件句柄都用宽字符。
    const std::wstring wideOutPath = utf8ToWide(outPath);
    HANDLE hOut = CreateFileW(wideOutPath.c_str(), GENERIC_WRITE, FILE_SHARE_READ, &sa, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
    HANDLE hIn = CreateFileW(L"NUL", GENERIC_READ, FILE_SHARE_READ, &sa, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
    if (hOut == INVALID_HANDLE_VALUE || hIn == INVALID_HANDLE_VALUE)
    {
        if (hOut != INVALID_HANDLE_VALUE) CloseHandle(hOut);
        if (hIn != INVALID_HANDLE_VALUE) CloseHandle(hIn);
        DeleteFileW(wideOutPath.c_str());
        return "Tool Error: 无法创建命令输出文件";
    }

    STARTUPINFOW si;
    ZeroMemory(&si, sizeof(si));
    si.cb = sizeof(si);
    si.dwFlags = STARTF_USESTDHANDLES;
    si.hStdOutput = hOut;
    si.hStdError = hOut;
    si.hStdInput = hIn;

    PROCESS_INFORMATION pi;
    ZeroMemory(&pi, sizeof(pi));
    std::wstring cmdLine = L"cmd.exe /c " + utf8ToWide(command);

    // Job Object：关闭时终止整个进程树（cmd 及其子进程），避免超时后遗留子进程占用控制台
    HANDLE job = CreateJobObjectW(NULL, NULL);
    if (job)
    {
        JOBOBJECT_EXTENDED_LIMIT_INFORMATION ji;
        ZeroMemory(&ji, sizeof(ji));
        ji.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
        SetInformationJobObject(job, JobObjectExtendedLimitInformation, &ji, sizeof(ji));
    }

    BOOL created = CreateProcessW(NULL, &cmdLine[0], NULL, NULL, TRUE, CREATE_NO_WINDOW, NULL, NULL, &si, &pi);
    CloseHandle(hOut);
    CloseHandle(hIn);
    if (!created)
    {
        if (job) CloseHandle(job);
        DeleteFileW(wideOutPath.c_str());
        return "Tool Error: 无法启动命令进程";
    }
    if (job)
        AssignProcessToJobObject(job, pi.hProcess);

    DWORD waitResult = WaitForSingleObject(pi.hProcess, timeoutMs);
    if (waitResult == WAIT_TIMEOUT)
    {
        if (job)
            CloseHandle(job); // KILL_ON_JOB_CLOSE：终止整个进程树
        TerminateProcess(pi.hProcess, 1);
        WaitForSingleObject(pi.hProcess, 5000);
        CloseHandle(pi.hProcess);
        CloseHandle(pi.hThread);
        DeleteFileW(wideOutPath.c_str());
        return "Tool Error: 命令执行超时（已终止），请改用更快或非交互的命令。";
    }
    if (job)
        CloseHandle(job);
    CloseHandle(pi.hProcess);
    CloseHandle(pi.hThread);

    std::string result = normalizeCommandOutputEncoding(readFile(outPath));
    DeleteFileW(wideOutPath.c_str());
    // readFile 失败只会返回 "Tool Error: ..." 前缀；这里只认该前缀，
    // 命令自身输出以 "FAIL" 开头时不会被误清空。
    if (result.rfind("Tool Error", 0) == 0)
        result.clear();
    if (result.empty())
        return "(无输出)";
    if (result.length() > 8000)
        result = result.substr(0, 8000) + "\n... (截断)";
    return result;
}


// 路径统一走宽字符 API。
// 工具参数里的路径是 UTF-8（模型给的就是 UTF-8），而 Windows 的 *A 版 API 与
// std::ifstream 的窄字符构造按 **ANSI 代码页**解释它 —— 中文路径会变成别的字符，
// 于是"文件明明存在却报无法访问"（实测：read_file 读 样例文件.txt 失败）。
// 这里统一 utf8ToWide 后用 *W 版 API / 宽字符流。
std::string readFile(const std::string &filepath)
{
    const std::wstring widePath = utf8ToWide(filepath);
    struct _stat64 st;
    if (_wstat64(widePath.c_str(), &st) != 0) return "Tool Error: 无法访问文件 " + filepath;
    if (st.st_size > 10 * 1024 * 1024) return "Tool Error: 文件过大";
    std::ifstream file(widePath.c_str(), std::ios::in | std::ios::binary);
    if (!file.is_open()) return "Tool Error: 无法打开文件 " + filepath;
    std::string content((std::istreambuf_iterator<char>(file)), std::istreambuf_iterator<char>());
    file.close();
    if (content.empty()) return "(文件为空)";
    if (content.length() > 50000)
        content = content.substr(0, 50000) + "\n" + READ_TRUNCATION_MARKER +
                  " 文件超过 50000 字符，以上内容不完整。禁止把本结果写回 write_file（会截断源文件）；"
                  "请分段读取分段写入，或用 execute_command 调增量修改工具。";
    return content;
}

bool writeFile(const std::string &filepath, const std::string &content)
{
    std::ofstream file(utf8ToWide(filepath).c_str(), std::ios::out | std::ios::binary);
    if (!file.is_open()) return false;
    file << content;
    file.close();
    return !file.fail();
}

// 原子写：先写 .tmp 再整文件替换，避免进程中断把目标写成半截（会话文件尤其关键）。
bool writeFileAtomic(const std::string &filepath, const std::string &content)
{
    std::string tmp = filepath + ".tmp";
    if (!writeFile(tmp, content)) return false;
    if (MoveFileExW(utf8ToWide(tmp).c_str(), utf8ToWide(filepath).c_str(),
                    MOVEFILE_REPLACE_EXISTING | MOVEFILE_WRITE_THROUGH))
        return true;
    DeleteFileW(utf8ToWide(tmp).c_str());
    return false;
}

// 整份读取、不截断（供会话文件这类"必须完整"的场景）；超过 maxBytes 视为不可用。
std::string readFileWhole(const std::string &filepath, size_t maxBytes)
{
    const std::wstring widePath = utf8ToWide(filepath);
    struct _stat64 st;
    if (_wstat64(widePath.c_str(), &st) != 0) return "";
    if (st.st_size <= 0 || static_cast<size_t>(st.st_size) > maxBytes) return "";
    std::ifstream file(widePath.c_str(), std::ios::in | std::ios::binary);
    if (!file.is_open()) return "";
    std::string content((std::istreambuf_iterator<char>(file)), std::istreambuf_iterator<char>());
    return content;
}

std::string searchFiles(const std::string &pattern, const std::string &directory)
{
    if (hasShellMetacharacter(pattern) || hasShellMetacharacter(directory))
        return "Tool Error: pattern 或 directory 含非法字符，已拒绝执行（避免命令注入）。";
    std::string cmd = "dir /s /b \"" + directory + "\\" + pattern + "\" 2>nul";
    std::string result = executeShellCommand(cmd);
    if (result == "(无输出)" || result.empty()) return "未找到: " + pattern;
    return result;
}

std::string searchContent(const std::string &text, const std::string &filePattern, const std::string &directory)
{
    if (hasShellMetacharacter(text) || hasShellMetacharacter(filePattern) || hasShellMetacharacter(directory))
        return "Tool Error: text、filePattern 或 directory 含非法字符，已拒绝执行（避免命令注入）。";
    std::string cmd = "findstr /s /n /i /c:\"" + text + "\" \"" + directory + "\\" + filePattern + "\" 2>nul";
    std::string result = executeShellCommand(cmd);
    if (result == "(无输出)" || result.empty()) return "未找到包含 \"" + text + "\" 的内容";
    return result;
}

std::string createDirectory(const std::string &path)
{
    if (CreateDirectoryW(utf8ToWide(path).c_str(), NULL)) return "OK 目录创建: " + path;
    if (GetLastError() == ERROR_ALREADY_EXISTS) return "目录已存在: " + path;
    return "FAIL 目录创建: " + path;
}

std::string deleteFile(const std::string &path)
{
    const std::wstring widePath = utf8ToWide(path);
    DWORD attrs = GetFileAttributesW(widePath.c_str());
    if (attrs == INVALID_FILE_ATTRIBUTES) return "FAIL: 不存在 " + path;
    if (attrs & FILE_ATTRIBUTE_DIRECTORY)
    {
        if (hasShellMetacharacter(path))
            return "FAIL 目录删除: 路径含非法字符，已拒绝执行（避免命令注入）";
        std::string cmd = "rmdir /s /q \"" + path + "\" 2>nul";
        // system() 在 Windows 上按 ANSI 传命令 → 中文路径会失效；_wsystem 走宽字符
        _wsystem(utf8ToWide(cmd).c_str());
        if (GetFileAttributesW(widePath.c_str()) == INVALID_FILE_ATTRIBUTES) return "OK 目录删除: " + path;
        return "FAIL 目录删除: " + path;
    }
    if (DeleteFileW(widePath.c_str())) return "OK 文件删除: " + path;
    return "FAIL 文件删除: " + path;
}

std::string moveFile(const std::string &source, const std::string &destination)
{
    if (MoveFileW(utf8ToWide(source).c_str(), utf8ToWide(destination).c_str()))
        return "OK 移动: " + source + " -> " + destination;
    return "FAIL 移动: " + source;
}

std::string getFileInfo(const std::string &path)
{
    struct _stat64 st;
    if (_wstat64(utf8ToWide(path).c_str(), &st) != 0) return "FAIL 获取信息: " + path;
    std::ostringstream info;
    info << "文件: " << path << "\n";
    info << "大小: " << st.st_size << " 字节";
    double size = (double)st.st_size;
    const char *units[] = {"B", "KB", "MB", "GB"};
    int unitIdx = 0;
    while (size >= 1024.0 && unitIdx < 3) { size /= 1024.0; unitIdx++; }
    info << " (" << std::fixed << std::setprecision(2) << size << " " << units[unitIdx] << ")\n";
    char timeBuf[64];
    struct tm timeinfo;
    localtime_s(&timeinfo, &st.st_mtime);
    strftime(timeBuf, sizeof(timeBuf), "%Y-%m-%d %H:%M:%S", &timeinfo);
    info << "修改: " << timeBuf << "\n";
    info << "类型: " << ((st.st_mode & _S_IFDIR) ? "目录" : "文件");
    return info.str();
}

// ============================================================
// 路径助手：本程序自身路径一律取宽字符再转 UTF-8
// （GetModuleFileNameA 给的是 ANSI 字节，中文目录会直接变乱码）
// ============================================================
std::string modulePathUtf8()
{
    wchar_t buffer[MAX_PATH];
    DWORD len = GetModuleFileNameW(NULL, buffer, MAX_PATH);
    if (len == 0 || len >= MAX_PATH) return "";
    return wstringToUtf8(buffer);
}

std::string moduleDirUtf8()
{
    const std::string full = modulePathUtf8();
    if (full.empty()) return "";
    const size_t pos = full.find_last_of("\\/");
    return pos == std::string::npos ? std::string(".") : full.substr(0, pos);
}

// ============================================================
// 崩溃诊断：把未处理异常的类型与调用栈写入 agent_crash.log
// ============================================================
// 崩溃日志固定写到 agent.exe 所在目录：原先用相对路径，取决于 CWD，客户端托管时会落到工作区里
std::string crashLogPath()
{
    const std::string dir = moduleDirUtf8();
    if (!dir.empty())
        return dir + "\\agent_crash.log";
    return "agent_crash.log";
}

void logCrashStack(FILE *f)
{
    static bool symInit = false;
    if (!symInit)
        symInit = SymInitialize(GetCurrentProcess(), NULL, TRUE) == TRUE;

    void *stack[32];
    USHORT frames = CaptureStackBackTrace(0, 32, stack, NULL);
    for (USHORT i = 0; i < frames; i++)
    {
        DWORD64 addr = (DWORD64)stack[i];
        if (symInit)
        {
            char buf[sizeof(SYMBOL_INFO) + 256] = {0};
            SYMBOL_INFO *sym = (SYMBOL_INFO *)buf;
            sym->SizeOfStruct = sizeof(SYMBOL_INFO);
            sym->MaxNameLen = 255;
            DWORD64 disp = 0;
            if (SymFromAddr(GetCurrentProcess(), addr, &disp, sym))
                fprintf(f, "  [%u] %p %s+0x%llX\n", i, stack[i], sym->Name, (unsigned long long)disp);
            else
                fprintf(f, "  [%u] %p\n", i, stack[i]);
        }
        else
        {
            fprintf(f, "  [%u] %p\n", i, stack[i]);
        }
    }
}

void terminateHandler()
{
    FILE *f = fopen(crashLogPath().c_str(), "a");
    if (f)
    {
        fprintf(f, "\n=== std::terminate %s ===\n", __TIMESTAMP__);
        if (std::current_exception())
        {
            try { std::rethrow_exception(std::current_exception()); }
            catch (const std::exception &e) { fprintf(f, "exception: %s\n", e.what()); }
            catch (...) { fprintf(f, "exception: (未知类型)\n"); }
        }
        else
        {
            fprintf(f, "exception: (无活动异常)\n");
        }
        logCrashStack(f);
        fclose(f);
    }
    std::abort();
}

LONG WINAPI CrashVectoredHandler(PEXCEPTION_POINTERS info)
{
    const DWORD code = info->ExceptionRecord->ExceptionCode;
    // 这些不是崩溃，日常流程会大量触发；原先不过滤，导致 agent_crash.log 里几乎全是噪声：
    //   0xE06D7363 MSVC C++ 异常（throw） 0x40010006 OutputDebugString
    //   0x406D1388 SetThreadName        0x4001000A 调试打印
    if (code == 0xE06D7363 || code == 0x40010006 || code == 0x406D1388 || code == 0x4001000A)
        return EXCEPTION_CONTINUE_SEARCH;

    FILE *f = fopen(crashLogPath().c_str(), "a");
    if (f)
    {
        fprintf(f, "\n=== vectored exception 0x%08lX at %p %s ===\n",
                (unsigned long)code,
                info->ExceptionRecord->ExceptionAddress, __TIMESTAMP__);
        logCrashStack(f);
        fclose(f);
    }
    return EXCEPTION_CONTINUE_SEARCH;
}

// ============================================================
// 重启 Agent — 发送信号给 YCode 客户端来接管重启
// ============================================================
// 新的重启策略：
//   agent.exe 不再自己创建批处理脚本重启，而是向 YCode 客户端
//   发送特殊信号 "SIGNAL:RESTART_AGENT"，然后干净退出。
//   YCode 的 AgentManager 检测到该信号后，会杀掉旧进程并启动新的 agent.exe，
//   这样新的 agent.exe 仍然由 YCode 管理，聊天功能不会断连。
//   如果在独立模式运行（无 YCode），则回退到旧的批处理自重启方式。
// ============================================================
std::string restartAgent(bool standalone)
{
    const char *managed = std::getenv("YCODE_MANAGED");
    bool managedByYCode = !standalone && managed && std::string(managed) == "1";

    if (managedByYCode)
    {
        // 由 YCode 客户端托管：发送重启信号，让 YCode 来重启 agent
        std::cout << "\n  [Agent 请求 YCode 客户端执行重启...]" << std::endl;
        std::cout << SIGNAL_RESTART_AGENT << std::endl;
        // 保存会话后再退出
        exit(0);
        return "OK 重启信号已发送";  // 不会执行到这里
    }

    // 独立模式：回退到批处理自重启
    const std::string fullPath = modulePathUtf8();
    if (fullPath.empty())
        return "FAIL: 无法获取当前程序路径";

    size_t lastSlash = fullPath.find_last_of("\\/");
    std::string exeDir = (lastSlash != std::string::npos) ? fullPath.substr(0, lastSlash) : ".";
    std::string exeName = (lastSlash != std::string::npos) ? fullPath.substr(lastSlash + 1) : fullPath;

    std::string batPath = exeDir + "\\_agent_restart.bat";
    std::ofstream bat(utf8ToWide(batPath).c_str());
    if (!bat.is_open())
        return "FAIL: 无法创建重启脚本 " + batPath;

    bat << "@echo off\n";
    bat << "chcp 65001 >nul\n";
    bat << "echo 正在重启 YCode Agent...\n";
    bat << "ping 127.0.0.1 -n 2 >nul\n";
    bat << "cd /d \"" << exeDir << "\"\n";
    // 不再在 cmd /k 里重复 cd：子进程继承上面的工作目录，
    // 原先 "cd /d " << exeDir << " 未加引号，安装路径含空格时自重启会失败。
    bat << "start \"YCode Agent v2.0\" cmd /k \"" << exeName << "\"\n";
    bat << "ping 127.0.0.1 -n 1 >nul\n";
    bat << "del \"%~f0\" >nul 2>&1\n";
    bat.close();

    std::cout << "\n  [Agent 即将重启... 请稍候]" << std::endl;
    ShellExecuteW(NULL, L"open", utf8ToWide(batPath).c_str(), NULL, utf8ToWide(exeDir).c_str(), SW_HIDE);
    exit(0);
    return "OK 重启中...";  // 实际上不会执行到这里
}

std::string requestYCodeRebuildAndRestart()
{
    const char *managed = std::getenv("YCODE_MANAGED");
    if (managed && std::string(managed) == "1")
    {
        std::cout << "\n  [Agent 请求 YCode 执行完整自更新：重建 Agent、客户端并更新快捷方式...]" << std::endl;
        std::cout << SIGNAL_REBUILD_RESTART_YCODE << std::endl;
        exit(0);
        return "OK YCode 自更新信号已发送";
    }

    // GetModuleFileNameA 返回的是 ANSI 字节，而本项目内部一律 UTF-8：
    // 直接当 UTF-8 用会把中文路径变成乱码 → 用 W 版再转 UTF-8。
    wchar_t exePathW[MAX_PATH];
    DWORD len = GetModuleFileNameW(NULL, exePathW, MAX_PATH);
    if (len == 0 || len >= MAX_PATH)
        return "FAIL: 无法获取当前程序路径";

    std::string fullPath = wstringToUtf8(exePathW);
    size_t lastSlash = fullPath.find_last_of("\\/");
    std::string exeDir = (lastSlash != std::string::npos) ? fullPath.substr(0, lastSlash) : ".";
    std::string scriptPath = exeDir + "\\ycode_self_update.bat";

    const std::wstring wideScriptPath = utf8ToWide(scriptPath);
    if (GetFileAttributesW(wideScriptPath.c_str()) == INVALID_FILE_ATTRIBUTES)
        return "FAIL: 找不到自更新脚本 " + scriptPath;

    ShellExecuteW(NULL, L"open", wideScriptPath.c_str(), NULL, utf8ToWide(exeDir).c_str(), SW_HIDE);
    exit(0);
    return "OK YCode 自更新脚本已启动";
}
std::string wstringToUtf8(const std::wstring &w)
{
    if (w.empty()) return "";
    int size = WideCharToMultiByte(CP_UTF8, 0, w.c_str(), (int)w.size(), nullptr, 0, nullptr, nullptr);
    std::string out(size, 0);
    WideCharToMultiByte(CP_UTF8, 0, w.c_str(), (int)w.size(), &out[0], size, nullptr, nullptr);
    return out;
}

std::wstring utf8ToWide(const std::string &s)
{
    if (s.empty()) return L"";
    int size = MultiByteToWideChar(CP_UTF8, 0, s.c_str(), (int)s.size(), nullptr, 0);
    std::wstring out(size, 0);
    MultiByteToWideChar(CP_UTF8, 0, s.c_str(), (int)s.size(), &out[0], size);
    return out;
}
// agent.exe 所在目录（正常安装即 YCode 源码根），小写 + 正斜杠；失败返回空串。
// 用于在"工作区不是 YCode 根目录"时也能识别自修改路径。
std::string ycodeInstallDirLower()
{
    std::string dir = moduleDirUtf8();
    if (dir.empty()) return "";
    std::replace(dir.begin(), dir.end(), '\\', '/');
    // UTF-8 字节里 >127 的字节不会被 tolower 改动，所以中文目录名保持原样、只折 ASCII 部分
    std::transform(dir.begin(), dir.end(), dir.begin(),
                   [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
    return dir;
}

// ============================================================
// 构建单个工具定义的辅助函数
// ============================================================
json makeTool(const std::string &name, const std::string &desc, const json &properties, const json &required)
{
    json tool;
    tool["type"] = "function";
    tool["function"]["name"] = name;
    tool["function"]["description"] = desc;
    tool["function"]["parameters"]["type"] = "object";
    tool["function"]["parameters"]["properties"] = properties;
    tool["function"]["parameters"]["required"] = required;
    return tool;
}

json makeProp(const std::string &name, const std::string &type, const std::string &desc)
{
    json prop;
    prop["type"] = type;
    prop["description"] = desc;
    return prop;
}