Improved how AI Studio saves its files (#1201)

app/MindWork AI Studio/Tools/AtomicFile.cs

@@ -0,0 +1,196 @@
+using System.Text;
+using System.Text.Json;
+using System.Text.RegularExpressions;
+
+namespace AIStudio.Tools;
+
+/// <summary>
+/// Writes files so that readers see either the old or the new content, never a part of it.
+/// </summary>
+/// <remarks>
+/// The new content goes into a temporary file next to the target first, which then replaces the target by a rename.
+/// The temporary file has to live in the directory of the target: a rename across a file system boundary falls back
+/// to copying, which is exactly what this avoids, and the temporary directory of the operating system is such another
+/// file system under Flatpak. The temporary file is named after its target, followed by a GUID and ".tmp", so that
+/// a crash leaves a file behind which DeleteTemporaryFiles, DataSourceTranscriptStore, and MediaSidecarTranscripts
+/// recognize.<br/><br/>
+/// This works only where AI Studio may create files next to the target: in its data and configuration directories,
+/// or in a folder the user chose as a whole, e.g., the output folder of the Batch Processing assistant. It does not
+/// work for a single file the user chose in a save or an open dialog or dropped onto AI Studio: under Flatpak, the
+/// document portal shares only that one file, so nothing can be created next to it. Prepare such a file in the
+/// temporary directory and write it into the chosen file with FileMode.Create, as PresentationFile and
+/// PluginShareService do. Keep in mind, too, that the Flatpak sees the home directory read-only (--filesystem=home:ro).
+/// <br/><br/>
+/// The rename protects against a crash of AI Studio. Against a power failure, the content has to reach the disk
+/// before the rename, which flushToDisk asks for. That costs time on every write, especially on macOS and on SD cards,
+/// so only files which are hard to get back use it.
+/// </remarks>
+public static partial class AtomicFile
+{
+    private const int BUFFER_SIZE = 65_536;
+
+    /// <summary>
+    /// Writes a file, unless the writer decides to discard what it wrote.
+    /// </summary>
+    /// <param name="targetPath">The file to write. Its directory has to exist.</param>
+    /// <param name="write">Writes the new content into the stream and returns whether it replaces the target.</param>
+    /// <param name="token">The cancellation token.</param>
+    /// <param name="overwrite">Whether an existing target is replaced; otherwise, an existing target fails the write with an IOException.</param>
+    /// <param name="flushToDisk">Whether the content has to reach the disk before it replaces the target.</param>
+    /// <param name="logger">Where to report a temporary file which cannot be removed.</param>
+    /// <returns>True when the new content replaced the target, false when the writer discarded it.</returns>
+    public static async Task<bool> WriteOrDiscardAsync(string targetPath, Func<Stream, CancellationToken, Task<bool>> write, CancellationToken token, bool overwrite = true, bool flushToDisk = false, ILogger? logger = null)
+    {
+        var temporaryPath = $"{targetPath}.{Guid.NewGuid():N}.tmp";
+        try
+        {
+            bool keep;
+            await using (var stream = new FileStream(temporaryPath, FileMode.CreateNew, FileAccess.Write, FileShare.None, BUFFER_SIZE, FileOptions.Asynchronous))
+            {
+                keep = await write(stream, token);
+                if (keep && flushToDisk)
+                    await FlushToDiskAsync(stream, token);
+            }
+
+            if (!keep)
+            {
+                TryDelete(temporaryPath, logger);
+                return false;
+            }
+
+            File.Move(temporaryPath, targetPath, overwrite);
+            return true;
+        }
+        catch
+        {
+            // Deleting never throws, so the original error is the one which reaches the caller:
+            TryDelete(temporaryPath, logger);
+            throw;
+        }
+    }
+
+    /// <summary>
+    /// Writes a file.
+    /// </summary>
+    /// <param name="targetPath">The file to write. Its directory has to exist.</param>
+    /// <param name="write">Writes the new content into the stream.</param>
+    /// <param name="token">The cancellation token.</param>
+    /// <param name="overwrite">Whether an existing target is replaced; otherwise, an existing target fails the write with an IOException.</param>
+    /// <param name="flushToDisk">Whether the content has to reach the disk before it replaces the target.</param>
+    /// <param name="logger">Where to report a temporary file which cannot be removed.</param>
+    public static Task WriteAsync(string targetPath, Func<Stream, CancellationToken, Task> write, CancellationToken token, bool overwrite = true, bool flushToDisk = false, ILogger? logger = null) => WriteOrDiscardAsync(targetPath, async (stream, writeToken) =>
+    {
+        await write(stream, writeToken);
+        return true;
+    }, token, overwrite, flushToDisk, logger);
+
+    /// <summary>
+    /// Writes a text file.
+    /// </summary>
+    /// <remarks>
+    /// The file starts with the byte order mark of the encoding, if it has one, just like with File.WriteAllTextAsync:
+    /// Encoding.UTF8 writes one, new UTF8Encoding(false) does not.
+    /// </remarks>
+    /// <param name="targetPath">The file to write. Its directory has to exist.</param>
+    /// <param name="content">The new content.</param>
+    /// <param name="encoding">The encoding of the file.</param>
+    /// <param name="token">The cancellation token.</param>
+    /// <param name="overwrite">Whether an existing target is replaced; otherwise, an existing target fails the write with an IOException.</param>
+    /// <param name="flushToDisk">Whether the content has to reach the disk before it replaces the target.</param>
+    /// <param name="logger">Where to report a temporary file which cannot be removed.</param>
+    public static Task WriteTextAsync(string targetPath, string content, Encoding encoding, CancellationToken token, bool overwrite = true, bool flushToDisk = false, ILogger? logger = null) =>
+        WriteAsync(targetPath, (stream, writeToken) => WriteTextAsync(stream, content, encoding, writeToken), token, overwrite, flushToDisk, logger);
+
+    /// <summary>
+    /// Writes a text into the stream of a file, for a writer passed to WriteAsync or WriteOrDiscardAsync.
+    /// </summary>
+    /// <remarks>
+    /// The stream gets the byte order mark of the encoding, if it has one, just like a file written by File.WriteAllTextAsync.
+    /// </remarks>
+    /// <param name="stream">The stream of the file.</param>
+    /// <param name="content">The text.</param>
+    /// <param name="encoding">The encoding of the file.</param>
+    /// <param name="token">The cancellation token.</param>
+    public static async Task WriteTextAsync(Stream stream, string content, Encoding encoding, CancellationToken token)
+    {
+        await using var writer = new StreamWriter(stream, encoding, BUFFER_SIZE, leaveOpen: true);
+        await writer.WriteAsync(content.AsMemory(), token);
+        await writer.FlushAsync(token);
+    }
+
+    /// <summary>
+    /// Writes a JSON file, as UTF-8 without a byte order mark.
+    /// </summary>
+    /// <remarks>
+    /// The value goes into the file piece by piece, so that the complete JSON is never in memory at once. The file
+    /// holds the same bytes as when the value was serialized into a string and written by File.WriteAllTextAsync.
+    /// </remarks>
+    /// <typeparam name="T">The type of the value.</typeparam>
+    /// <param name="targetPath">The file to write. Its directory has to exist.</param>
+    /// <param name="value">The value to write.</param>
+    /// <param name="options">The JSON options.</param>
+    /// <param name="token">The cancellation token.</param>
+    /// <param name="overwrite">Whether an existing target is replaced; otherwise, an existing target fails the write with an IOException.</param>
+    /// <param name="flushToDisk">Whether the content has to reach the disk before it replaces the target.</param>
+    /// <param name="logger">Where to report a temporary file which cannot be removed.</param>
+    public static Task WriteJsonAsync<T>(string targetPath, T value, JsonSerializerOptions options, CancellationToken token, bool overwrite = true, bool flushToDisk = false, ILogger? logger = null) =>
+        WriteAsync(targetPath, (stream, writeToken) => JsonSerializer.SerializeAsync(stream, value, options, writeToken), token, overwrite, flushToDisk, logger);
+
+    /// <summary>
+    /// Deletes the temporary files which a crash left behind anywhere below a directory.
+    /// </summary>
+    /// <remarks>
+    /// Besides the temporary files of this class, it also deletes those which AI Studio named ".tmp-" followed by a
+    /// GUID before, so that the leftovers of older versions disappear as well. Call it only for a directory which
+    /// belongs to AI Studio alone.
+    /// </remarks>
+    /// <param name="rootDirectory">The directory to clean up.</param>
+    /// <param name="logger">Where to report a temporary file which cannot be removed.</param>
+    public static void DeleteTemporaryFiles(string rootDirectory, ILogger? logger = null)
+    {
+        foreach (var path in Directory.EnumerateFiles(rootDirectory, "*.tmp*", SearchOption.AllDirectories))
+        {
+            var fileName = Path.GetFileName(path);
+            if (TemporaryFileName().IsMatch(fileName) || PreviousTemporaryFileName().IsMatch(fileName))
+                TryDelete(path, logger);
+        }
+    }
+
+    /// <summary>
+    /// Makes sure the content of a file reached the disk.
+    /// </summary>
+    /// <remarks>
+    /// FlushAsync hands the buffered content to the operating system only; just Flush(true) waits for the disk, and
+    /// that only synchronously. It runs on the thread pool, so that it does not block the caller, e.g., the renderer.
+    /// </remarks>
+    /// <param name="stream">The stream of the file.</param>
+    /// <param name="token">The cancellation token.</param>
+    private static async Task FlushToDiskAsync(FileStream stream, CancellationToken token)
+    {
+        await stream.FlushAsync(token);
+        await Task.Run(() => stream.Flush(flushToDisk: true), token);
+    }
+
+    /// <summary>
+    /// Deletes a file, without failing.
+    /// </summary>
+    /// <param name="path">The file to delete; a file which does not exist is fine.</param>
+    /// <param name="logger">Where to report a file which cannot be removed.</param>
+    private static void TryDelete(string path, ILogger? logger)
+    {
+        try
+        {
+            File.Delete(path);
+        }
+        catch (Exception exception) when (exception is IOException or UnauthorizedAccessException)
+        {
+            logger?.LogWarning(exception, "Could not remove the temporary file '{TemporaryPath}'.", path);
+        }
+    }
+
+    [GeneratedRegex(@"\.[0-9a-f]{32}\.tmp$")]
+    private static partial Regex TemporaryFileName();
+
+    [GeneratedRegex(@"\.tmp-[0-9a-f]{32}$")]
+    private static partial Regex PreviousTemporaryFileName();
+}

app/MindWork AI Studio/Tools/Media/PersistentMediaTranscript.cs

@@ -119,7 +119,6 @@ internal static class PersistentMediaTranscript
         if (!sourceBefore.StillMatches(mediaPath))
             return SourceChanged(slot);
 
-        var temporaryPath = $"{slot.Path}.{Guid.NewGuid():N}.tmp";
         try
         {
             var directory = Path.GetDirectoryName(slot.Path);
@@ -128,13 +127,18 @@ internal static class PersistentMediaTranscript
 
             //
             // The transcription is paid for at this point. Storing it is quick, so it is completed
-            // even when the caller cancelled meanwhile, and the next attempt can use it:
+            // even when the caller cancelled meanwhile, and the next attempt can use it. A transcript
+            // of a media file which changed meanwhile is discarded instead of replacing the old one:
             //
-            await File.WriteAllTextAsync(temporaryPath, transcription.Text, new UTF8Encoding(false), CancellationToken.None);
-            if (!sourceBefore.StillMatches(mediaPath))
+            var stored = await AtomicFile.WriteOrDiscardAsync(slot.Path, async (stream, writeToken) =>
+            {
+                await AtomicFile.WriteTextAsync(stream, transcription.Text, new UTF8Encoding(false), writeToken);
+                return sourceBefore.StillMatches(mediaPath);
+            }, CancellationToken.None, logger: logger);
+
+            if (!stored)
                 return SourceChanged(slot);
 
-            File.Move(temporaryPath, slot.Path, true);
             logger.LogInformation("Stored the transcript '{TranscriptFilePath}' of the media file '{MediaFilePath}'.", slot.Path, mediaPath);
             return PersistentMediaTranscriptResult.Ready(slot, reused: false);
         }
@@ -143,18 +147,6 @@ internal static class PersistentMediaTranscript
             logger.LogWarning(exception, "Could not store the transcript '{TranscriptFilePath}' of the media file '{MediaFilePath}'.", slot.Path, mediaPath);
             return PersistentMediaTranscriptResult.Unavailable(PersistentMediaTranscriptStatus.STORE_FAILED, slot, string.Format(TB("Was not able to store the transcript: {0}"), exception.Message));
         }
-        finally
-        {
-            try
-            {
-                if (File.Exists(temporaryPath))
-                    File.Delete(temporaryPath);
-            }
-            catch (Exception exception)
-            {
-                logger.LogWarning(exception, "Could not remove the temporary transcript '{TempFilePath}'.", temporaryPath);
-            }
-        }
     }
 
     /// <summary>

app/MindWork AI Studio/Tools/WorkspaceBehaviour.RelativeDataPaths.cs

@@ -81,19 +81,8 @@ public static partial class WorkspaceBehaviour
                 return true;
 
             var lastEditTime = File.GetLastWriteTimeUtc(chatPath);
-            var temporaryPath = Path.Join(chatDirectory, $".thread-{Guid.NewGuid():N}.tmp");
-            try
-            {
-                await File.WriteAllTextAsync(temporaryPath, json, Encoding.UTF8, token);
-                File.Move(temporaryPath, chatPath, overwrite: true);
-                File.SetLastWriteTimeUtc(chatPath, lastEditTime);
-            }
-            finally
-            {
-                if (File.Exists(temporaryPath))
-                    File.Delete(temporaryPath);
-            }
-
+            await AtomicFile.WriteTextAsync(chatPath, json, Encoding.UTF8, token, logger: LOG);
+            File.SetLastWriteTimeUtc(chatPath, lastEditTime);
             return true;
         }
         catch (Exception exception) when (exception is JsonException or NotSupportedException)

app/MindWork AI Studio/Tools/WorkspaceBehaviour.cs

@@ -938,13 +938,9 @@ public static partial class WorkspaceBehaviour
             Directory.CreateDirectory(chatDirectory);
 
             FinalizeStagedTranscripts(chat, chatDirectory);
-            var chatNamePath = Path.Join(chatDirectory, "name");
-            await File.WriteAllTextAsync(chatNamePath, chat.Name);
+            await StoreChatFilesAsync(chat, chatDirectory);
 
-            var chatPath = Path.Join(chatDirectory, "thread.json");
-            await File.WriteAllTextAsync(chatPath, JsonSerializer.Serialize(chat, JSON_OPTIONS), Encoding.UTF8);
-
-            var lastEditTime = File.GetLastWriteTimeUtc(chatPath);
+            var lastEditTime = File.GetLastWriteTimeUtc(Path.Join(chatDirectory, "thread.json"));
             await UpdateCacheAfterChatStored(chat.WorkspaceId, chat.ChatId, chatDirectory, chat.Name, lastEditTime);
             return true;
         }
@@ -1174,7 +1170,7 @@ public static partial class WorkspaceBehaviour
             movedChat.WorkspaceId = targetWorkspaceId;
             
             FinalizeStagedTranscripts(movedChat, targetDirectory);
-            await StoreMovedChatFilesAsync(movedChat, targetDirectory);
+            await StoreChatFilesAsync(movedChat, targetDirectory);
 
             var lastEditTime = File.GetLastWriteTimeUtc(Path.Join(targetDirectory, "thread.json"));
             await UpdateCacheAfterChatMoved(sourceWorkspaceId, targetWorkspaceId, chat.ChatId, targetDirectory, movedChat.Name, lastEditTime);
@@ -1202,7 +1198,7 @@ public static partial class WorkspaceBehaviour
                         Directory.Move(targetDirectory, sourceDirectory);
 
                     if (movedChat is not null && Directory.Exists(sourceDirectory))
-                        await StoreMovedChatFilesAsync(movedChat, sourceDirectory);
+                        await StoreChatFilesAsync(movedChat, sourceDirectory);
                 }
                 catch (Exception rollbackError)
                 {
@@ -1229,23 +1225,17 @@ public static partial class WorkspaceBehaviour
         return JsonSerializer.Deserialize<ChatThread>(chatData, JSON_OPTIONS) ?? throw new InvalidDataException($"The stored chat could not be read: '{chatDirectory}'.");
     }
 
-    /// <summary>Atomically stores the name and thread after a directory move.</summary>
-    private static async Task StoreMovedChatFilesAsync(ChatThread chat, string chatDirectory)
+    /// <summary>Stores the name and the thread of a chat into its directory.</summary>
+    /// <remarks>
+    /// Both files are written atomically: a crash of AI Studio while storing leaves each of them as it was before,
+    /// never half written. The thread keeps its UTF-8 byte order mark, the name has none, as before.
+    /// </remarks>
+    /// <param name="chat">The chat to store.</param>
+    /// <param name="chatDirectory">The directory of the chat, which has to exist.</param>
+    private static async Task StoreChatFilesAsync(ChatThread chat, string chatDirectory)
     {
-        await File.WriteAllTextAsync(Path.Join(chatDirectory, "name"), chat.Name);
-        var chatPath = Path.Join(chatDirectory, "thread.json");
-        var temporaryPath = Path.Join(chatDirectory, $".thread-{Guid.NewGuid():N}.tmp");
-        
-        try
-        {
-            await File.WriteAllTextAsync(temporaryPath, JsonSerializer.Serialize(chat, JSON_OPTIONS), Encoding.UTF8);
-            File.Move(temporaryPath, chatPath, overwrite: true);
-        }
-        finally
-        {
-            if (File.Exists(temporaryPath))
-                File.Delete(temporaryPath);
-        }
+        await AtomicFile.WriteTextAsync(Path.Join(chatDirectory, "name"), chat.Name, new UTF8Encoding(false), CancellationToken.None, logger: LOG);
+        await AtomicFile.WriteTextAsync(Path.Join(chatDirectory, "thread.json"), JsonSerializer.Serialize(chat, JSON_OPTIONS), Encoding.UTF8, CancellationToken.None, logger: LOG);
     }
 
     /// <summary>Rewrites absolute file paths after moving the complete chat directory.</summary>