← Back to Files
MemoryRebuildService.cs
using System.Diagnostics;
using System.Text;
using System.Text.RegularExpressions;
namespace KaiChat.Services;
public interface IMemoryRebuildService
{
Task<IReadOnlyList<MemoryRebuildSourceInfo>> ListSourcesAsync(int chunkTokens = 600_000);
MemoryRebuildJobStatus GetStatus();
Task<MemoryRebuildJobStatus> StartAsync(MemoryRebuildStartOptions options);
Task<MemoryRebuildJobStatus> CancelAsync();
Task<IReadOnlyList<MemoryRebuildFileInfo>> ListLatestFilesAsync();
Task<string> ReadLatestFileAsync(string relativePath);
}
public record MemoryRebuildSourceInfo(
string Key,
string DisplayName,
string FileName,
int Order,
int? PartNumber,
int MessageBlocks,
int ChunkCount,
int EstimatedTokens,
bool SnapshotExists,
bool IsNext);
public record MemoryRebuildStartOptions(
string? SourceKey = null,
bool ProcessRemaining = false,
bool Apply = false,
int ChunkTokens = 600_000,
int MaxOutputTokens = 384_000,
int MaxChunks = 0,
bool Resume = true,
bool CommitAfterSource = true,
bool RetentionAudit = true,
bool ChangeAudit = true);
public record MemoryRebuildJobStatus(
string? Id,
bool Running,
bool Apply,
string? SourceKey,
string? SourceName,
DateTime? StartedAtUtc,
DateTime? CompletedAtUtc,
int? ExitCode,
string State,
string Log,
string OutputRoot,
string LatestDir,
string SnapshotsDir);
public record MemoryRebuildFileInfo(
string Path,
long Size,
DateTime UpdatedAt);
public class MemoryRebuildService : IMemoryRebuildService
{
private const double TokensPerCharacter = 0.3;
private readonly string _baseFolder;
private readonly string _sourceDir;
private readonly string _outputRoot;
private readonly string _latestDir;
private readonly string _snapshotsDir;
private readonly string _toolProject;
private readonly ILogger<MemoryRebuildService> _logger;
private readonly object _sync = new();
private readonly StringBuilder _log = new();
private Process? _process;
private MemoryRebuildJobStatus _status;
private static readonly string[] OrderedSourceFiles =
[
"001_The Chat Part 1.md",
"002_The Chat Part 2.md",
"003_The Chat Part 3.md",
"004_The Chat Part 4.md",
"005_The Chat Part 5.md",
"006_The Chat Part 6.md",
"007_The Chat Part 7.md",
"008_The Chat Part 8.md",
"009_The Chat Part 9.md",
"010_The Chat Part 10.md",
"T01_Test Chat 1.md",
"T02_Test Chat 2.md",
"011_The Chat Part 11.md",
"012_The Chat Part 12.md",
"013_The Chat Part 13.md",
"014_The Chat Part 14.md",
"015_The Chat Part 15.md",
"016_The Chat Part 16.md",
"017_The Chat Part 17.md",
"018_The Chat Part 18.md",
"019_The Chat Part 19.md",
"020_The Chat Part 20.md",
"021_The Chat Part 21.md",
"022_The Chat Part 22.md",
"023_The Chat Part 23.md",
"024_The Chat Part 24.md",
"025_The Chat Part 25.md",
"026_The Chat Part 26.md",
"027_The Chat Part 27.md",
"028_The Chat Part 28.md",
"029_The Chat Part 29.md",
"030_The Chat Part 30.md",
"031_The Chat Part 31.md"
];
public MemoryRebuildService(IConfiguration config, ILogger<MemoryRebuildService> logger)
{
_baseFolder = config.GetValue<string>("KaiChat:BaseFolder") ?? ".";
_sourceDir = Path.Combine(_baseFolder, "Extracted The Chat Project");
_outputRoot = Path.Combine(_baseFolder, "KaiChat", "Memory_Rebuild");
_latestDir = Path.Combine(_outputRoot, "latest");
_snapshotsDir = Path.Combine(_outputRoot, "snapshots");
_toolProject = Path.Combine(_baseFolder, "MemoryRebuildTool", "MemoryRebuildTool.csproj");
_logger = logger;
_status = EmptyStatus();
}
public async Task<IReadOnlyList<MemoryRebuildSourceInfo>> ListSourcesAsync(int chunkTokens = 600_000)
{
var maxChars = Math.Max(4_000, (int)(chunkTokens / TokensPerCharacter));
var sources = new List<MemoryRebuildSourceInfo>();
foreach (var (fileName, index) in OrderedSourceFiles.Select((file, idx) => (file, idx + 1)))
{
var path = Path.Combine(_sourceDir, fileName);
if (!File.Exists(path))
continue;
var source = SourceMetadata.FromFile(path, index);
if (source == null)
continue;
var text = await File.ReadAllTextAsync(path);
var blocks = ExtractMessageBlocks(text);
var chunks = BuildChunkCount(blocks, maxChars);
sources.Add(new MemoryRebuildSourceInfo(
source.Key,
source.DisplayName,
fileName,
source.Order,
source.PartNumber,
blocks.Count,
chunks,
EstimateTokens(text),
Directory.Exists(Path.Combine(_snapshotsDir, source.Key)),
false));
}
var nextKey = sources.FirstOrDefault(s => !s.SnapshotExists)?.Key;
return sources
.Select(s => s with { IsNext = string.Equals(s.Key, nextKey, StringComparison.OrdinalIgnoreCase) })
.ToList();
}
public MemoryRebuildJobStatus GetStatus()
{
lock (_sync)
return _status with { Log = _log.ToString() };
}
public async Task<MemoryRebuildJobStatus> StartAsync(MemoryRebuildStartOptions options)
{
var sources = await ListSourcesAsync(options.ChunkTokens);
var remainingSources = sources.Where(s => !s.SnapshotExists).ToList();
var source = options.ProcessRemaining
? null
: sources.FirstOrDefault(s => string.Equals(s.Key, options.SourceKey, StringComparison.OrdinalIgnoreCase))
?? throw new ArgumentException($"Unknown memory rebuild source: {options.SourceKey}", nameof(options));
if (options.ProcessRemaining && remainingSources.Count == 0)
throw new InvalidOperationException("There are no remaining memory rebuild sources to process.");
lock (_sync)
{
if (_process is { HasExited: false })
throw new InvalidOperationException("A memory rebuild job is already running.");
_log.Clear();
var id = DateTime.UtcNow.ToString("yyyyMMdd-HHmmss");
_status = new MemoryRebuildJobStatus(
id,
Running: true,
options.Apply,
options.ProcessRemaining ? "remaining" : source!.Key,
options.ProcessRemaining ? $"Remaining sources ({remainingSources.Count})" : source!.DisplayName,
DateTime.UtcNow,
null,
null,
options.Apply ? "running apply" : "running dry-run",
"",
_outputRoot,
_latestDir,
_snapshotsDir);
var startInfo = new ProcessStartInfo
{
FileName = "dotnet",
WorkingDirectory = _baseFolder,
RedirectStandardOutput = true,
RedirectStandardError = true,
UseShellExecute = false,
CreateNoWindow = true
};
startInfo.ArgumentList.Add("run");
startInfo.ArgumentList.Add("--project");
startInfo.ArgumentList.Add(_toolProject);
startInfo.ArgumentList.Add("--");
if (options.Apply)
startInfo.ArgumentList.Add("--apply");
if (!options.Resume)
startInfo.ArgumentList.Add("--no-resume");
if (options.ProcessRemaining)
{
startInfo.ArgumentList.Add("--skip-snapshotted");
}
else
{
startInfo.ArgumentList.Add("--source-key");
startInfo.ArgumentList.Add(source!.Key);
}
if (options.Apply && options.CommitAfterSource)
startInfo.ArgumentList.Add("--commit-after-source");
if (!options.RetentionAudit)
startInfo.ArgumentList.Add("--no-retention-audit");
if (!options.ChangeAudit)
startInfo.ArgumentList.Add("--no-change-audit");
startInfo.ArgumentList.Add("--chunk-tokens");
startInfo.ArgumentList.Add(options.ChunkTokens.ToString());
startInfo.ArgumentList.Add("--max-output-tokens");
startInfo.ArgumentList.Add(options.MaxOutputTokens.ToString());
if (options.MaxChunks > 0)
{
startInfo.ArgumentList.Add("--max-chunks");
startInfo.ArgumentList.Add(options.MaxChunks.ToString());
}
AppendLog(options.ProcessRemaining
? $"Starting {(options.Apply ? "apply" : "dry-run")} for remaining sources ({remainingSources.Count})"
: $"Starting {(options.Apply ? "apply" : "dry-run")} for {source!.DisplayName} ({source.Key})");
AppendLog($"Output root: {_outputRoot}");
_process = new Process { StartInfo = startInfo, EnableRaisingEvents = true };
_process.OutputDataReceived += (_, e) => { if (e.Data != null) AppendLog(e.Data); };
_process.ErrorDataReceived += (_, e) => { if (e.Data != null) AppendLog("[err] " + e.Data); };
_process.Exited += (_, _) => CompleteCurrentProcess();
if (!_process.Start())
throw new InvalidOperationException("Failed to start memory rebuild process.");
_process.BeginOutputReadLine();
_process.BeginErrorReadLine();
return _status with { Log = _log.ToString() };
}
}
public Task<MemoryRebuildJobStatus> CancelAsync()
{
lock (_sync)
{
if (_process is { HasExited: false })
{
AppendLog("Cancelling memory rebuild process...");
_process.Kill(entireProcessTree: true);
_status = _status with
{
Running = false,
CompletedAtUtc = DateTime.UtcNow,
State = "cancelled"
};
}
return Task.FromResult(_status with { Log = _log.ToString() });
}
}
public Task<IReadOnlyList<MemoryRebuildFileInfo>> ListLatestFilesAsync()
{
if (!Directory.Exists(_latestDir))
return Task.FromResult<IReadOnlyList<MemoryRebuildFileInfo>>([]);
var files = Directory.GetFiles(_latestDir, "*.md", SearchOption.AllDirectories)
.Where(path =>
{
var rel = Path.GetRelativePath(_latestDir, path).Replace('\\', '/');
return !rel.Split('/', StringSplitOptions.RemoveEmptyEntries).Any(part => part.StartsWith("_", StringComparison.Ordinal));
})
.Select(path =>
{
var info = new FileInfo(path);
return new MemoryRebuildFileInfo(
Path.GetRelativePath(_latestDir, path).Replace('\\', '/'),
info.Length,
info.LastWriteTimeUtc);
})
.OrderBy(f => LoadOrder(f.Path))
.ThenBy(f => f.Path, StringComparer.OrdinalIgnoreCase)
.ToList();
return Task.FromResult<IReadOnlyList<MemoryRebuildFileInfo>>(files);
}
public async Task<string> ReadLatestFileAsync(string relativePath)
{
var path = ResolveLatestPath(relativePath);
if (!File.Exists(path))
throw new FileNotFoundException("Rebuilt memory file not found.", relativePath);
return await File.ReadAllTextAsync(path);
}
private MemoryRebuildJobStatus EmptyStatus() => new(
null,
Running: false,
Apply: false,
SourceKey: null,
SourceName: null,
StartedAtUtc: null,
CompletedAtUtc: null,
ExitCode: null,
State: "idle",
Log: "",
_outputRoot,
_latestDir,
_snapshotsDir);
private void CompleteCurrentProcess()
{
lock (_sync)
{
if (_process == null)
return;
var exitCode = _process.ExitCode;
AppendLog($"Process exited with code {exitCode}");
_status = _status with
{
Running = false,
CompletedAtUtc = DateTime.UtcNow,
ExitCode = exitCode,
State = exitCode == 0 ? "completed" : "failed"
};
_process.Dispose();
_process = null;
}
}
private void AppendLog(string line)
{
lock (_sync)
{
_log.AppendLine(line);
if (_log.Length > 80_000)
_log.Remove(0, _log.Length - 80_000);
}
_logger.LogInformation("Memory rebuild: {Line}", line);
}
private string ResolveLatestPath(string relativePath)
{
relativePath = relativePath.Replace('\\', '/').Trim().TrimStart('/');
var parts = relativePath.Split('/', StringSplitOptions.RemoveEmptyEntries);
if (parts.Length == 0 || parts.Any(part => part == "." || part == ".." || part.StartsWith("_", StringComparison.Ordinal)))
throw new ArgumentException("Invalid rebuilt memory path.", nameof(relativePath));
var full = Path.GetFullPath(Path.Combine(_latestDir, Path.Combine(parts)));
var root = Path.GetFullPath(_latestDir);
if (!full.StartsWith(root + Path.DirectorySeparatorChar, StringComparison.OrdinalIgnoreCase))
throw new UnauthorizedAccessException("Rebuilt memory path is outside latest memory directory.");
return full;
}
private static List<string> ExtractMessageBlocks(string text)
{
var matches = Regex.Matches(text, @"^### (Raymond|Claude|Pyrite|Kai) - \d{4}-\d{2}-\d{2}T.*$", RegexOptions.Multiline);
var blocks = new List<string>();
for (var i = 0; i < matches.Count; i++)
{
var start = matches[i].Index;
var end = i + 1 < matches.Count ? matches[i + 1].Index : text.Length;
var block = text[start..end].Trim();
if (!string.IsNullOrWhiteSpace(block))
blocks.Add(block);
}
if (blocks.Count == 0 && !string.IsNullOrWhiteSpace(text))
blocks.Add(text.Trim());
return blocks;
}
private static int BuildChunkCount(List<string> blocks, int maxChars)
{
var chunks = 0;
var currentChars = 0;
foreach (var block in blocks)
{
if (block.Length > maxChars)
{
if (currentChars > 0)
{
chunks++;
currentChars = 0;
}
chunks += (int)Math.Ceiling(block.Length / (double)maxChars);
continue;
}
if (currentChars > 0 && currentChars + block.Length > maxChars)
{
chunks++;
currentChars = 0;
}
currentChars += block.Length;
}
if (currentChars > 0)
chunks++;
return chunks;
}
private static int EstimateTokens(string text) => (int)Math.Ceiling(text.Length * TokensPerCharacter);
private static int LoadOrder(string relativePath)
{
var preferred = new[]
{
"core.md", "recent.md", "history.md", "history.rebuild.md", "communication.md", "raymond.md", "kai.md",
"health.md", "story-world.md", "projects.md", "preferences.md",
"index.md", "manifest.md"
};
var idx = Array.FindIndex(preferred, p => string.Equals(p, relativePath, StringComparison.OrdinalIgnoreCase));
return idx >= 0 ? idx : 1_000;
}
private sealed record SourceMetadata(string Key, int Order, int? PartNumber, string DisplayName)
{
public static SourceMetadata? FromFile(string path, int order)
{
var file = Path.GetFileName(path);
var part = Regex.Match(file, @"^(\d{3})_The Chat Part (\d+)\.md$", RegexOptions.IgnoreCase);
if (part.Success && int.TryParse(part.Groups[2].Value, out var partNumber))
{
var display = $"The Chat Part {partNumber}";
return new($"{order:D3}_{Slugify(display)}", order, partNumber, display);
}
var test = Regex.Match(file, @"^T(\d{2})_Test Chat (\d+)\.md$", RegexOptions.IgnoreCase);
if (test.Success && int.TryParse(test.Groups[2].Value, out var testNumber))
{
var display = $"Test Chat {testNumber}";
return new($"{order:D3}_{Slugify(display)}", order, null, display);
}
return null;
}
private static string Slugify(string value)
{
value = Regex.Replace(value.ToLowerInvariant(), @"[^a-z0-9]+", "-").Trim('-');
return string.IsNullOrWhiteSpace(value) ? "conversation" : value;
}
}
}