434 lines
14 KiB
C#
434 lines
14 KiB
C#
using System;
|
|
using System.Collections.Generic;
|
|
using System.IO;
|
|
using System.Linq;
|
|
using System.Runtime.InteropServices;
|
|
using System.Security.Cryptography;
|
|
using System.Text;
|
|
using System.Text.Encodings.Web;
|
|
using System.Text.Json;
|
|
using System.Threading;
|
|
using System.Threading.Tasks;
|
|
|
|
namespace AxCopilot.Services.Agent;
|
|
|
|
public class CheckpointTool : IAgentTool
|
|
{
|
|
private class CheckpointManifest
|
|
{
|
|
public string Name { get; set; } = "";
|
|
|
|
public string CreatedAt { get; set; } = "";
|
|
|
|
public string WorkFolder { get; set; } = "";
|
|
|
|
public List<FileEntry> Files { get; set; } = new List<FileEntry>();
|
|
}
|
|
|
|
private class FileEntry
|
|
{
|
|
public string Path { get; set; } = "";
|
|
|
|
public string Hash { get; set; } = "";
|
|
}
|
|
|
|
private const int MaxCheckpoints = 10;
|
|
|
|
private const long MaxFileSize = 5242880L;
|
|
|
|
private static readonly HashSet<string> SkipDirs = new HashSet<string>(StringComparer.OrdinalIgnoreCase)
|
|
{
|
|
"bin", "obj", "node_modules", ".git", ".vs", ".ax", "packages", "Debug", "Release", "TestResults",
|
|
".idea", "__pycache__"
|
|
};
|
|
|
|
private static readonly HashSet<string> TextExtensions = new HashSet<string>(StringComparer.OrdinalIgnoreCase)
|
|
{
|
|
".cs", ".py", ".js", ".ts", ".java", ".cpp", ".c", ".h", ".hpp", ".xml",
|
|
".json", ".yaml", ".yml", ".md", ".txt", ".csv", ".html", ".htm", ".css", ".sql",
|
|
".sh", ".bat", ".ps1", ".config", ".ini", ".xaml", ".csproj", ".sln", ".props", ".targets",
|
|
".editorconfig", ".gitignore", ".tsx", ".jsx", ".vue", ".svelte", ".scss", ".less", ".toml", ".env",
|
|
".razor", ".proto", ".graphql", ".rs", ".go", ".rb", ".php", ".swift"
|
|
};
|
|
|
|
public string Name => "checkpoint";
|
|
|
|
public string Description => "Create, list, or restore file system checkpoints for undo/rollback. Checkpoints capture text files in the working folder as snapshots. - action=\"create\": Create a new checkpoint (name optional)\n- action=\"list\": List all checkpoints\n- action=\"restore\": Restore files from a checkpoint (id or name required, requires user approval)\n- action=\"delete\": Delete a checkpoint (id required)";
|
|
|
|
public ToolParameterSchema Parameters
|
|
{
|
|
get
|
|
{
|
|
ToolParameterSchema toolParameterSchema = new ToolParameterSchema();
|
|
Dictionary<string, ToolProperty> dictionary = new Dictionary<string, ToolProperty>();
|
|
ToolProperty obj = new ToolProperty
|
|
{
|
|
Type = "string",
|
|
Description = "create | list | restore | delete"
|
|
};
|
|
int num = 4;
|
|
List<string> list = new List<string>(num);
|
|
CollectionsMarshal.SetCount(list, num);
|
|
Span<string> span = CollectionsMarshal.AsSpan(list);
|
|
span[0] = "create";
|
|
span[1] = "list";
|
|
span[2] = "restore";
|
|
span[3] = "delete";
|
|
obj.Enum = list;
|
|
dictionary["action"] = obj;
|
|
dictionary["name"] = new ToolProperty
|
|
{
|
|
Type = "string",
|
|
Description = "Checkpoint name (for create/restore)"
|
|
};
|
|
dictionary["id"] = new ToolProperty
|
|
{
|
|
Type = "integer",
|
|
Description = "Checkpoint ID (for restore/delete)"
|
|
};
|
|
toolParameterSchema.Properties = dictionary;
|
|
num = 1;
|
|
List<string> list2 = new List<string>(num);
|
|
CollectionsMarshal.SetCount(list2, num);
|
|
CollectionsMarshal.AsSpan(list2)[0] = "action";
|
|
toolParameterSchema.Required = list2;
|
|
return toolParameterSchema;
|
|
}
|
|
}
|
|
|
|
public async Task<ToolResult> ExecuteAsync(JsonElement args, AgentContext context, CancellationToken ct = default(CancellationToken))
|
|
{
|
|
if (!args.TryGetProperty("action", out var actionEl))
|
|
{
|
|
return ToolResult.Fail("action이 필요합니다.");
|
|
}
|
|
string action = actionEl.GetString() ?? "";
|
|
if (string.IsNullOrEmpty(context.WorkFolder))
|
|
{
|
|
return ToolResult.Fail("작업 폴더가 설정되지 않았습니다.");
|
|
}
|
|
string checkpointDir = Path.Combine(context.WorkFolder, ".ax", "checkpoints");
|
|
if (1 == 0)
|
|
{
|
|
}
|
|
ToolResult result = action switch
|
|
{
|
|
"create" => await CreateCheckpoint(args, context, checkpointDir, ct),
|
|
"list" => ListCheckpoints(checkpointDir),
|
|
"restore" => await RestoreCheckpoint(args, context, checkpointDir, ct),
|
|
"delete" => DeleteCheckpoint(args, checkpointDir),
|
|
_ => ToolResult.Fail("알 수 없는 액션: " + action + ". create | list | restore | delete 중 선택하세요."),
|
|
};
|
|
if (1 == 0)
|
|
{
|
|
}
|
|
return result;
|
|
}
|
|
|
|
private async Task<ToolResult> CreateCheckpoint(JsonElement args, AgentContext context, string checkpointDir, CancellationToken ct)
|
|
{
|
|
string name = (args.TryGetProperty("name", out var n) ? (n.GetString() ?? "unnamed") : "unnamed");
|
|
name = string.Join("_", name.Split(Path.GetInvalidFileNameChars()));
|
|
string timestamp = DateTime.Now.ToString("yyyyMMdd_HHmmss");
|
|
string folderName = timestamp + "_" + name;
|
|
string cpDir = Path.Combine(checkpointDir, folderName);
|
|
try
|
|
{
|
|
Directory.CreateDirectory(cpDir);
|
|
List<string> files = CollectTextFiles(context.WorkFolder);
|
|
if (files.Count == 0)
|
|
{
|
|
return ToolResult.Fail("체크포인트할 텍스트 파일이 없습니다.");
|
|
}
|
|
CheckpointManifest manifest = new CheckpointManifest
|
|
{
|
|
Name = name,
|
|
CreatedAt = DateTime.Now.ToString("yyyy-MM-dd HH:mm:ss"),
|
|
WorkFolder = context.WorkFolder,
|
|
Files = new List<FileEntry>()
|
|
};
|
|
foreach (string file in files)
|
|
{
|
|
ct.ThrowIfCancellationRequested();
|
|
string relativePath = Path.GetRelativePath(context.WorkFolder, file);
|
|
string destPath = Path.Combine(cpDir, relativePath);
|
|
string destDir = Path.GetDirectoryName(destPath);
|
|
if (!string.IsNullOrEmpty(destDir) && !Directory.Exists(destDir))
|
|
{
|
|
Directory.CreateDirectory(destDir);
|
|
}
|
|
await CopyFileAsync(file, destPath, ct);
|
|
string hash = await ComputeHashAsync(file, ct);
|
|
manifest.Files.Add(new FileEntry
|
|
{
|
|
Path = relativePath,
|
|
Hash = hash
|
|
});
|
|
}
|
|
await File.WriteAllTextAsync(contents: JsonSerializer.Serialize(manifest, new JsonSerializerOptions
|
|
{
|
|
WriteIndented = true,
|
|
Encoder = JavaScriptEncoder.UnsafeRelaxedJsonEscaping
|
|
}), path: Path.Combine(cpDir, "manifest.json"), cancellationToken: ct);
|
|
CleanupOldCheckpoints(checkpointDir);
|
|
return ToolResult.Ok($"체크포인트 생성 완료: {folderName}\n파일 수: {manifest.Files.Count}");
|
|
}
|
|
catch (Exception ex)
|
|
{
|
|
return ToolResult.Fail("체크포인트 생성 오류: " + ex.Message);
|
|
}
|
|
}
|
|
|
|
private static ToolResult ListCheckpoints(string checkpointDir)
|
|
{
|
|
if (!Directory.Exists(checkpointDir))
|
|
{
|
|
return ToolResult.Ok("저장된 체크포인트가 없습니다.");
|
|
}
|
|
List<string> list = (from d in Directory.GetDirectories(checkpointDir)
|
|
orderby d descending
|
|
select d).ToList();
|
|
if (list.Count == 0)
|
|
{
|
|
return ToolResult.Ok("저장된 체크포인트가 없습니다.");
|
|
}
|
|
StringBuilder stringBuilder = new StringBuilder();
|
|
StringBuilder stringBuilder2 = stringBuilder;
|
|
StringBuilder stringBuilder3 = stringBuilder2;
|
|
StringBuilder.AppendInterpolatedStringHandler handler = new StringBuilder.AppendInterpolatedStringHandler(8, 1, stringBuilder2);
|
|
handler.AppendLiteral("체크포인트 ");
|
|
handler.AppendFormatted(list.Count);
|
|
handler.AppendLiteral("개:");
|
|
stringBuilder3.AppendLine(ref handler);
|
|
for (int num = 0; num < list.Count; num++)
|
|
{
|
|
string fileName = Path.GetFileName(list[num]);
|
|
string path = Path.Combine(list[num], "manifest.json");
|
|
int value = 0;
|
|
string text = "";
|
|
if (File.Exists(path))
|
|
{
|
|
try
|
|
{
|
|
string json = File.ReadAllText(path);
|
|
CheckpointManifest checkpointManifest = JsonSerializer.Deserialize<CheckpointManifest>(json);
|
|
value = checkpointManifest?.Files.Count ?? 0;
|
|
text = checkpointManifest?.CreatedAt ?? "";
|
|
}
|
|
catch
|
|
{
|
|
}
|
|
}
|
|
stringBuilder2 = stringBuilder;
|
|
StringBuilder stringBuilder4 = stringBuilder2;
|
|
handler = new StringBuilder.AppendInterpolatedStringHandler(12, 4, stringBuilder2);
|
|
handler.AppendLiteral(" [");
|
|
handler.AppendFormatted(num);
|
|
handler.AppendLiteral("] ");
|
|
handler.AppendFormatted(fileName);
|
|
handler.AppendLiteral(" — 파일 ");
|
|
handler.AppendFormatted(value);
|
|
handler.AppendLiteral("개");
|
|
handler.AppendFormatted(string.IsNullOrEmpty(text) ? "" : (", " + text));
|
|
stringBuilder4.AppendLine(ref handler);
|
|
}
|
|
return ToolResult.Ok(stringBuilder.ToString());
|
|
}
|
|
|
|
private async Task<ToolResult> RestoreCheckpoint(JsonElement args, AgentContext context, string checkpointDir, CancellationToken ct)
|
|
{
|
|
if (!Directory.Exists(checkpointDir))
|
|
{
|
|
return ToolResult.Fail("저장된 체크포인트가 없습니다.");
|
|
}
|
|
List<string> dirs = (from d in Directory.GetDirectories(checkpointDir)
|
|
orderby d descending
|
|
select d).ToList();
|
|
if (dirs.Count == 0)
|
|
{
|
|
return ToolResult.Fail("저장된 체크포인트가 없습니다.");
|
|
}
|
|
string targetDir = null;
|
|
if (args.TryGetProperty("id", out var idEl))
|
|
{
|
|
int parsed;
|
|
int id = ((idEl.ValueKind == JsonValueKind.Number) ? idEl.GetInt32() : (int.TryParse(idEl.GetString(), out parsed) ? parsed : (-1)));
|
|
if (id >= 0 && id < dirs.Count)
|
|
{
|
|
targetDir = dirs[id];
|
|
}
|
|
}
|
|
if (targetDir == null && args.TryGetProperty("name", out var nameEl))
|
|
{
|
|
string name = nameEl.GetString() ?? "";
|
|
targetDir = dirs.FirstOrDefault((string d) => Path.GetFileName(d).Contains(name, StringComparison.OrdinalIgnoreCase));
|
|
}
|
|
if (targetDir == null)
|
|
{
|
|
return ToolResult.Fail("체크포인트를 찾을 수 없습니다. id 또는 name을 확인하세요.");
|
|
}
|
|
if (context.AskPermission != null && !(await context.AskPermission("checkpoint_restore", "체크포인트 복원: " + Path.GetFileName(targetDir))))
|
|
{
|
|
return ToolResult.Ok("사용자가 복원을 거부했습니다.");
|
|
}
|
|
try
|
|
{
|
|
string manifestPath = Path.Combine(targetDir, "manifest.json");
|
|
if (!File.Exists(manifestPath))
|
|
{
|
|
return ToolResult.Fail("체크포인트 매니페스트를 찾을 수 없습니다.");
|
|
}
|
|
CheckpointManifest manifest = JsonSerializer.Deserialize<CheckpointManifest>(await File.ReadAllTextAsync(manifestPath, ct));
|
|
if (manifest == null)
|
|
{
|
|
return ToolResult.Fail("매니페스트 파싱 오류");
|
|
}
|
|
int restoredCount = 0;
|
|
foreach (FileEntry entry in manifest.Files)
|
|
{
|
|
ct.ThrowIfCancellationRequested();
|
|
string srcPath = Path.Combine(targetDir, entry.Path);
|
|
string destPath = Path.Combine(context.WorkFolder, entry.Path);
|
|
if (File.Exists(srcPath))
|
|
{
|
|
string destDir = Path.GetDirectoryName(destPath);
|
|
if (!string.IsNullOrEmpty(destDir) && !Directory.Exists(destDir))
|
|
{
|
|
Directory.CreateDirectory(destDir);
|
|
}
|
|
await CopyFileAsync(srcPath, destPath, ct);
|
|
restoredCount++;
|
|
}
|
|
}
|
|
return ToolResult.Ok($"체크포인트 복원 완료: {Path.GetFileName(targetDir)}\n복원 파일 수: {restoredCount}/{manifest.Files.Count}");
|
|
}
|
|
catch (Exception ex)
|
|
{
|
|
return ToolResult.Fail("체크포인트 복원 오류: " + ex.Message);
|
|
}
|
|
}
|
|
|
|
private static ToolResult DeleteCheckpoint(JsonElement args, string checkpointDir)
|
|
{
|
|
if (!Directory.Exists(checkpointDir))
|
|
{
|
|
return ToolResult.Fail("저장된 체크포인트가 없습니다.");
|
|
}
|
|
List<string> list = (from d in Directory.GetDirectories(checkpointDir)
|
|
orderby d descending
|
|
select d).ToList();
|
|
if (!args.TryGetProperty("id", out var value))
|
|
{
|
|
return ToolResult.Fail("삭제할 체크포인트 id가 필요합니다.");
|
|
}
|
|
int result;
|
|
int num = ((value.ValueKind == JsonValueKind.Number) ? value.GetInt32() : (int.TryParse(value.GetString(), out result) ? result : (-1)));
|
|
if (num < 0 || num >= list.Count)
|
|
{
|
|
return ToolResult.Fail($"잘못된 체크포인트 ID: {num}. 0~{list.Count - 1} 범위를 사용하세요.");
|
|
}
|
|
try
|
|
{
|
|
string path = list[num];
|
|
string fileName = Path.GetFileName(path);
|
|
Directory.Delete(path, recursive: true);
|
|
return ToolResult.Ok("체크포인트 삭제됨: " + fileName);
|
|
}
|
|
catch (Exception ex)
|
|
{
|
|
return ToolResult.Fail("체크포인트 삭제 오류: " + ex.Message);
|
|
}
|
|
}
|
|
|
|
private List<string> CollectTextFiles(string workFolder)
|
|
{
|
|
List<string> list = new List<string>();
|
|
CollectFilesRecursive(workFolder, workFolder, list);
|
|
return list;
|
|
}
|
|
|
|
private void CollectFilesRecursive(string dir, string rootDir, List<string> files)
|
|
{
|
|
string fileName = Path.GetFileName(dir);
|
|
if (dir != rootDir && SkipDirs.Contains(fileName))
|
|
{
|
|
return;
|
|
}
|
|
try
|
|
{
|
|
string[] files2 = Directory.GetFiles(dir);
|
|
foreach (string text in files2)
|
|
{
|
|
string extension = Path.GetExtension(text);
|
|
if (TextExtensions.Contains(extension))
|
|
{
|
|
FileInfo fileInfo = new FileInfo(text);
|
|
if (fileInfo.Length <= 5242880)
|
|
{
|
|
files.Add(text);
|
|
}
|
|
}
|
|
}
|
|
string[] directories = Directory.GetDirectories(dir);
|
|
foreach (string dir2 in directories)
|
|
{
|
|
CollectFilesRecursive(dir2, rootDir, files);
|
|
}
|
|
}
|
|
catch (UnauthorizedAccessException)
|
|
{
|
|
}
|
|
}
|
|
|
|
private static async Task CopyFileAsync(string src, string dest, CancellationToken ct)
|
|
{
|
|
byte[] buffer = new byte[81920];
|
|
await using FileStream srcStream = new FileStream(src, FileMode.Open, FileAccess.Read, FileShare.Read, buffer.Length, useAsync: true);
|
|
await using FileStream destStream = new FileStream(dest, FileMode.Create, FileAccess.Write, FileShare.None, buffer.Length, useAsync: true);
|
|
while (true)
|
|
{
|
|
int num;
|
|
int bytesRead = (num = await srcStream.ReadAsync(buffer, ct));
|
|
if (num <= 0)
|
|
{
|
|
break;
|
|
}
|
|
await destStream.WriteAsync(buffer.AsMemory(0, bytesRead), ct);
|
|
}
|
|
}
|
|
|
|
private static async Task<string> ComputeHashAsync(string filePath, CancellationToken ct)
|
|
{
|
|
using SHA256 sha256 = SHA256.Create();
|
|
string result;
|
|
await using (FileStream stream = new FileStream(filePath, FileMode.Open, FileAccess.Read, FileShare.Read, 81920, useAsync: true))
|
|
{
|
|
result = Convert.ToHexString(await sha256.ComputeHashAsync(stream, ct)).ToLowerInvariant();
|
|
}
|
|
return result;
|
|
}
|
|
|
|
private static void CleanupOldCheckpoints(string checkpointDir)
|
|
{
|
|
if (!Directory.Exists(checkpointDir))
|
|
{
|
|
return;
|
|
}
|
|
List<string> list = (from d in Directory.GetDirectories(checkpointDir)
|
|
orderby d
|
|
select d).ToList();
|
|
while (list.Count > 10)
|
|
{
|
|
try
|
|
{
|
|
Directory.Delete(list[0], recursive: true);
|
|
list.RemoveAt(0);
|
|
}
|
|
catch
|
|
{
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
}
|