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

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);