Merge branch 'main' into feature/token-count-from-usage

app/Tests/Chat/ChatThreadRollbackTests.cs

@@ -0,0 +1,225 @@
+using AIStudio.Agents;
+using AIStudio.Chat;
+using AIStudio.Components;
+using AIStudio.Provider;
+using AIStudio.Settings.DataModel;
+
+namespace AIStudio.Tests.Chat;
+
+/// <summary>
+/// Checks what rolling a chat back to an earlier answer removes and what it keeps.
+/// </summary>
+/// <remarks>
+/// A rollback has to remove more than the user sees. The prompts an assistant sends into a chat are
+/// hidden, yet they are part of the conversation; a chat which kept them would continue with
+/// messages the user can neither see nor remove. The example conversation of a chat template is
+/// hidden as well, but it comes before everything else and has to keep working.
+///
+/// What a rollback must not remove is anything which tells the providers what this chat has seen.
+/// Removing the message which brought confidential data in does not unsee it, so the chat keeps
+/// demanding a self-hosted provider, and the confidence which that data asked for.
+/// </remarks>
+[TestFixture]
+public sealed class ChatThreadRollbackTests
+{
+    private static readonly DateTimeOffset START = new(2026, 9, 23, 10, 0, 0, TimeSpan.Zero);
+
+    [Test]
+    public void HiddenPromptsAfterTheAnswerGoAsWell()
+    {
+        var answer = Block(ChatRole.AI, "First answer", 2);
+        var thread = new ChatThread
+        {
+            Blocks =
+            [
+                Block(ChatRole.USER, "First question", 1),
+                answer,
+                Block(ChatRole.USER, "Prompt of an assistant", 3, hidden: true),
+                Block(ChatRole.AI, "Answer to the assistant", 4),
+            ],
+        };
+
+        var rolledBack = thread.RollBackTo(answer.Content!);
+
+        Assert.Multiple(() =>
+        {
+            Assert.That(rolledBack, Is.True, "Two blocks came after the answer, so the rollback removed something.");
+            Assert.That(Texts(thread), Is.EqualTo(new[] { "First question", "First answer" }), "The hidden prompt goes with the answer to it, although the user never saw it.");
+        });
+    }
+
+    [Test]
+    public void TheExampleConversationOfAChatTemplateStays()
+    {
+        var answer = Block(ChatRole.AI, "First answer", 4);
+        var thread = new ChatThread
+        {
+            Blocks =
+            [
+                Block(ChatRole.USER, "Example question", 1, hidden: true),
+                Block(ChatRole.AI, "Example answer", 2, hidden: true),
+                Block(ChatRole.USER, "First question", 3),
+                answer,
+                Block(ChatRole.USER, "Second question", 5),
+                Block(ChatRole.AI, "Second answer", 6),
+            ],
+        };
+
+        thread.RollBackTo(answer.Content!);
+
+        Assert.That(Texts(thread), Is.EqualTo(new[] { "Example question", "Example answer", "First question", "First answer" }), "Hidden blocks before the answer belong to what the user rolls back to, so the chat template keeps its example conversation.");
+    }
+
+    [Test]
+    public void TheOrderIsTheOneOfTheTimeStampsNotTheOneOfTheList()
+    {
+        var answer = Block(ChatRole.AI, "First answer", 2);
+        var thread = new ChatThread
+        {
+            Blocks =
+            [
+                Block(ChatRole.USER, "Second question", 3),
+                Block(ChatRole.AI, "Second answer", 4),
+                answer,
+                Block(ChatRole.USER, "First question", 1),
+            ],
+        };
+
+        thread.RollBackTo(answer.Content!);
+
+        Assert.That(Texts(thread), Is.EqualTo(new[] { "First question", "First answer" }), "The chat shows its blocks by time, so a rollback has to count by time as well, wherever a block sits in the list.");
+    }
+
+    [Test]
+    public void TheLastRetrievalIsDroppedButTheChoiceOfDataSourcesStays()
+    {
+        var answer = Block(ChatRole.AI, "First answer", 2);
+        var options = new DataSourceOptions
+        {
+            DisableDataSources = false,
+            AutomaticDataSourceSelection = true,
+            PreselectedDataSourceIds = ["handbook"],
+        };
+
+        var thread = new ChatThread
+        {
+            DataSourceOptions = options,
+            AISelectedDataSources = [SelectedHandbook()],
+            AugmentedData = "A chunk from the handbook",
+            Blocks =
+            [
+                Block(ChatRole.USER, "First question", 1),
+                answer,
+                Block(ChatRole.USER, "Second question", 3),
+                Block(ChatRole.AI, "Second answer", 4),
+            ],
+        };
+
+        thread.RollBackTo(answer.Content!);
+
+        Assert.Multiple(() =>
+        {
+            Assert.That(thread.AugmentedData, Is.Empty, "Nobody knows whether the retrieved data belongs to a kept or to a removed message, so it must not reach the next system prompt.");
+            Assert.That(thread.AISelectedDataSources, Is.Empty, "The data sources an agent picked belong to the same retrieval as the data.");
+            Assert.That(thread.DataSourceOptions, Is.SameAs(options), "The data source options are the user's choice, not the result of a message.");
+            Assert.That(options.DisableDataSources, Is.False);
+            Assert.That(options.AutomaticDataSourceSelection, Is.True);
+            Assert.That(options.PreselectedDataSourceIds, Is.EqualTo(new[] { "handbook" }));
+        });
+    }
+
+    [Test]
+    public void WhatTheChatHasSeenKeepsRestrictingTheProviders()
+    {
+        var answer = Block(ChatRole.AI, "First answer", 2);
+        var thread = new ChatThread
+        {
+            DataSecurity = DataSourceSecurity.SELF_HOSTED,
+            Blocks =
+            [
+                Block(ChatRole.USER, "First question", 1),
+                answer,
+                Block(ChatRole.USER, "Question about the confidential handbook", 3),
+                Block(ChatRole.AI, "Answer from the confidential handbook", 4),
+            ],
+        };
+
+        thread.RequireProviderConfidence(ConfidenceLevel.HIGH);
+
+        thread.RollBackTo(answer.Content!);
+
+        Assert.Multiple(() =>
+        {
+            Assert.That(thread.DataSecurity, Is.EqualTo(DataSourceSecurity.SELF_HOSTED), "The confidential data was seen by this chat, so no cloud provider may continue it.");
+            Assert.That(thread.RequiredProviderConfidence, Is.EqualTo(ConfidenceLevel.HIGH), "The confidence the data demanded stays, although the message which brought it in is gone.");
+        });
+    }
+
+    [Test]
+    public void RollingBackToTheLastAnswerChangesNothing()
+    {
+        var answer = Block(ChatRole.AI, "First answer", 2);
+        var thread = ThreadWithRetrieval(Block(ChatRole.USER, "First question", 1), answer);
+
+        var rolledBack = thread.RollBackTo(answer.Content!);
+
+        AssertUnchanged(thread, rolledBack, "First question", "First answer");
+    }
+
+    [Test]
+    public void RollingBackToUnknownContentChangesNothing()
+    {
+        var thread = ThreadWithRetrieval(Block(ChatRole.USER, "First question", 1), Block(ChatRole.AI, "First answer", 2));
+
+        var rolledBack = thread.RollBackTo(new ContentText { Text = "Answer of another chat" });
+
+        AssertUnchanged(thread, rolledBack, "First question", "First answer");
+    }
+
+    /// <summary>
+    /// A chat whose last retrieval is still in place, so that a rollback which should do nothing
+    /// shows when it resets it anyway.
+    /// </summary>
+    /// <param name="blocks">The blocks of the chat.</param>
+    /// <returns>The chat.</returns>
+    private static ChatThread ThreadWithRetrieval(params ContentBlock[] blocks) => new()
+    {
+        AISelectedDataSources = [SelectedHandbook()],
+        AugmentedData = "A chunk from the handbook",
+        Blocks = [..blocks],
+    };
+
+    private static void AssertUnchanged(ChatThread thread, bool rolledBack, params string[] expectedTexts)
+    {
+        Assert.Multiple(() =>
+        {
+            Assert.That(rolledBack, Is.False, "Nothing came after the content, so there was nothing to roll back.");
+            Assert.That(Texts(thread), Is.EqualTo(expectedTexts), "A rollback without anything to remove leaves the blocks alone.");
+            Assert.That(thread.AugmentedData, Is.EqualTo("A chunk from the handbook"), "A rollback without anything to remove keeps the last retrieval, because it still belongs to the last message.");
+            Assert.That(thread.AISelectedDataSources, Has.Count.EqualTo(1));
+        });
+    }
+
+    private static ContentBlock Block(ChatRole role, string text, int minute, bool hidden = false) => new()
+    {
+        Time = START.AddMinutes(minute),
+        ContentType = ContentType.TEXT,
+        Content = new ContentText { Text = text },
+        Role = role,
+        HideFromUser = hidden,
+    };
+
+    private static DataSourceAgentSelected SelectedHandbook() => new()
+    {
+        DataSource = new DataSourceLocalDirectory { Id = "handbook", Name = "Handbook" },
+        AIDecision = new SelectedDataSource("handbook", "The question is about the handbook.", 0.9f),
+        Selected = true,
+    };
+
+    /// <summary>
+    /// The texts of the chat in the order the chat shows them.
+    /// </summary>
+    /// <param name="thread">The chat.</param>
+    /// <returns>The texts.</returns>
+    private static string[] Texts(ChatThread thread) => thread.Blocks.OrderBy(x => x.Time).Select(x => ((ContentText)x.Content!).Text).ToArray();
+}

app/Tests/Chat/IContentExtensionsTests.cs

@@ -162,16 +162,85 @@ public sealed class IContentExtensionsTests
             "|---|---|",
             "| Q1 | 100 |"), TOOL_SOURCE);
 
-        content.TryGetMarkdownText(out var markdown);
-        var tables = PlainFileExport.ExtractTables(markdown, ',');
+        var tables = FilesOf(content);
 
         Assert.Multiple(() =>
         {
             Assert.That(tables, Has.Count.EqualTo(1), "One table in the message, one table offered for it.");
-            Assert.That(tables[0].Content, Does.Not.Contain("example.org"), "A data table has no column a link list would fit into.");
+            Assert.That(content.ToExportContent(tables[0]), Is.EqualTo(tables[0].Content), "A data table has no column a link list would fit into.");
         });
     }
 
+    [Test]
+    public void AMarkdownBlockCarriesTheSourcesVisibly()
+    {
+        var content = TextWith(Lines(
+            "Here are your notes:",
+            string.Empty,
+            "```markdown",
+            "# Notes",
+            string.Empty,
+            "The notes.",
+            "```"), TOOL_SOURCE);
+
+        var exported = content.ToExportContent(FilesOf(content).Single());
+
+        Assert.That(TopLevelBlocksOf(exported), Is.EqualTo(new[] { "h1", "ParagraphBlock", "h1", "h2", "ListBlock" }), "The notes without the text around them, followed by the source list just as the entire answer carries it.");
+    }
+
+    [Test]
+    public void AWebPageCarriesTheSourcesInAComment()
+    {
+        var content = TextWith(Lines(
+            "```html",
+            "<!DOCTYPE html>",
+            "<html><body><p>Hello</p></body></html>",
+            "```"), TOOL_SOURCE);
+
+        var file = FilesOf(content).Single();
+        var exported = content.ToExportContent(file);
+        var appended = exported[file.Content.Length..].Trim();
+
+        Assert.Multiple(() =>
+        {
+            Assert.That(exported, Does.StartWith(file.Content), "The page stays as the model wrote it.");
+            Assert.That(appended, Does.StartWith("<!--").And.EndWith("-->"), "Below the page stands one comment and nothing a browser would show.");
+            Assert.That(appended, Does.Contain(TOOL_SOURCE.URL));
+        });
+    }
+
+    [Test]
+    public void ALatexBlockCarriesTheSourcesInComments()
+    {
+        var content = TextWith(Lines(
+            "```latex",
+            @"\section{Results}",
+            "```"), TOOL_SOURCE);
+
+        var file = FilesOf(content).Single();
+        var exported = content.ToExportContent(file);
+        var appended = exported[file.Content.Length..].Trim();
+
+        Assert.Multiple(() =>
+        {
+            Assert.That(exported, Does.StartWith(file.Content), "The document stays as the model wrote it.");
+            Assert.That(appended.Split(Environment.NewLine), Has.All.StartWith("%"), "A line LaTeX would read could stop the whole run.");
+            Assert.That(appended, Does.Contain(TOOL_SOURCE.URL));
+        });
+    }
+
+    [TestCase("markdown")]
+    [TestCase("html")]
+    [TestCase("latex")]
+    public void WithoutSourcesACodeBlockStaysAsTheModelWroteIt(string language)
+    {
+        var content = TextWith(Lines($"```{language}", "The content.", "```"));
+
+        var file = FilesOf(content).Single();
+
+        Assert.That(content.ToExportContent(file), Is.EqualTo(file.Content), "No comment, no heading, no empty line.");
+    }
+
     /// <summary>
     /// A text message with the given sources hanging on it.
     /// </summary>
@@ -184,6 +253,17 @@ public sealed class IContentExtensionsTests
         Sources = [..sources],
     };
 
+    /// <summary>
+    /// Reads the files of a message the way the export menu does.
+    /// </summary>
+    /// <param name="content">The content to read.</param>
+    /// <returns>The files the export menu offers for it.</returns>
+    private static IReadOnlyList<MessageFile> FilesOf(IContent content)
+    {
+        content.TryGetMarkdownText(out var markdown);
+        return PlainFileExport.ExtractFiles(markdown, ',');
+    }
+
     /// <summary>
     /// Names the blocks a Markdown text is made of, headings by their level.
     /// </summary>