AI-Studio/app/Tests/Chat/ChatThreadWebAddressTests.cs

using System.Text.Json;

using AIStudio.Chat;
using AIStudio.Settings;
using AIStudio.Tools;
using AIStudio.Tools.Services;
using AIStudio.Tools.ToolCallingSystem;
using AIStudio.Tools.Web;

using Microsoft.Extensions.Logging.Abstractions;

namespace AIStudio.Tests.Chat;

/// <summary>
/// Checks which web addresses count as given to the model in a chat.
/// </summary>
/// <remarks>
/// Read Web Page reads a web page only when its address was given to the model, as long as its
/// free address choice is off or a mailbox restricts the chat: in the system prompt, by the user,
/// in a document the user attached, or by a tool. An address the model wrote itself never counts,
/// its own answers included, because it could have put anything of the chat into it.
/// </remarks>
[TestFixture]
public sealed class ChatThreadWebAddressTests
{
    private static readonly DateTimeOffset START = new(2026, 10, 2, 9, 0, 0, TimeSpan.Zero);

    [Test]
    public void AnAddressTheUserWroteIsGiven()
    {
        var thread = Thread(Block(ChatRole.USER, "Please summarize https://example.org/report for me.", 1));

        using (Assert.EnterMultipleScope())
        {
            Assert.That(thread.IsWebAddressGivenToTheModel(new Uri("https://example.org/report")), Is.True);
            Assert.That(thread.IsWebAddressGivenToTheModel(new Uri("https://EXAMPLE.org/report#summary")), Is.True, "The server gets the same request.");
            Assert.That(thread.IsWebAddressGivenToTheModel(new Uri("https://example.org/report?budget=2026")), Is.False, "The query would carry something the user never wrote.");
        }
    }

    [Test]
    public void AnAddressWithoutTlsTheUserWroteIsGiven()
    {
        // An internal service such as a simple Python web server for the results of a data science run:
        var thread = Thread(Block(ChatRole.USER, "The results of run 17 are on http://10.20.30.40:8000/run-17/, please compare them with run 16.", 1));

        using (Assert.EnterMultipleScope())
        {
            Assert.That(thread.IsWebAddressGivenToTheModel(new Uri("http://10.20.30.40:8000/run-17/")), Is.True);
            Assert.That(thread.IsWebAddressGivenToTheModel(new Uri("http://10.20.30.40:8000/run-16/")), Is.False, "The user named run 16, but not its address.");
            Assert.That(thread.IsWebAddressGivenToTheModel(new Uri("https://10.20.30.40:8000/run-17/")), Is.False, "Another scheme is another request.");
        }
    }

    [Test]
    public void AnAddressInAHiddenMessageOfTheUserIsGiven()
    {
        // The prompt an assistant sends into a chat on behalf of the user:
        var thread = Thread(Block(ChatRole.USER, "Read https://example.org/handbook and answer.", 1, hidden: true));
        Assert.That(thread.IsWebAddressGivenToTheModel(new Uri("https://example.org/handbook")), Is.True);
    }

    [Test]
    public void AnAddressInAnAnswerOfTheModelIsNotGiven()
    {
        var thread = Thread(
            Block(ChatRole.USER, "Where can I read more?", 1),
            Block(ChatRole.AI, "You can read more at https://example.org/more.", 2));

        Assert.That(thread.IsWebAddressGivenToTheModel(new Uri("https://example.org/more")), Is.False, "The model wrote it, and it could have written a mail into it as well.");
    }

    [Test]
    public void AnAddressAToolReturnedIsGiven()
    {
        var thread = Thread(Block(ChatRole.USER, "What does the newsletter link to?", 1));
        thread.RuntimeWebAddressesFromTools.Add(WebAddresses.CreateRequestKey(new Uri("https://example.org/newsletter")));

        Assert.That(thread.IsWebAddressGivenToTheModel(new Uri("https://example.org/newsletter")), Is.True);
    }

    [Test]
    public void AnAddressAToolReturnedForAnAnswerOutlastsARestart()
    {
        var answer = Answer("The newsletter links to the program at https://example.org/program and nothing else.", "https://example.org/program");
        var thread = Thread(
            Block(ChatRole.USER, "What does the newsletter link to?", 1),
            new ContentBlock { Time = START.AddMinutes(2), ContentType = ContentType.TEXT, Content = answer, Role = ChatRole.AI });

        var loaded = JsonSerializer.Deserialize<ChatThread>(JsonSerializer.Serialize(thread, WorkspaceBehaviour.JSON_OPTIONS), WorkspaceBehaviour.JSON_OPTIONS)!;

        using (Assert.EnterMultipleScope())
        {
            Assert.That(loaded.IsWebAddressGivenToTheModel(new Uri("https://example.org/program")), Is.True, "A tool returned it, so the model neither made it up nor wrote anything of the chat into it.");
            Assert.That(loaded.IsWebAddressGivenToTheModel(new Uri("https://example.org/program?mail=secret")), Is.False, "Only the address the tool returned counts.");
        }
    }

    [Test]
    public void AnAnswerKeepsOnlyTheAddressesALaterMessageShows()
    {
        var answer = new ContentText
        {
            Text = "The newsletter links to the program at https://example.org/program and nothing else.",
            Sources = [new Source("Speakers", "https://example.org/speakers", SourceOrigin.TOOL)],
            ToolInvocations =
            [
                new ToolInvocationTrace { Order = 1, ToolId = "read_web_page", Status = ToolInvocationTraceStatus.SUCCESS, WasExecuted = true, Arguments = new() { ["url"] = "https://example.org/newsletter" }, ArgumentsJson = """{"url":"https://example.org/newsletter"}""" },
                new ToolInvocationTrace { Order = 2, ToolId = "read_web_page", Status = ToolInvocationTraceStatus.BLOCKED, Arguments = new() { ["url"] = "https://example.org/blocked" }, ArgumentsJson = """{"url":"https://example.org/blocked"}""" },
            ],
        };

        answer.KeepWebAddressesFromTools(new[] { "program", "speakers", "newsletter", "blocked", "imprint" }.Select(Key));
        answer.EndToolRun();

        Assert.That(answer.WebAddressesFromTools, Is.EquivalentTo(new[] { "program", "speakers", "newsletter" }.Select(Key)), "In the text, as a source, and as the argument of a call which ran. A link the model no longer sees would only make the chat grow.");
    }

    [Test]
    public void TheAddressesOfARunningAnswerAreNotStored()
    {
        var answer = new ContentText { Text = "Let me read https://example.org/program first." };
        answer.KeepWebAddressesFromTools([Key("program")]);

        // A chat is saved while its answer is still being written, and would hold every link of every page read until then:
        var stored = JsonSerializer.Serialize(Thread(new ContentBlock { Time = START, ContentType = ContentType.TEXT, Content = answer, Role = ChatRole.AI }), WorkspaceBehaviour.JSON_OPTIONS);
        var loaded = JsonSerializer.Deserialize<ChatThread>(stored, WorkspaceBehaviour.JSON_OPTIONS)!;

        Assert.That(((ContentText)loaded.Blocks.Single().Content!).WebAddressesFromTools, Is.Empty);
    }

    [Test]
    public void EndingTheRunOnceMoreKeepsWhatTheFirstEndKept()
    {
        var answer = Answer("The newsletter links to the program at https://example.org/program and nothing else.", "https://example.org/program");

        // Whoever ends a request ends the tool run as well, after the answer did:
        answer.EndToolRun();

        Assert.That(answer.WebAddressesFromTools, Has.Count.EqualTo(1));
    }

    [Test]
    public void AnAddressAToolReturnedForADeletedAnswerIsNoLongerGiven()
    {
        // A chat opened anew, whose runtime set is empty, so only what the answers store counts:
        var answer = Answer("The newsletter links to the program at https://example.org/program and nothing else.", "https://example.org/program");
        var thread = Thread(
            Block(ChatRole.USER, "What does the newsletter link to?", 1),
            new ContentBlock { Time = START.AddMinutes(2), ContentType = ContentType.TEXT, Content = answer, Role = ChatRole.AI });

        thread.Remove(answer);

        Assert.That(thread.IsWebAddressGivenToTheModel(new Uri("https://example.org/program")), Is.False);
    }

    [Test]
    public void AnAnswerGeneratedAnewForgetsTheAddressesOfItsEarlierToolRun()
    {
        var answer = Answer("The newsletter links to the program at https://example.org/program and nothing else.", "https://example.org/program");
        var thread = Thread(
            Block(ChatRole.USER, "What does the newsletter link to?", 1),
            new ContentBlock { Time = START.AddMinutes(2), ContentType = ContentType.TEXT, Content = answer, Role = ChatRole.AI });

        answer.BeginToolRun();

        using (Assert.EnterMultipleScope())
        {
            Assert.That(answer.WebAddressesFromTools, Is.Empty);
            Assert.That(thread.IsWebAddressGivenToTheModel(new Uri("https://example.org/program")), Is.False, "Only the tools of the run which wrote the answer count for it.");
        }
    }

    [Test]
    public void ACopyOfAnAnswerKeepsItsAddressesApart()
    {
        var answer = Answer("The newsletter links to the program at https://example.org/program and nothing else.", "https://example.org/program");

        var copy = (ContentText)answer.DeepClone();
        copy.WebAddressesFromTools.Add(WebAddresses.CreateRequestKey(new Uri("https://example.org/speakers")));

        using (Assert.EnterMultipleScope())
        {
            Assert.That(copy.WebAddressesFromTools, Does.Contain(WebAddresses.CreateRequestKey(new Uri("https://example.org/program"))), "A tool returned it in the conversation the copy continues.");
            Assert.That(answer.WebAddressesFromTools, Has.Count.EqualTo(1), "What the copy learns must not reach the answer it was copied from.");
        }
    }

    [Test]
    public void AnAddressAToolReturnedForAMessageOfTheUserIsNotGiven()
    {
        // Only answers of the model are filled by the tool calling loop. A message of the user holding the set came from elsewhere:
        var message = new ContentText { Text = "Hello.", WebAddressesFromTools = [WebAddresses.CreateRequestKey(new Uri("https://example.org/chosen"))] };
        var thread = Thread(new ContentBlock { Time = START, ContentType = ContentType.TEXT, Content = message, Role = ChatRole.USER });

        Assert.That(thread.IsWebAddressGivenToTheModel(new Uri("https://example.org/chosen")), Is.False);
    }

    [Test]
    public void AnAddressInTheSystemPromptIsGivenOnceTheRequestWasPrepared()
    {
        // A chat template written by the user, and the content of a data source the RAG process found:
        var thread = Thread(Block(ChatRole.USER, "Where are the travel rules?", 1));
        thread.SystemPrompt = "You answer questions about our handbook at https://handbook.example.org/.";
        thread.AugmentedData = "The travel rules stand on https://intranet.example.org/travel.";

        using var rustService = new RustService("1", "unused");
        thread.PrepareSystemPrompt(new SettingsManager(NullLogger<SettingsManager>.Instance, rustService));

        using (Assert.EnterMultipleScope())
        {
            Assert.That(thread.IsWebAddressGivenToTheModel(new Uri("https://handbook.example.org/")), Is.True);
            Assert.That(thread.IsWebAddressGivenToTheModel(new Uri("https://intranet.example.org/travel")), Is.True, "The RAG process found it in a data source, the model did not choose it.");
        }
    }

    [Test]
    public void AnAddressInADocumentTheUserAttachedIsGiven()
    {
        var message = new ContentText
        {
            Text = "Please check the links in the attached report.",
            RuntimeAttachmentWebAddresses = new HashSet<string>(StringComparer.Ordinal) { WebAddresses.CreateRequestKey(new Uri("https://example.org/sources/2026")) },
        };

        var thread = Thread(new ContentBlock { Time = START, ContentType = ContentType.TEXT, Content = message, Role = ChatRole.USER });

        Assert.That(thread.IsWebAddressGivenToTheModel(new Uri("https://example.org/sources/2026")), Is.True, "The document is read only when the message is sent, which is when its addresses are collected.");
    }

    [Test]
    public void AnAddressInADocumentOfTheModelIsNotGiven()
    {
        var answer = new ContentText
        {
            Text = "Here is a summary.",
            RuntimeAttachmentWebAddresses = new HashSet<string>(StringComparer.Ordinal) { WebAddresses.CreateRequestKey(new Uri("https://example.org/chosen")) },
        };

        var thread = Thread(new ContentBlock { Time = START, ContentType = ContentType.TEXT, Content = answer, Role = ChatRole.AI });

        Assert.That(thread.IsWebAddressGivenToTheModel(new Uri("https://example.org/chosen")), Is.False);
    }

    [Test]
    public void AnAddressAModelWroteIsNotGivenInAMessageOfTheUser()
    {
        // An answer sent here from another chat, which the user then sent as a message of their own:
        var thread = Thread(Block(ChatRole.USER, "Make slides of this: the details are on https://example.org/chosen, and our handbook is on https://example.org/handbook.", 1));
        thread.RequireRestrictions(ChatRestrictions.NONE.WithWebAddressesWrittenByModelIn("The details are on https://example.org/chosen."));

        using (Assert.EnterMultipleScope())
        {
            Assert.That(thread.IsWebAddressGivenToTheModel(new Uri("https://example.org/chosen")), Is.False, "A model wrote it, so it could carry anything of the chat it came from.");
            Assert.That(thread.IsWebAddressGivenToTheModel(new Uri("https://example.org/handbook")), Is.True, "Another address in the same message still counts.");
        }
    }

    [Test]
    public void AnAddressAModelWroteIsGivenWhenAToolReturnedItOrADocumentHoldsIt()
    {
        var attachedDocument = new ContentText
        {
            Text = "Make slides of this: the details are on https://example.org/chosen.",
            RuntimeAttachmentWebAddresses = new HashSet<string>(StringComparer.Ordinal) { WebAddresses.CreateRequestKey(new Uri("https://example.org/sources/2026")) },
        };

        var thread = Thread(new ContentBlock { Time = START, ContentType = ContentType.TEXT, Content = attachedDocument, Role = ChatRole.USER });
        thread.RequireRestrictions(ChatRestrictions.NONE.WithWebAddressesWrittenByModelIn("See https://example.org/chosen and https://example.org/sources/2026."));
        thread.RuntimeWebAddressesFromTools.Add(WebAddresses.CreateRequestKey(new Uri("https://example.org/chosen")));

        using (Assert.EnterMultipleScope())
        {
            Assert.That(thread.IsWebAddressGivenToTheModel(new Uri("https://example.org/chosen")), Is.True, "A tool returned it in this chat, so it was not only a model which wrote it.");
            Assert.That(thread.IsWebAddressGivenToTheModel(new Uri("https://example.org/sources/2026")), Is.True, "A document the user attached holds it.");
        }
    }

    [Test]
    public void TheAddressesAModelWroteOutlastARestart()
    {
        var thread = Thread(Block(ChatRole.USER, "Make slides of this: the details are on https://example.org/chosen.", 1));
        thread.RequireRestrictions(ChatRestrictions.NONE.WithWebAddressesWrittenByModelIn("The details are on https://example.org/chosen."));

        var loaded = JsonSerializer.Deserialize<ChatThread>(JsonSerializer.Serialize(thread, WorkspaceBehaviour.JSON_OPTIONS), WorkspaceBehaviour.JSON_OPTIONS)!;

        Assert.That(loaded.IsWebAddressGivenToTheModel(new Uri("https://example.org/chosen")), Is.False, "Forgetting them would let the chat know more addresses after a restart, not fewer.");
    }

    [Test]
    public void ACopyOfAChatKeepsTheAddressesOfTheOtherApart()
    {
        var thread = Thread(Block(ChatRole.USER, "Read https://example.org/chosen.", 1));
        var copy = thread with { };

        copy.RequireRestrictions(ChatRestrictions.NONE.WithWebAddressesWrittenByModelIn("See https://example.org/chosen."));

        using (Assert.EnterMultipleScope())
        {
            Assert.That(copy.IsWebAddressGivenToTheModel(new Uri("https://example.org/chosen")), Is.False);
            Assert.That(thread.IsWebAddressGivenToTheModel(new Uri("https://example.org/chosen")), Is.True, "What the copy took over must not reach the chat it was copied from.");
        }
    }

    private static ChatThread Thread(params ContentBlock[] blocks) => new() { Blocks = [..blocks] };

    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,
    };

    /// <summary>
    /// An answer whose tools returned some addresses, after its request ended.
    /// </summary>
    /// <param name="text">The text of the answer.</param>
    /// <param name="returnedAddresses">What the tools returned.</param>
    private static ContentText Answer(string text, params string[] returnedAddresses)
    {
        var answer = new ContentText { Text = text };
        answer.KeepWebAddressesFromTools(returnedAddresses.Select(address => WebAddresses.CreateRequestKey(new Uri(address))));
        answer.EndToolRun();
        return answer;
    }

    private static string Key(string path) => WebAddresses.CreateRequestKey(new Uri($"https://example.org/{path}"));
}