swarm-assistent/AssistentPersist.cs

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

namespace Mrleo1nid.SwarmAssistent;

/// <summary>Chat sessions and UI state in <c>Assistent/memory/assistent.sqlite</c>.
/// Survives browser wipes and follows the data volume across gpu-rent VMs.</summary>
public partial class SwarmAssistentExtension
{
    const int MaxChatMessagesOnDisk = 40;
    const int MaxChatMessageChars = 4000;

    static readonly System.Text.RegularExpressions.Regex ChatIdRe = new(@"^[A-Za-z0-9][A-Za-z0-9_\-]{0,63}$", System.Text.RegularExpressions.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",
    ];

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

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

    /// <summary>All chats, newest first. <paramref name="q"/> searches title+body via FTS.</summary>
    public async Task<JObject> AssistentListChats(Session session, bool with_messages = false, int limit = 60, string q = null)
    {
        await Task.CompletedTask;
        try
        {
            int take = Math.Clamp(limit, 1, AssistentMemory.MaxChatsStored);
            List<JObject> chats = Memory.ListChats(with_messages, take, q);
            JArray list = [];
            foreach (JObject chat in chats)
            {
                list.Add(with_messages ? chat : ChatSummary(chat));
            }
            return new JObject
            {
                ["success"] = true,
                ["chats"] = list,
                ["total"] = string.IsNullOrWhiteSpace(q) ? Memory.CountChats() : list.Count,
                ["path"] = "Assistent/memory/assistent.sqlite",
            };
        }
        catch (Exception ex)
        {
            return new JObject { ["error"] = $"chats list: {ex.Message}" };
        }
    }

    public async Task<JObject> AssistentGetChat(Session session, string id)
    {
        await Task.CompletedTask;
        string safe = SafeChatId(id);
        if (safe is null)
        {
            return new JObject { ["error"] = "valid id required" };
        }
        try
        {
            JObject chat = Memory.GetChat(safe);
            return new JObject
            {
                ["success"] = true,
                ["id"] = safe,
                ["found"] = chat is not null,
                ["chat"] = chat,
            };
        }
        catch (Exception ex)
        {
            return new JObject { ["error"] = $"chat get: {ex.Message}" };
        }
    }

    /// <summary>Writes one chat into sqlite.
    /// 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();
        JObject existing = null;
        try
        {
            existing = Memory.GetChat(safe);
        }
        catch (Exception ex)
        {
            Logs.Debug($"AssistentSaveChat existing: {ex.Message}");
        }
        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"] as JObject,
        };
        try
        {
            Memory.SaveChat(chat);
            return new JObject { ["success"] = true, ["id"] = safe, ["path"] = "Assistent/memory/assistent.sqlite" };
        }
        catch (Exception ex)
        {
            return new JObject { ["error"] = $"chat save: {ex.Message}" };
        }
    }

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

    public async Task<JObject> AssistentGetUiState(Session session)
    {
        await Task.CompletedTask;
        try
        {
            return new JObject
            {
                ["success"] = true,
                ["ui_state"] = Memory.GetKvObject(AssistentMemory.KvUiState),
            };
        }
        catch (Exception ex)
        {
            return new JObject { ["error"] = $"ui-state: {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
        {
            Memory.SetKvObject(AssistentMemory.KvUiState, clean);
            return new JObject { ["success"] = true, ["path"] = "Assistent/memory/assistent.sqlite" };
        }
        catch (Exception ex)
        {
            return new JObject { ["error"] = $"ui-state save: {ex.Message}" };
        }
    }
}