swarm-assistent/AssistentPersist.cs

using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Text;
using System.Text.RegularExpressions;
using System.Threading.Tasks;
using Newtonsoft.Json.Linq;
using SwarmUI.Accounts;
using SwarmUI.Utils;

namespace Mrleo1nid.SwarmAssistent;

/// <summary>Disk persistence for chat sessions and UI state under <c>DataRoot()/Assistent/</c>.
/// Chats survive browser storage wipes and follow the data volume across gpu-rent VMs.</summary>
public partial class SwarmAssistentExtension
{
    const int MaxChatsOnDisk = 60;
    const int MaxChatMessagesOnDisk = 40;
    const int MaxChatMessageChars = 4000;

    static readonly Regex ChatIdRe = new(@"^[A-Za-z0-9][A-Za-z0-9_\-]{0,63}$", RegexOptions.Compiled);

    /// <summary>UI-state keys accepted from the browser — anything else is dropped.</summary>
    static readonly string[] UiStateKeys =
    [
        "pack", "persona", "auto_vision", "auto_apply", "auto_generate", "auto_critique",
        "auto_download", "pane_width", "embed_model", "base_url", "model", "view", "board_tab",
    ];

    string AssistentDataDir() => Path.Combine(DataRoot(), "Assistent");

    string AssistentChatsDir() => Path.Combine(AssistentDataDir(), "chats");

    string AssistentUiStatePath() => Path.Combine(AssistentDataDir(), "ui-state.json");

    static string SafeChatId(string id)
    {
        string s = (id ?? "").Trim();
        return ChatIdRe.IsMatch(s) ? s : null;
    }

    string ChatFilePath(string id)
    {
        string safe = SafeChatId(id);
        return safe is null ? null : Path.Combine(AssistentChatsDir(), $"{safe}.json");
    }

    JObject ReadChatFile(string path)
    {
        if (string.IsNullOrWhiteSpace(path) || !File.Exists(path))
        {
            return null;
        }
        try
        {
            JObject chat = JObject.Parse(File.ReadAllText(path, Encoding.UTF8));
            string id = SafeChatId(chat["id"]?.ToString() ?? Path.GetFileNameWithoutExtension(path));
            if (id is null)
            {
                return null;
            }
            chat["id"] = id;
            return chat;
        }
        catch (Exception ex)
        {
            Logs.Debug($"AssistentPersist read {path}: {ex.Message}");
            return null;
        }
    }

    static JObject ChatSummary(JObject chat)
    {
        JArray messages = chat["messages"] as JArray ?? [];
        return new JObject
        {
            ["id"] = chat["id"],
            ["title"] = chat["title"]?.ToString() ?? "Новый чат",
            ["createdAt"] = chat["createdAt"] ?? 0,
            ["updatedAt"] = chat["updatedAt"] ?? 0,
            ["messages_count"] = messages.Count,
            ["params"] = chat["params"],
        };
    }

    List<JObject> LoadAllChats()
    {
        string dir = AssistentChatsDir();
        if (!Directory.Exists(dir))
        {
            return [];
        }
        List<JObject> chats = [];
        foreach (string file in Directory.EnumerateFiles(dir, "*.json"))
        {
            JObject chat = ReadChatFile(file);
            if (chat is not null)
            {
                chats.Add(chat);
            }
        }
        return chats
            .OrderByDescending(c => c["updatedAt"]?.Value<long?>() ?? 0)
            .ToList();
    }

    /// <summary>All chats on disk, newest first. Pass <c>with_messages</c> to get full transcripts.</summary>
    public async Task<JObject> AssistentListChats(Session session, bool with_messages = false, int limit = MaxChatsOnDisk)
    {
        await Task.CompletedTask;
        try
        {
            List<JObject> chats = LoadAllChats();
            int take = Math.Clamp(limit, 1, MaxChatsOnDisk);
            JArray list = [];
            foreach (JObject chat in chats.Take(take))
            {
                list.Add(with_messages ? chat : ChatSummary(chat));
            }
            return new JObject
            {
                ["success"] = true,
                ["chats"] = list,
                ["total"] = chats.Count,
                ["path"] = AssistentChatsDir(),
            };
        }
        catch (Exception ex)
        {
            return new JObject { ["error"] = $"chats list: {ex.Message}" };
        }
    }

    public async Task<JObject> AssistentGetChat(Session session, string id)
    {
        await Task.CompletedTask;
        string path = ChatFilePath(id);
        if (path is null)
        {
            return new JObject { ["error"] = "valid id required" };
        }
        JObject chat = ReadChatFile(path);
        return new JObject
        {
            ["success"] = true,
            ["id"] = SafeChatId(id),
            ["found"] = chat is not null,
            ["chat"] = chat,
        };
    }

    /// <summary>Writes one chat to <c>Assistent/chats/&lt;id&gt;.json</c>.
    /// SwarmUI hands the whole request body to a JObject param, so <c>messages</c> (array)
    /// and <c>params</c> (object) are read out of <paramref name="raw"/>.</summary>
    public async Task<JObject> AssistentSaveChat(Session session, string id, string title, JObject raw)
    {
        await Task.CompletedTask;
        string safe = SafeChatId(id);
        if (safe is null)
        {
            return new JObject { ["error"] = "valid id required" };
        }
        JArray messages = raw?["messages"] as JArray ?? [];
        JObject chatParams = raw?["params"] as JObject;
        JArray trimmed = [];
        foreach (JToken msg in messages.Skip(Math.Max(0, messages.Count - MaxChatMessagesOnDisk)))
        {
            if (msg is not JObject mo)
            {
                continue;
            }
            JObject copy = new()
            {
                ["role"] = mo["role"]?.ToString() ?? "user",
                ["content"] = Clip(mo["content"]?.ToString() ?? "", MaxChatMessageChars),
            };
            if (!string.IsNullOrWhiteSpace(mo["persona"]?.ToString()))
            {
                copy["persona"] = mo["persona"];
            }
            if (!string.IsNullOrWhiteSpace(mo["pack"]?.ToString()))
            {
                copy["pack"] = mo["pack"];
            }
            trimmed.Add(copy);
        }

        long now = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds();
        string path = ChatFilePath(safe);
        JObject existing = ReadChatFile(path);
        JObject chat = new()
        {
            ["id"] = safe,
            ["title"] = string.IsNullOrWhiteSpace(title) ? (existing?["title"]?.ToString() ?? "Новый чат") : title.Trim(),
            ["createdAt"] = raw?["createdAt"]?.Value<long?>() ?? existing?["createdAt"]?.Value<long?>() ?? now,
            ["updatedAt"] = raw?["updatedAt"]?.Value<long?>() ?? now,
            ["messages"] = trimmed,
            ["params"] = chatParams ?? existing?["params"],
        };
        try
        {
            Directory.CreateDirectory(AssistentChatsDir());
            File.WriteAllText(path, chat.ToString(Newtonsoft.Json.Formatting.Indented), Encoding.UTF8);
            PruneChatsOnDisk();
            return new JObject { ["success"] = true, ["id"] = safe, ["path"] = path };
        }
        catch (Exception ex)
        {
            return new JObject { ["error"] = $"chat save: {ex.Message}" };
        }
    }

    public async Task<JObject> AssistentDeleteChat(Session session, string id)
    {
        await Task.CompletedTask;
        string path = ChatFilePath(id);
        if (path is null)
        {
            return new JObject { ["error"] = "valid id required" };
        }
        try
        {
            bool existed = File.Exists(path);
            if (existed)
            {
                File.Delete(path);
            }
            return new JObject { ["success"] = true, ["deleted"] = existed, ["id"] = SafeChatId(id) };
        }
        catch (Exception ex)
        {
            return new JObject { ["error"] = $"chat delete: {ex.Message}" };
        }
    }

    void PruneChatsOnDisk()
    {
        try
        {
            List<JObject> chats = LoadAllChats();
            if (chats.Count <= MaxChatsOnDisk)
            {
                return;
            }
            foreach (JObject stale in chats.Skip(MaxChatsOnDisk))
            {
                string path = ChatFilePath(stale["id"]?.ToString());
                if (path is not null && File.Exists(path))
                {
                    File.Delete(path);
                }
            }
        }
        catch (Exception ex)
        {
            Logs.Debug($"AssistentPersist prune: {ex.Message}");
        }
    }

    public async Task<JObject> AssistentGetUiState(Session session)
    {
        await Task.CompletedTask;
        string path = AssistentUiStatePath();
        if (!File.Exists(path))
        {
            return new JObject { ["success"] = true, ["ui_state"] = null };
        }
        try
        {
            return new JObject
            {
                ["success"] = true,
                ["ui_state"] = JObject.Parse(File.ReadAllText(path, Encoding.UTF8)),
            };
        }
        catch (Exception ex)
        {
            return new JObject { ["error"] = $"ui-state.json: {ex.Message}" };
        }
    }

    /// <summary>Persists the whitelisted UI preferences. Reads <c>ui_state</c> from the body,
    /// falling back to the flat body for convenience.</summary>
    public async Task<JObject> AssistentSaveUiState(Session session, JObject raw)
    {
        await Task.CompletedTask;
        JObject source = raw?["ui_state"] as JObject ?? raw;
        if (source is null)
        {
            return new JObject { ["error"] = "ui_state required" };
        }
        JObject clean = new();
        foreach (string key in UiStateKeys)
        {
            JToken value = source[key];
            if (value is null || value.Type == JTokenType.Null)
            {
                continue;
            }
            clean[key] = value.DeepClone();
        }
        if (clean.Count == 0)
        {
            return new JObject { ["error"] = "ui_state has no known keys" };
        }
        clean["updated"] = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds();
        try
        {
            Directory.CreateDirectory(AssistentDataDir());
            string path = AssistentUiStatePath();
            File.WriteAllText(path, clean.ToString(Newtonsoft.Json.Formatting.Indented), Encoding.UTF8);
            return new JObject { ["success"] = true, ["path"] = path };
        }
        catch (Exception ex)
        {
            return new JObject { ["error"] = $"ui-state save: {ex.Message}" };
        }
    }
}