Reduced memory usage and fixed several memory leaks (#933)

app/MindWork AI Studio/Tools/AIJobs/AIJobService.cs

@@ -16,7 +16,13 @@ public sealed class AIJobService(SettingsManager settingsManager, MessageBus mes
 
         public required CancellationToken CancellationToken { get; init; }
 
-        public required ChatGenerationRequest ChatGenerationRequest { get; init; }
+        /// <summary>
+        /// What the job works on. This is the heavy part of a job: it holds the entire chat thread.
+        /// We release it once the job is done, so a finished job does not keep a chat alive for as
+        /// long as the app runs. Everything a finished job still has to answer lives in the
+        /// snapshot, which is small.
+        /// </summary>
+        public ChatGenerationRequest? ChatGenerationRequest { get; set; }
 
         public required AIJobSnapshot Snapshot { get; set; }
 
@@ -70,7 +76,7 @@ public sealed class AIJobService(SettingsManager settingsManager, MessageBus mes
         if (!this.activeChatJobsByChatId.TryGetValue(chatId, out var jobId))
             return null;
 
-        return this.jobs.TryGetValue(jobId, out var job) ? job.ChatGenerationRequest.ChatThread : null;
+        return this.jobs.TryGetValue(jobId, out var job) ? job.ChatGenerationRequest?.ChatThread : null;
     }
 
     public async Task<AIJobSnapshot?> TryStartChatGenerationAsync(ChatGenerationRequest request)
@@ -185,6 +191,9 @@ public sealed class AIJobService(SettingsManager settingsManager, MessageBus mes
     private async Task RunChatGenerationAsync(AIJobState state)
     {
         var request = state.ChatGenerationRequest;
+        if (request is null)
+            return;
+
         var token = state.CancellationToken;
 
         try
@@ -281,7 +290,11 @@ public sealed class AIJobService(SettingsManager settingsManager, MessageBus mes
             state.IsCompletionStarted = true;
         }
 
-        var aiText = state.ChatGenerationRequest.AIText;
+        var request = state.ChatGenerationRequest;
+        if (request is null)
+            return;
+
+        var aiText = request.AIText;
         aiText.InitialRemoteWait = false;
         aiText.IsStreaming = false;
         aiText.Text = aiText.Text.RemoveThinkTags().Trim();
@@ -298,31 +311,72 @@ public sealed class AIJobService(SettingsManager settingsManager, MessageBus mes
             };
         }
 
-        this.activeChatJobsByChatId.TryRemove(state.ChatGenerationRequest.ChatThread.ChatId, out _);
+        this.activeChatJobsByChatId.TryRemove(request.ChatThread.ChatId, out _);
         await CheckpointChatAsync(state, force: true);
         await this.NotifyChangedAsync(state);
         await messageBus.SendMessage(null, Event.AI_JOB_FINISHED, state.Snapshot);
         state.CancellationTokenSource.Dispose();
+
+        //
+        // The chat is stored and everyone was told about it, so nothing needs the request anymore.
+        // Releasing it here is what keeps a finished job from holding an entire chat thread — even
+        // one the user has deleted in the meantime. We do it under the lock, because that is where
+        // every other access to the state happens:
+        //
+        lock (state.SyncRoot)
+        {
+            state.ChatGenerationRequest = null;
+        }
+
+        this.PruneCompletedJobs(state.Snapshot);
+    }
+
+    /// <summary>
+    /// Drops the finished jobs which nothing needs anymore.
+    /// </summary>
+    /// <remarks>
+    /// What the app asks for is the outcome of the last generation of a chat, cf. TryGetChatSnapshot.
+    /// Everything older than that is a history no one reads, and it would grow for as long as the
+    /// app runs. Active jobs are never touched, and neither is the job we just finished.
+    /// </remarks>
+    /// <param name="latest">The snapshot of the job which just finished.</param>
+    private void PruneCompletedJobs(AIJobSnapshot latest)
+    {
+        var supersededJobIds = this.jobs.Values
+            .Select(job => job.Snapshot)
+            .Where(snapshot => snapshot.Kind == latest.Kind)
+            .Where(snapshot => snapshot.SubjectId == latest.SubjectId)
+            .Where(snapshot => snapshot.JobId != latest.JobId)
+            .Where(snapshot => !snapshot.IsActive)
+            .Select(snapshot => snapshot.JobId)
+            .ToList();
+
+        foreach (var jobId in supersededJobIds)
+            this.jobs.TryRemove(jobId, out _);
     }
 
     private static void RemoveEmptyAIResponse(AIJobState state)
     {
-        var aiText = state.ChatGenerationRequest.AIText;
+        var request = state.ChatGenerationRequest;
+        if (request is null)
+            return;
+
+        var aiText = request.AIText;
         if (!string.IsNullOrWhiteSpace(aiText.Text))
             return;
 
-        var aiBlock = state.ChatGenerationRequest.ChatThread.Blocks
+        var aiBlock = request.ChatThread.Blocks
             .LastOrDefault(block => ReferenceEquals(block.Content, aiText));
 
         if (aiBlock is not null)
-            state.ChatGenerationRequest.ChatThread.Blocks.Remove(aiBlock);
+            request.ChatThread.Blocks.Remove(aiBlock);
     }
 
     private static bool TrySetWaitingForRemote(AIJobState state, CancellationToken token)
     {
         lock (state.SyncRoot)
         {
-            if (state.IsCompletionStarted || token.IsCancellationRequested)
+            if (state.IsCompletionStarted || token.IsCancellationRequested || state.ChatGenerationRequest is null)
                 return false;
 
             state.ChatGenerationRequest.AIText.InitialRemoteWait = true;
@@ -334,7 +388,7 @@ public sealed class AIJobService(SettingsManager settingsManager, MessageBus mes
     {
         lock (state.SyncRoot)
         {
-            if (state.IsCompletionStarted || token.IsCancellationRequested)
+            if (state.IsCompletionStarted || token.IsCancellationRequested || state.ChatGenerationRequest is null)
                 return false;
 
             var aiText = state.ChatGenerationRequest.AIText;
@@ -360,9 +414,13 @@ public sealed class AIJobService(SettingsManager settingsManager, MessageBus mes
     {
         lock (state.SyncRoot)
         {
+            //
+            // A released request keeps its last known title: the job is done, so there is nothing
+            // left to read a newer one from.
+            //
             state.Snapshot = state.Snapshot with
             {
-                Title = state.ChatGenerationRequest.ChatThread.Name,
+                Title = state.ChatGenerationRequest?.ChatThread.Name ?? state.Snapshot.Title,
                 UpdatedAt = DateTimeOffset.Now,
             };
         }
@@ -376,8 +434,12 @@ public sealed class AIJobService(SettingsManager settingsManager, MessageBus mes
         if (!force && now - state.LastCheckpoint < CHECKPOINT_MIN_TIME)
             return;
 
+        var request = state.ChatGenerationRequest;
+        if (request is null)
+            return;
+
         state.LastCheckpoint = now;
-        await WorkspaceBehaviour.StoreChatAsync(state.ChatGenerationRequest.ChatThread);
+        await WorkspaceBehaviour.StoreChatAsync(request.ChatThread);
     }
 
     private static bool ModelsMatch(Model modelA, Model modelB)

app/MindWork AI Studio/Tools/FileExtractionErrorCode.cs

@@ -83,4 +83,11 @@ public enum FileExtractionErrorCode
     /// The extraction finished without reporting a failure, but produced no content at all.
     /// </summary>
     NO_CONTENT,
+
+    /// <summary>
+    /// The caller no longer needs the content, e.g. because the user closed the dialog which
+    /// asked for it. This is not a failure: nobody has to be told about it, which is why there
+    /// is no user-facing message for this code.
+    /// </summary>
+    CANCELLED,
 }

app/MindWork AI Studio/Tools/MessageBus.cs

@@ -93,22 +93,47 @@ public sealed class MessageBus
 
     public Task SendInfo(DataInfoMessage dataInfoMessage) => this.SendMessage(null, Event.SHOW_INFO, dataInfoMessage);
 
+    /// <summary>
+    /// Stores a message until someone asks for it, cf. TakeDeferredMessages. This is how a
+    /// component hands data to a component which does not exist yet, e.g. an assistant which
+    /// sends its result to the chat before the user gets there.
+    /// </summary>
+    /// <param name="sendingComponent">That's you, the sender.</param>
+    /// <param name="triggeredEvent">The event this message belongs to.</param>
+    /// <param name="data">The data to hand over.</param>
     public void DeferMessage<T>(ComponentBase? sendingComponent, Event triggeredEvent, T? data = default)
     {
-        if (this.deferredMessages.TryGetValue(triggeredEvent, out var queue))
-            queue.Enqueue(new Message(sendingComponent, triggeredEvent, data));
-        else
-        {
-            this.deferredMessages[triggeredEvent] = new();
-            this.deferredMessages[triggeredEvent].Enqueue(new Message(sendingComponent, triggeredEvent, data));
-        }
+        var queue = this.deferredMessages.GetOrAdd(triggeredEvent, _ => new());
+        queue.Enqueue(new Message(sendingComponent, triggeredEvent, data));
     }
-    
-    public IEnumerable<T?> CheckDeferredMessages<T>(Event triggeredEvent)
+
+    /// <summary>
+    /// Takes all deferred messages of an event out of the bus.
+    /// </summary>
+    /// <remarks>
+    /// This empties the queue and returns what was in it. It used to be a lazy iterator, which
+    /// meant that a caller stopping after the first message left the rest of the queue behind:
+    /// those messages were never delivered, and the data they carry — a complete chat thread, for
+    /// instance — stayed alive for as long as the app ran. Returning a list makes that impossible.
+    /// Callers who expect a single message take the last one, since that is the most recent thing
+    /// the user asked for.
+    /// </remarks>
+    /// <param name="triggeredEvent">The event whose messages you want.</param>
+    /// <returns>The deferred messages, oldest first. Empty when there are none.</returns>
+    public IReadOnlyList<T?> TakeDeferredMessages<T>(Event triggeredEvent)
     {
-        if (this.deferredMessages.TryGetValue(triggeredEvent, out var queue))
-            while (queue.TryDequeue(out var message))
-                yield return message.Data is T data ? data : default;
+        //
+        // Removing the queue along with its messages is what keeps the dictionary from growing:
+        // otherwise, every event which ever deferred a message would keep an empty queue forever.
+        //
+        if (!this.deferredMessages.TryRemove(triggeredEvent, out var queue))
+            return [];
+
+        var messages = new List<T?>();
+        while (queue.TryDequeue(out var message))
+            messages.Add(message.Data is T data ? data : default);
+
+        return messages;
     }
     
     public async Task<TResult?> SendMessageUseFirstResult<TPayload, TResult>(ComponentBase? sendingComponent, Event triggeredEvent, TPayload? data = default)

app/MindWork AI Studio/Tools/PluginSystem/PluginBase.cs

@@ -6,7 +6,7 @@ namespace AIStudio.Tools.PluginSystem;
 /// <summary>
 /// Represents the base of any AI Studio plugin.
 /// </summary>
-public abstract partial class PluginBase : IPluginMetadata
+public abstract partial class PluginBase : IPluginMetadata, IDisposable
 {
     private static string TB(string fallbackEN) => I18N.I.T(fallbackEN, typeof(PluginBase).Namespace, nameof(PluginBase));
     
@@ -546,4 +546,18 @@ public abstract partial class PluginBase : IPluginMetadata
     }
 
     #endregion
+
+    #region Implementation of IDisposable
+
+    /// <summary>
+    /// Releases the Lua runtime of this plugin.
+    /// </summary>
+    /// <remarks>
+    /// Every plugin owns a Lua state, which is an entire scripting runtime. Dropping a plugin
+    /// without disposing it leaves that runtime behind: before this existed, each hot reload added
+    /// another set of them for as long as the app was running.
+    /// </remarks>
+    public void Dispose() => this.State.Dispose();
+
+    #endregion
 }

app/MindWork AI Studio/Tools/PluginSystem/PluginFactory.HotReload.cs

@@ -21,6 +21,16 @@ public static partial class PluginFactory
         AutoReset = false,
     };
 
+    /// <summary>
+    /// Whether hot reloading was set up already.
+    /// </summary>
+    /// <remarks>
+    /// The timer and the watcher are static, while this method is called from a component. Calling
+    /// it twice would add a second handler to each of them, and every change in the plugins
+    /// directory would then trigger as many reloads as there were calls.
+    /// </remarks>
+    private static bool IS_HOT_RELOADING_SET_UP;
+
     public static void SetUpHotReloading()
     {
         if (!IsInitialized)
@@ -29,6 +39,14 @@ public static partial class PluginFactory
             return;
         }
 
+        if (IS_HOT_RELOADING_SET_UP)
+        {
+            LOG.LogInformation("Hot reloading is already set up. Skipping.");
+            return;
+        }
+
+        IS_HOT_RELOADING_SET_UP = true;
+
         LOG.LogInformation($"Start hot reloading plugins for path '{HOT_RELOAD_WATCHER.Path}'.");
         try
         {

app/MindWork AI Studio/Tools/PluginSystem/PluginFactory.Remove.cs

@@ -69,8 +69,13 @@ public static partial class PluginFactory
             AVAILABLE_PLUGINS.Remove(plugin);
 
             if (RUNNING_PLUGINS.FirstOrDefault(runningPlugin => runningPlugin.Id == plugin.Id) is { } runningPluginToRemove)
+            {
                 RUNNING_PLUGINS.Remove(runningPluginToRemove);
 
+                // The plugin is unloaded, so its Lua runtime is of no use anymore:
+                runningPluginToRemove.Dispose();
+            }
+
             LOG.LogInformation("Unloaded the plugin '{PluginName}' ({PluginId}). Reason: {Reason}.", plugin.Name, plugin.Id, reason);
         }
 

app/MindWork AI Studio/Tools/PluginSystem/PluginFactory.Starting.cs

@@ -18,6 +18,15 @@ public static partial class PluginFactory
     {
         LOG.LogInformation("Try to start or restart all plugins.");
         var configObjects = new List<PluginConfigurationObject>();
+
+        //
+        // Dropping the plugins is not enough: each one owns a Lua runtime, which we have to release
+        // ourselves. Otherwise, every restart — above all every hot reload during development —
+        // leaves another set of runtimes behind:
+        //
+        foreach (var runningPlugin in RUNNING_PLUGINS)
+            runningPlugin.Dispose();
+
         RUNNING_PLUGINS.Clear();
 
         //

app/MindWork AI Studio/Tools/Services/RustService.Retrieval.cs

@@ -16,7 +16,19 @@ public sealed partial class RustService
     /// </remarks>
     private static readonly TimeSpan EXTRACTION_TIMEOUT = TimeSpan.FromMinutes(10);
 
-    public async Task<FileExtractionResult> ReadArbitraryFileData(string path, int maxChunks, bool extractImages = false)
+    /// <summary>
+    /// Reads the content of an arbitrary file through the Rust runtime.
+    /// </summary>
+    /// <param name="path">The path of the file to read.</param>
+    /// <param name="maxChunks">How many chunks of the content stream we read at most.</param>
+    /// <param name="extractImages">Whether we want the images of the file as well.</param>
+    /// <param name="token">
+    /// Cancels the extraction when the caller no longer needs the content. Reading a large document
+    /// takes a while, and without this, the runtime would keep streaming into a caller which is
+    /// already gone.
+    /// </param>
+    /// <returns>The result of reading the file.</returns>
+    public async Task<FileExtractionResult> ReadArbitraryFileData(string path, int maxChunks, bool extractImages = false, CancellationToken token = default)
     {
         //
         // The runtime filters prompt injections while it streams the file. Doing it there rather
@@ -28,8 +40,13 @@ public sealed partial class RustService
         var streamId = Guid.NewGuid().ToString();
         var requestUri = $"/retrieval/fs/extract?path={Uri.EscapeDataString(path)}&stream_id={streamId}&extract_images={extractImages}";
 
+        //
+        // Both reasons to stop end the same read, so we combine them: our own timeout bounds the
+        // operation, and the caller's token ends it as soon as nobody needs the content anymore.
+        //
         using var timeoutTokenSource = new CancellationTokenSource(EXTRACTION_TIMEOUT);
-        var cancellationToken = timeoutTokenSource.Token;
+        using var cancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(timeoutTokenSource.Token, token);
+        var cancellationToken = cancellationTokenSource.Token;
 
         var resultBuilder = new StringBuilder();
         var failedPages = new List<int>();
@@ -162,6 +179,16 @@ public sealed partial class RustService
                 }
             }
         }
+        catch (OperationCanceledException) when (token.IsCancellationRequested)
+        {
+            //
+            // The caller dropped out, e.g. because the user closed the dialog which asked for this
+            // file. That is not a failure, so we log it as information and leave it to the caller
+            // to stay silent about it.
+            //
+            this.logger?.LogInformation("Reading the file '{Path}' was cancelled by the caller.", path);
+            return FileExtractionResult.Failed(FileExtractionErrorCode.CANCELLED, "The caller cancelled reading the file.");
+        }
         catch (OperationCanceledException) when (timeoutTokenSource.IsCancellationRequested)
         {
             this.logger?.LogError("Reading the file '{Path}' timed out after {Timeout}.", path, EXTRACTION_TIMEOUT);

app/MindWork AI Studio/Tools/UserFile.cs

@@ -22,8 +22,9 @@ public static class UserFile
     /// <param name="filePath">The full path to the file to be read. Must not be null or empty.</param>
     /// <param name="rustService">Rust service used to read file content.</param>
     /// <param name="dialogService">Dialogservice used to display the Pandoc installation dialog if needed.</param>
+    /// <param name="token">Cancels the extraction when the caller no longer needs the content.</param>
     /// <returns>The result of reading the file.</returns>
-    public static async Task<FileExtractionResult> LoadFileData(string filePath, RustService rustService, IDialogService dialogService)
+    public static async Task<FileExtractionResult> LoadFileData(string filePath, RustService rustService, IDialogService dialogService, CancellationToken token = default)
     {
         if (string.IsNullOrEmpty(filePath))
         {
@@ -61,7 +62,15 @@ public static class UserFile
             }
         }
 
-        var result = await rustService.ReadArbitraryFileData(filePath, int.MaxValue);
+        var result = await rustService.ReadArbitraryFileData(filePath, int.MaxValue, token: token);
+
+        //
+        // Nobody wants to read that their own cancellation failed. We hand the result back so the
+        // caller can tell the two apart, but we report nothing to the user:
+        //
+        if (result.ErrorCode is FileExtractionErrorCode.CANCELLED)
+            return result;
+
         if (!result.HasUsableContent)
         {
             LOGGER.LogError("Reading the file '{FilePath}' failed: code={ErrorCode}, message='{ErrorMessage}'.", filePath, result.ErrorCode, result.ErrorMessage);

app/MindWork AI Studio/Tools/WorkspaceBehaviour.cs

@@ -82,11 +82,23 @@ public static class WorkspaceBehaviour
 
     private static readonly string TEMPORARY_CHATS_ROOT_DIRECTORY = Path.Join(SettingsManager.DataDirectory, "tempChats");
 
-    private static SemaphoreSlim GetChatSemaphore(Guid workspaceId, Guid chatId)
-    {
-        var key = $"{workspaceId}_{chatId}";
-        return CHAT_STORAGE_SEMAPHORES.GetOrAdd(key, _ => new SemaphoreSlim(1, 1));
-    }
+    private static string ChatSemaphoreKey(Guid workspaceId, Guid chatId) => $"{workspaceId}_{chatId}";
+
+    private static SemaphoreSlim GetChatSemaphore(Guid workspaceId, Guid chatId) =>
+        CHAT_STORAGE_SEMAPHORES.GetOrAdd(ChatSemaphoreKey(workspaceId, chatId), _ => new SemaphoreSlim(1, 1));
+
+    /// <summary>
+    /// Drops the storage semaphore of a chat which does not exist anymore.
+    /// </summary>
+    /// <remarks>
+    /// Deleting the chat is the one moment where we know that nobody will ask for this semaphore
+    /// again; without this, the dictionary would keep one entry per chat the app ever touched. We
+    /// do not dispose the semaphore, though: another operation might still be waiting on it, and
+    /// disposing it under their feet would turn a deleted chat into an exception somewhere else.
+    /// The garbage collector takes care of it once the last waiter is gone.
+    /// </remarks>
+    private static void ForgetChatSemaphore(Guid workspaceId, Guid chatId) =>
+        CHAT_STORAGE_SEMAPHORES.TryRemove(ChatSemaphoreKey(workspaceId, chatId), out _);
 
     private static async Task<(bool Acquired, SemaphoreSlim Semaphore)> TryAcquireChatSemaphoreAsync(Guid workspaceId, Guid chatId, string callerName)
     {
@@ -1114,6 +1126,7 @@ public static class WorkspaceBehaviour
         finally
         {
             semaphore.Release();
+            ForgetChatSemaphore(workspaceId, chatId);
         }
     }