Size: 16.4 KB Modified: 6/07/2026 3:22 AM
using System.Text.Json;
using System.Text.Json.Serialization;
using System.Text.RegularExpressions;
using KaiChat.Models;

namespace KaiChat.Services;

public interface IMemoryManagementService
{
    Task<MemoryManagementResult> ManageAsync(string instruction, CancellationToken ct = default);
}

public record MemoryManagementResult(
    bool Changed,
    string RawResponse,
    IReadOnlyList<MemoryOperationResult> Operations);

public record MemoryOperationResult(
    string Action,
    string Path,
    bool Success,
    string Message,
    string? Reason = null);

public class MemoryManagementService : IMemoryManagementService
{
    private readonly IMemoryService _memory;
    private readonly IChatService _chat;
    private readonly ILogger<MemoryManagementService> _logger;
    private static readonly JsonSerializerOptions JsonOptions = new()
    {
        PropertyNameCaseInsensitive = true,
        WriteIndented = true
    };
    private static readonly Regex TruncationPlaceholderRegex = new(
        @"(?m)^\s*[-*]?\s*\.\.\.\s*\([^)]*(?:rest|remaining|preserve|unchanged|omitted|truncated)[^)]*\)\s*\.{0,3}\s*$",
        RegexOptions.IgnoreCase | RegexOptions.CultureInvariant);

    public MemoryManagementService(
        IMemoryService memory,
        IChatService chat,
        ILogger<MemoryManagementService> logger)
    {
        _memory = memory;
        _chat = chat;
        _logger = logger;
    }

    public async Task<MemoryManagementResult> ManageAsync(string instruction, CancellationToken ct = default)
    {
        var prompt = await BuildPromptAsync(instruction);
        ct.ThrowIfCancellationRequested();

        var msgs = new List<ChatMessage>
        {
            new() { Role = "system", Content = prompt }
        };

        var raw = await _chat.ChatAsync(msgs, ct);
        ct.ThrowIfCancellationRequested();

        MemoryChangeSet changeSet;
        try
        {
            changeSet = ParseChangeSet(raw);
        }
        catch (JsonException ex)
        {
            _logger.LogWarning(ex, "Memory management response was not valid JSON");
            return new MemoryManagementResult(
                false,
                raw,
                [new("invalid_response", "", false, "AI returned invalid memory JSON; no changes were applied.")]);
        }

        var results = new List<MemoryOperationResult>();

        foreach (var op in changeSet.Operations)
        {
            ct.ThrowIfCancellationRequested();

            var action = op.Action.Trim().ToLowerInvariant();
            try
            {
                switch (action)
                {
                    case "save":
                    case "write":
                    case "upsert":
                        if (string.IsNullOrWhiteSpace(op.Path))
                            throw new InvalidOperationException("Save operation requires a path.");
                        var content = NormalizeManagedContent(op.Path, op.Content ?? "");
                        if (LooksLikeTruncationPlaceholder(content, out var guardReason))
                        {
                            var preview = GetContentPreview(op.Content);
                            var existingSize = await ReadExistingContentLengthAsync(op.Path);
                            _logger.LogWarning(
                                "Memory save blocked by truncation guard for {Path}: {Reason}. Previous bytes: {ExistingBytes}, new content preview: {Preview}",
                                op.Path,
                                guardReason,
                                existingSize,
                                preview);
                            results.Add(new(action, op.Path, false, "Memory write blocked by guardrail: truncated omission markers detected.", "truncation_placeholder_detected"));
                            break;
                        }
                        if (string.IsNullOrWhiteSpace(content))
                            throw new InvalidOperationException("Save operation produced empty memory content.");
                        await _memory.SaveMemoryFileAsync(op.Path, content);
                        results.Add(new(action, op.Path, true, "Saved memory file.", op.Reason));
                        break;

                    case "soft_delete":
                    case "delete":
                        if (string.IsNullOrWhiteSpace(op.Path))
                            throw new InvalidOperationException("Delete operation requires a path.");
                        if (string.IsNullOrWhiteSpace(op.Reason))
                            throw new InvalidOperationException("Delete operation requires a reason.");
                        await _memory.SoftDeleteMemoryFileAsync(op.Path, op.Reason);
                        results.Add(new("soft_delete", op.Path, true, "Soft-deleted memory file.", op.Reason));
                        break;

                    case "noop":
                    case "none":
                        results.Add(new("noop", op.Path ?? "", true, op.Reason ?? "No memory changes needed.", op.Reason));
                        break;

                    default:
                        results.Add(new(action, op.Path ?? "", false, $"Unknown memory operation: {op.Action}"));
                        break;
                }
            }
            catch (Exception ex)
            {
                _logger.LogWarning(ex, "Memory operation failed: {Action} {Path}", op.Action, op.Path);
                results.Add(new(action, op.Path ?? "", false, ex.Message));
            }
        }

        return new MemoryManagementResult(
            results.Any(r => r.Success && r.Action is not "noop" and not "none"),
            raw,
            results);
    }

    private async Task<string> BuildPromptAsync(string instruction)
    {
        var files = await _memory.ListMemoryFilesAsync(includeTrash: false);
        var fileBlocks = new List<string>();

        foreach (var file in files.Where(f => f.IsLoaded))
        {
            var content = await _memory.ReadMemoryFileAsync(file.Path);
            fileBlocks.Add($"""
            Memory file `{file.Path}`:
            {content.Trim()}
            """);
        }

        var currentFiles = fileBlocks.Count == 0
            ? "(no memory files yet)"
            : string.Join("\n\n", fileBlocks);
        var memoryDirectory = _memory.MemoryDirectoryRelativePath;

        return $$"""
        You are KaiChat's memory librarian. Update the memory directory cautiously and precisely.

        Memory is stored as markdown files under `{memoryDirectory}`.
        Loaded files are shown below. Files in `_trash` are not loaded into model context.

        Current memory files:

        {{currentFiles}}

        User instruction:
        {{instruction}}

        Rules:
        - Preserve durable facts, relationship/lore state, active projects, health context, preferences, corrections, and unresolved tasks.
        - Prefer editing an existing memory file over creating near-duplicate files.
        - Default to broad memory files such as `core.md`, `recent.md`, `history.md`, `communication.md`, `raymond.md`, `kai.md`, `health.md`, `story-world.md`, `projects.md`, and `preferences.md`.
        - You may create a new memory file only when it is a broad, durable category likely to stay useful across many future sessions.
        - Do not create tiny files for individual chats, days, scenes, moods, incidents, single jokes, or narrow one-off topics. Merge those into an existing broad category.
        - Prefer a small, stable memory directory over many precise files. If in doubt, update an existing broad category.
        - Broad category files may be as long and detailed as needed. Preserve granular concrete details, dates, names, preferences, recurring jokes, corrections, failure modes, unresolved threads, and emotionally important context.
        - Do not compress away specifics merely to keep files short. Only omit content that is clearly transient, duplicated, contradicted later, or not useful for future context.
        - Treat existing memory files as a curated baseline. A `save` operation replaces the complete file, so preserve every still-valid durable detail unless it is clearly duplicated elsewhere, contradicted by the new instruction, or intentionally moved to a better file.
        - Prefer surgical additive edits over rewriting a whole file from scratch. If you move information between files, ensure the final memory directory still contains the same durable facts exactly once.
        - When new information corrects, refines, or supersedes an existing rule/fact, reconcile the memory directory in the same response: find older or conflicting versions of that same point across all loaded files and update or remove those stale versions instead of adding a parallel newer copy.
        - Cross-file repetition is allowed only when each file genuinely needs the point; repeated entries must agree with the latest correction/refinement and must not leave older contradictory wording behind.
        - Do not replace specific earlier facts with vague summaries. If the old memory says a date, name, tool, dosage, schedule, failure mode, or concrete incident, keep that specificity unless it is wrong.
        - If you are uncertain whether an earlier detail still matters, keep it.
        - Be very cautious with deletion. Prefer updating stale content. If deleting a file, use `soft_delete`, include a concrete reason, and only delete files that are clearly obsolete, duplicate, empty, or superseded.
        - Include a concise `reason` for every operation. For save/write/upsert, explain what durable information changed and why. For noop, explain why no memory change is needed.
        - Never write to `_trash` directly.
        - Keep `core.md` compact and high-signal. Put detailed topic context into broad category files.
        - `recent.md` should contain current active state and unresolved threads likely needed at the start of new chats or after compaction.
        - `history.md` is a lookup/reference timeline. Add durable historical milestones there when needed, but keep current active state in `recent.md` and topic files.
        - The `content` field must contain only the markdown that belongs inside that file. Do not include wrapper labels, `FILE:` lines, delimiter-only `---` lines, or code fences.
        - Do not invent facts. If the provided context is insufficient, use `noop`.
        - Never replace omitted sections with shorthand placeholder text. Never include constructs like `... (rest ...) ...` or other ellipsis-based omissions in saved content.
        - If content is too long to include fully in your reasoning context, return `noop` and wait for a follow-up rather than truncating memory files.

        Respond with ONLY valid JSON in this schema:
        {
          "operations": [
            {
              "action": "save",
              "path": "recent.md",
              "content": "complete markdown content for that file",
              "reason": "what durable information was added, changed, moved, or removed and why"
            },
            {
              "action": "soft_delete",
              "path": "obsolete.md",
              "reason": "why this file is obsolete or superseded"
            },
            {
              "action": "noop",
              "reason": "why no memory change is needed"
            }
          ]
        }
        """;
    }

    private static string NormalizeManagedContent(string path, string content)
    {
        content = StripWholeCodeFence(content.Trim());

        var lines = content
            .Replace("\r\n", "\n")
            .Replace('\r', '\n')
            .Split('\n')
            .Select(line => line.TrimEnd())
            .ToList();

        TrimOuterBlankLines(lines);

        if (lines.Count > 0 && IsCopiedPromptWrapperLine(lines[0], path))
        {
            lines.RemoveAt(0);
            TrimOuterBlankLines(lines);
        }

        while (lines.Count > 0 && lines[0].Trim() == "---")
        {
            lines.RemoveAt(0);
            TrimOuterBlankLines(lines);
        }

        while (lines.Count > 0 && IsCopiedPromptWrapperLine(lines[^1], path))
        {
            lines.RemoveAt(lines.Count - 1);
            TrimOuterBlankLines(lines);
        }

        while (lines.Count > 0 && lines[^1].Trim() == "---")
        {
            lines.RemoveAt(lines.Count - 1);
            TrimOuterBlankLines(lines);
        }

        return string.Join(Environment.NewLine, lines).TrimEnd();
    }

    private static string StripWholeCodeFence(string content)
    {
        if (!content.StartsWith("```", StringComparison.Ordinal))
            return content;

        var firstNewline = content.IndexOf('\n');
        var lastFence = content.LastIndexOf("```", StringComparison.Ordinal);
        return firstNewline >= 0 && lastFence > firstNewline
            ? content[(firstNewline + 1)..lastFence].Trim()
            : content;
    }

    private static bool IsCopiedPromptWrapperLine(string line, string path)
    {
        var trimmed = line.Trim();
        return trimmed.StartsWith("FILE:", StringComparison.OrdinalIgnoreCase)
            || trimmed.StartsWith("Memory file `", StringComparison.OrdinalIgnoreCase)
            || trimmed.StartsWith("[Memory file", StringComparison.OrdinalIgnoreCase)
            || trimmed.StartsWith("<memory-file", StringComparison.OrdinalIgnoreCase)
            || trimmed.StartsWith("</memory-file", StringComparison.OrdinalIgnoreCase)
            || trimmed.Equals(path, StringComparison.OrdinalIgnoreCase);
    }

    private static void TrimOuterBlankLines(List<string> lines)
    {
        while (lines.Count > 0 && string.IsNullOrWhiteSpace(lines[0]))
            lines.RemoveAt(0);
        while (lines.Count > 0 && string.IsNullOrWhiteSpace(lines[^1]))
            lines.RemoveAt(lines.Count - 1);
    }

    private static MemoryChangeSet ParseChangeSet(string raw)
    {
        var json = ExtractJson(raw);
        var changeSet = JsonSerializer.Deserialize<MemoryChangeSet>(json, JsonOptions);
        if (changeSet?.Operations == null || changeSet.Operations.Count == 0)
            return new MemoryChangeSet { Operations = [new() { Action = "noop", Reason = "AI returned no operations." }] };
        return changeSet;
    }

    private static string ExtractJson(string raw)
    {
        raw = raw.Trim();
        if (raw.StartsWith("```", StringComparison.Ordinal))
        {
            var firstNewline = raw.IndexOf('\n');
            var lastFence = raw.LastIndexOf("```", StringComparison.Ordinal);
            if (firstNewline >= 0 && lastFence > firstNewline)
                raw = raw[(firstNewline + 1)..lastFence].Trim();
        }

        var firstBrace = raw.IndexOf('{');
        var lastBrace = raw.LastIndexOf('}');
        if (firstBrace >= 0 && lastBrace > firstBrace)
            return raw[firstBrace..(lastBrace + 1)];

        return raw;
    }

    private static bool LooksLikeTruncationPlaceholder(string content, out string reason)
    {
        var lines = content
            .Replace("\r\n", "\n")
            .Replace('\r', '\n')
            .Split('\n');

        foreach (var line in lines)
        {
            if (!TruncationPlaceholderRegex.IsMatch(line))
                continue;

            reason = "Detected ellipsis omission placeholder in memory content.";
            return true;
        }

        reason = "";
        return false;
    }

    private static string GetContentPreview(string? content)
    {
        if (string.IsNullOrWhiteSpace(content))
            return "<empty>";

        var compact = content.Replace('\r', '\n').Trim();
        var firstLine = compact.Split('\n')[0].Trim();
        return firstLine.Length > 240
            ? firstLine[..240] + "..."
            : firstLine;
    }

    private async Task<long> ReadExistingContentLengthAsync(string relativePath)
    {
        try
        {
            var existing = await _memory.ReadMemoryFileAsync(relativePath);
            return existing.Length;
        }
        catch
        {
            return 0;
        }
    }

    private sealed class MemoryChangeSet
    {
        [JsonPropertyName("operations")]
        public List<MemoryOperation> Operations { get; set; } = new();
    }

    private sealed class MemoryOperation
    {
        [JsonPropertyName("action")]
        public string Action { get; set; } = "noop";

        [JsonPropertyName("path")]
        public string? Path { get; set; }

        [JsonPropertyName("content")]
        public string? Content { get; set; }

        [JsonPropertyName("reason")]
        public string? Reason { get; set; }
    }
}
Offline