Size: 12.7 KB Modified: 3/07/2026 1:54 AM
using System.Text;
using System.Text.RegularExpressions;
using KaiChat.Models;

namespace KaiChat.Services;

public interface IMemorySweepService
{
    Task<MemorySweepResult> SweepBeforeCompactionAsync(
        Conversation conversation,
        IReadOnlyList<ChatMessage> messagesToSweep,
        CompactionPreview preview,
        int compactionNumber,
        CancellationToken ct = default);
}

public record MemorySweepResult(
    bool Changed,
    bool Success,
    string Status,
    string Report,
    string RawResponsePreview,
    IReadOnlyList<MemoryOperationResult> Operations,
    int MessagesSwept,
    int WriteCount,
    int DeleteCount,
    int FailedCount);

public class MemorySweepService : IMemorySweepService
{
    private const int RawPreviewLimit = 4000;
    private readonly IMemoryManagementService _memoryManager;
    private readonly ILogger<MemorySweepService> _logger;

    public MemorySweepService(
        IMemoryManagementService memoryManager,
        ILogger<MemorySweepService> logger)
    {
        _memoryManager = memoryManager;
        _logger = logger;
    }

    public async Task<MemorySweepResult> SweepBeforeCompactionAsync(
        Conversation conversation,
        IReadOnlyList<ChatMessage> messagesToSweep,
        CompactionPreview preview,
        int compactionNumber,
        CancellationToken ct = default)
    {
        var usableMessages = messagesToSweep
            .Where(HasMemorySweepContent)
            .OrderBy(m => NormalizeUtc(m.Timestamp))
            .ToList();

        if (usableMessages.Count == 0)
        {
            var noop = new MemoryOperationResult(
                "noop",
                "",
                true,
                "No message content was available for the pre-compaction memory sweep.",
                "No message content was available for the pre-compaction memory sweep.");

            return BuildResult(
                false,
                "",
                [noop],
                0,
                "Memory sweep skipped; no message content was available.");
        }

        var transcript = BuildTranscript(conversation, usableMessages, preview, compactionNumber);
        _logger.LogInformation(
            "Running pre-compaction memory sweep for {ConversationId}: {Count} messages before compaction #{CompactionNumber}",
            conversation.Id,
            usableMessages.Count,
            compactionNumber);

        var result = await _memoryManager.ManageAsync($"""
        Pre-compaction memory sweep for KaiChat.

        These messages are about to be compacted and removed from the live recent-message window. The compaction archive and compacted summary will preserve broad conversation continuity, but persistent memory should separately retain durable facts that Kai must have immediately available in future chats, after compactions, and after restarts.

        Integrate any missing durable or currently useful information into the memory directory.
        Use the same broad-category memory rules as the historical rebuild and idle memory updater: preserve specifics, avoid tiny files, avoid duplicate facts, and keep all still-valid details when saving a complete file.
        The current memory files may already include AI-added memories or idle update memories from this same conversation. Treat those as existing state; do not duplicate them.
        Only add missing durable details, materially changed context, corrections, active unresolved tasks, relationship/lore state, health context, preferences, project state, or recent-state context that should survive compaction.
        Do not add transient chatter merely because it is being compacted.
        If the existing memory already covers the durable information, use noop.

        {transcript}
        """, ct);

        return BuildResult(
            result.Changed,
            result.RawResponse,
            result.Operations,
            usableMessages.Count,
            "");
    }

    private static MemorySweepResult BuildResult(
        bool changed,
        string rawResponse,
        IReadOnlyList<MemoryOperationResult> operations,
        int messagesSwept,
        string statusOverride)
    {
        var writeCount = operations.Count(o => o.Success && IsWriteAction(o.Action));
        var deleteCount = operations.Count(o => o.Success && IsDeleteAction(o.Action));
        var failedCount = operations.Count(o => !o.Success);
        var success = failedCount == 0;
        var status = !string.IsNullOrWhiteSpace(statusOverride)
            ? statusOverride
            : success
                ? changed
                    ? $"Memory sweep complete. Applied {writeCount} write(s) and {deleteCount} removal(s)."
                    : "Memory sweep complete. No memory changes were needed."
                : $"Memory sweep completed with {failedCount} failed operation(s); compaction should be retried after this is fixed.";

        var report = BuildReport(status, messagesSwept, changed, operations, rawResponse, failedCount > 0);

        return new MemorySweepResult(
            changed,
            success,
            status,
            report,
            failedCount > 0 ? TrimPreview(rawResponse, RawPreviewLimit) : "",
            operations,
            messagesSwept,
            writeCount,
            deleteCount,
            failedCount);
    }

    private static string BuildTranscript(
        Conversation conversation,
        IReadOnlyList<ChatMessage> messages,
        CompactionPreview preview,
        int compactionNumber)
    {
        var builder = new StringBuilder();
        builder.AppendLine($"Conversation: {conversation.Title} ({conversation.Id})");
        builder.AppendLine($"Compaction number: {compactionNumber}");
        builder.AppendLine($"Messages selected for compaction: {preview.MessagesToCompact}");
        builder.AppendLine($"Messages kept live after compaction: {preview.MessagesToKeep}");
        builder.AppendLine($"Estimated API tokens before compaction: {preview.EstimatedApiTokens:N0}");
        builder.AppendLine($"Estimated API tokens after compaction: {preview.EstimatedPostCompactionApiTokens:N0}");
        builder.AppendLine();
        builder.AppendLine("Messages to sweep:");
        builder.AppendLine();

        foreach (var message in messages)
            AppendMessageForMemorySweep(builder, message);

        return builder.ToString();
    }

    private static string BuildReport(
        string status,
        int messagesSwept,
        bool changed,
        IReadOnlyList<MemoryOperationResult> operations,
        string rawResponse,
        bool includeRawPreview)
    {
        var builder = new StringBuilder();
        builder.AppendLine(status);
        builder.AppendLine();
        builder.AppendLine($"Messages swept: {messagesSwept:N0}");
        builder.AppendLine($"Memory changed: {(changed ? "yes" : "no")}");
        builder.AppendLine($"Operations: {operations.Count:N0}");

        if (operations.Count > 0)
        {
            builder.AppendLine();
            builder.AppendLine("Operation details:");
            foreach (var operation in operations)
            {
                var path = string.IsNullOrWhiteSpace(operation.Path) ? "(no path)" : operation.Path;
                var outcome = operation.Success ? "ok" : "failed";
                var reason = string.IsNullOrWhiteSpace(operation.Reason)
                    ? operation.Message
                    : operation.Reason;
                builder.AppendLine($"- {NormalizeAction(operation.Action)} {path} [{outcome}]: {TrimForDisplay(reason, 500)}");
            }
        }

        if (includeRawPreview && !string.IsNullOrWhiteSpace(rawResponse))
        {
            builder.AppendLine();
            builder.AppendLine("Raw response preview:");
            builder.AppendLine(TrimPreview(rawResponse, RawPreviewLimit));
        }

        return builder.ToString().TrimEnd();
    }

    private static bool HasMemorySweepContent(ChatMessage message)
    {
        return !string.IsNullOrWhiteSpace(message.Content)
            || !string.IsNullOrWhiteSpace(message.Thinking)
            || message.Tool is not null
            || message.Compaction is not null
            || message.MemorySweep is not null;
    }

    private static void AppendMessageForMemorySweep(StringBuilder builder, ChatMessage message)
    {
        builder.AppendLine($"[{FormatLocalTime(message.Timestamp)}] [{message.Role}] id={message.Id}");

        if (!string.IsNullOrWhiteSpace(message.Content))
        {
            var content = message.Role == "assistant"
                ? ToolResultTextSanitizer.RemoveEchoedToolResultBlocks(message.Content)
                : message.Content;
            builder.AppendLine(content.Trim());
            builder.AppendLine();
        }

        if (!string.IsNullOrWhiteSpace(message.Thinking))
        {
            builder.AppendLine("[thinking]");
            builder.AppendLine(message.Thinking.Trim());
            builder.AppendLine();
        }

        if (message.Tool is not null)
        {
            builder.AppendLine("[tool]");
            builder.AppendLine($"round: {message.Tool.Round}");
            builder.AppendLine($"request: {message.Tool.Method} {message.Tool.Path}");
            if (!string.IsNullOrWhiteSpace(message.Tool.Body))
                builder.AppendLine($"parameters: {message.Tool.Body.Trim()}");
            builder.AppendLine($"status: {message.Tool.Status}");
            builder.AppendLine($"is_error: {message.Tool.IsError}");
            if (!string.IsNullOrWhiteSpace(message.Tool.Result))
            {
                builder.AppendLine("result:");
                builder.AppendLine(message.Tool.Result.Trim());
            }
            builder.AppendLine();
        }

        if (message.Compaction is not null)
        {
            builder.AppendLine("[compaction]");
            builder.AppendLine($"number: {message.Compaction.CompactionNumber}");
            builder.AppendLine($"archive: {message.Compaction.ArchivePath}");
            builder.AppendLine($"previous_archive: {message.Compaction.PreviousArchivePath}");
            builder.AppendLine($"messages_to_compact: {message.Compaction.MessagesToCompact}");
            builder.AppendLine($"messages_to_keep: {message.Compaction.MessagesToKeep}");
            builder.AppendLine($"chunks: {message.Compaction.ChunkCount}");
            builder.AppendLine($"estimated_api_tokens: {message.Compaction.EstimatedApiTokens}");
            builder.AppendLine($"estimated_post_compaction_api_tokens: {message.Compaction.EstimatedPostCompactionApiTokens}");
            builder.AppendLine($"started_at: {FormatLocalTime(message.Compaction.StartedAt)}");
            builder.AppendLine($"completed_at: {FormatLocalTime(message.Compaction.CompletedAt)}");
            builder.AppendLine();
        }

        if (message.MemorySweep is not null)
        {
            builder.AppendLine("[memory sweep]");
            builder.AppendLine($"compaction_number: {message.MemorySweep.CompactionNumber}");
            builder.AppendLine($"messages_swept: {message.MemorySweep.MessagesSwept}");
            builder.AppendLine($"changed: {message.MemorySweep.Changed}");
            builder.AppendLine($"success: {message.MemorySweep.Success}");
            builder.AppendLine($"operations: {message.MemorySweep.OperationCount}");
            builder.AppendLine($"status: {message.MemorySweep.Status}");
            builder.AppendLine();
        }
    }

    private static bool IsWriteAction(string action)
    {
        return NormalizeAction(action) is "save" or "write" or "upsert";
    }

    private static bool IsDeleteAction(string action)
    {
        return NormalizeAction(action) is "soft_delete" or "delete";
    }

    private static string NormalizeAction(string action)
    {
        return (action ?? "").Trim().ToLowerInvariant();
    }

    private static string TrimPreview(string text, int maxLength)
    {
        if (string.IsNullOrWhiteSpace(text))
            return "";

        text = text.Trim();
        return text.Length <= maxLength
            ? text
            : text[..Math.Max(0, maxLength - 3)] + "...";
    }

    private static string TrimForDisplay(string? text, int maxLength)
    {
        if (string.IsNullOrWhiteSpace(text))
            return "";

        var compact = Regex.Replace(text.Replace("\r\n", " ").Replace('\r', ' ').Replace('\n', ' '), @"\s+", " ").Trim();
        return compact.Length <= maxLength
            ? compact
            : compact[..Math.Max(0, maxLength - 3)] + "...";
    }

    private static string FormatLocalTime(DateTime value)
    {
        return value.ToLocalTime().ToString("yyyy-MM-dd HH:mm:ss");
    }

    private static DateTime NormalizeUtc(DateTime value)
    {
        return value == DateTime.MinValue
            ? DateTime.MinValue
            : value.ToUniversalTime();
    }
}
Offline