Playa/tests/Volcanic.AgenticCore.Telegram.Tests/MessageContextChatTests.cs

using System.Text.RegularExpressions;
using Microsoft.Extensions.Time.Testing;
using Telegram.Bot.Requests;
using Telegram.Bot.Types;
using Volcanic.AgenticCore.Abstractions.Messages;
using Volcanic.AgenticCore.Telegram.Inbound;
using Volcanic.AgenticCore.Telegram.Tests.Support;
using Volcanic.AgenticCore.Testing;

namespace Volcanic.AgenticCore.Telegram.Tests;

/// <summary>Reply/forward headers, the timestamp line and grouping, end to end through the router and the faux model.</summary>
public sealed class MessageContextChatTests
{
    private static readonly TimeZoneInfo Tokyo = TimeZoneInfo.FindSystemTimeZoneById("Asia/Tokyo");
    private static readonly Regex TimestampLine = new(@"^\[(Mon|Tue|Wed|Thu|Fri|Sat|Sun) \d{4}-\d{2}-\d{2} \d{2}:\d{2} [^\]]+\]$");

    private static DateTimeOffset SentAt(Update update) => new(DateTime.SpecifyKind(update.Message!.Date, DateTimeKind.Utc));

    private static string Prompt(ChatHost host) => host.Faux.Requests.Single().Messages.Last().Text;

    /// <summary>Wait until <paramref name="parts"/> parts are composed, then let the (fake) window elapse.</summary>
    internal static async Task FlushAsync(ChatHost host, FakeTimeProvider time, int parts)
    {
        var chat = await host.ChatAsync();
        await Wait.UntilAsync(() => chat.Inbox.Buffered == parts && chat.Inbox.OpenLeases == 0, $"{parts} composed part(s)");
        time.Advance(host.Options.InputDebounce);
    }

    [Fact]
    public async Task A_reply_to_a_bot_message_reaches_the_model_with_the_header()
    {
        await using var host = new ChatHost(core: o => o.DefaultTimeZone = Tokyo);
        host.Faux.Enqueue(FauxTurn.Say("That one."));
        var reply = Updates.Reply("what about this?", Updates.BotMessage("Which plant do you mean?"));

        await host.RouteAsync(reply);
        await host.SettleAsync();

        Assert.Equal(
            $"{MessageContext.TimestampHeader(SentAt(reply), Tokyo)}\n[Reply to your message: \"Which plant do you mean?\"]\nwhat about this?",
            Prompt(host));
    }

    [Fact]
    public async Task The_first_line_is_one_timestamp_header_in_the_chat_zone()
    {
        var time = new FakeTimeProvider(DateTimeOffset.UtcNow);
        await using var host = new ChatHost(core: o => o.DefaultTimeZone = Tokyo, time: time);
        host.Faux.Enqueue(FauxTurn.Say("Yes."));
        var comment = Updates.Text("look at this");
        var forward = Updates.Forward("Geothermal plant opens", Updates.Channel("Daily News"));

        await host.RouteAsync(comment);
        await host.RouteAsync(forward);
        await FlushAsync(host, time, parts: 2);
        await host.SettleAsync();

        // The zone comes from the ChatRegistry.TimeZoneFor seam (the chat's zone from ChatSettingsStore, which defaults to AgenticCoreOptions.DefaultTimeZone).
        Assert.Equal("Asia/Tokyo", Volcanic.AgenticCore.Hosting.TimeZoneIds.ToIana(host.Chats.TimeZoneFor(Updates.AllowedUser)));
        var lines = Prompt(host).Split('\n');
        Assert.Equal(MessageContext.TimestampHeader(SentAt(comment), Tokyo), lines[0]);
        Assert.EndsWith(" Asia/Tokyo]", lines[0]);
        Assert.Single(lines, l => TimestampLine.IsMatch(l)); // one per batch, not per part
    }

    [Fact]
    public async Task After_timezone_Asia_Tokyo_the_next_messages_header_says_Tokyo()
    {
        var time = new FakeTimeProvider(DateTimeOffset.UtcNow);
        await using var host = new ChatHost(core: o => o.DefaultTimeZone = TimeZoneInfo.FindSystemTimeZoneById("Europe/Berlin"), time: time);
        host.Faux.Enqueue(FauxTurn.Say("It's evening in Tokyo."));

        await host.SendAsync("/timezone Asia/Tokyo"); // a bot command: handled at once, no inbox window
        var message = Updates.Text("what time is it?");
        await host.RouteAsync(message);
        await FlushAsync(host, time, parts: 1);
        await host.SettleAsync();

        var first = Prompt(host).Split('\n')[0];
        Assert.Equal(MessageContext.TimestampHeader(SentAt(message), Tokyo), first);
        Assert.EndsWith(" Asia/Tokyo]", first);
    }

    [Fact]
    public async Task Plain_text_waits_for_the_window_and_a_forward_inside_it_joins_the_request()
    {
        var time = new FakeTimeProvider(DateTimeOffset.UtcNow);
        await using var host = new ChatHost(time: time);
        host.Faux.Enqueue(FauxTurn.Say("No."));

        await host.SendAsync("is this true?");
        var chat = await host.ChatAsync();
        await Wait.UntilAsync(() => chat.Inbox.Buffered == 1 && chat.Inbox.OpenLeases == 0, "the text part");
        time.Advance(host.Options.InputDebounce - TimeSpan.FromMilliseconds(1));
        Assert.Equal(1, chat.Inbox.Buffered); // still waiting
        Assert.Empty(host.Faux.Requests);

        await host.RouteAsync(Updates.Forward("The moon is made of cheese.", Updates.UserOrigin("ada", "Ada", "Lovelace")));
        await FlushAsync(host, time, parts: 2);
        await host.SettleAsync();

        Assert.EndsWith("\nis this true?\n\n[Forwarded from @ada]\nThe moon is made of cheese.", Prompt(host));
    }

    [Fact]
    public async Task Comment_and_forward_sent_back_to_back_are_one_request_in_order()
    {
        var time = new FakeTimeProvider(DateTimeOffset.UtcNow);
        await using var host = new ChatHost(time: time);
        host.Faux.Enqueue(FauxTurn.Say("Yes, that one."));
        var comment = Updates.Text("look at this, is it the plant you mentioned?");
        var forward = Updates.Forward("Geothermal plant opens", Updates.Channel("Daily News", "Anna K."));

        await host.RouteAsync(comment);
        await host.RouteAsync(forward);
        await FlushAsync(host, time, parts: 2);
        await host.SettleAsync();

        Assert.Equal(
            $"{MessageContext.TimestampHeader(SentAt(comment), TimeZoneInfo.Local)}\nlook at this, is it the plant you mentioned?\n\n[Forwarded from Daily News (Anna K.)]\nGeothermal plant opens",
            Prompt(host));
    }

    [Fact]
    public async Task A_three_photo_album_with_one_caption_is_one_request_with_three_images_and_one_forward_label()
    {
        var time = new FakeTimeProvider(DateTimeOffset.UtcNow);
        await using var host = new ChatHost(time: time);
        host.Faux.Enqueue(FauxTurn.Say("Nice views."));
        var channel = Updates.Channel("Daily News");

        await host.RouteAsync(Updates.AlbumPhoto("g1", "three views", channel));
        await host.RouteAsync(Updates.AlbumPhoto("g1", forwardedFrom: channel));
        await host.RouteAsync(Updates.AlbumPhoto("g1", forwardedFrom: channel));
        await FlushAsync(host, time, parts: 3);
        await host.SettleAsync();

        var message = Assert.IsType<UserMessage>(host.Faux.Requests.Single().Messages.Last());
        Assert.Equal(3, message.Content.OfType<ImageContent>().Count());
        Assert.Single(Regex.Matches(message.Text, Regex.Escape("[Forwarded from Daily News]")));
        Assert.Contains("three views", message.Text);
    }

    [Fact]
    public async Task Stop_is_handled_while_a_photo_download_is_blocked()
    {
        await using var host = new ChatHost();
        var download = new TaskCompletionSource();
        host.Bot.OnSend = r => { if (r is GetFileRequest) download.Task.Wait(TimeSpan.FromSeconds(10)); };
        host.Faux.Enqueue(FauxTurn.Say("too late").WithDelay(TimeSpan.FromSeconds(30)), FauxTurn.Say("A photo."));

        await host.SendAsync("take your time");
        var chat = await host.ChatAsync();
        await Wait.UntilAsync(() => chat.Session.IsStreaming, "the run to start");

        await host.RouteAsync(Updates.Photo()).WaitAsync(TimeSpan.FromSeconds(2)); // returns while the download hangs
        await Wait.UntilAsync(() => host.Bot.Requests.OfType<GetFileRequest>().Any(), "the download to start");
        await host.SendAsync("/stop").WaitAsync(TimeSpan.FromSeconds(2));
        await Wait.UntilAsync(() => !chat.Session.IsStreaming, "the run to stop");
        Assert.False(download.Task.IsCompleted);

        download.SetResult();
        await host.SettleAsync();
        Assert.Contains("⏹ Stopped.", host.Texts);
    }

    [Fact]
    public async Task A_message_sent_during_a_run_gets_eyes_and_steers_it()
    {
        var started = new TaskCompletionSource();
        var gate = new TaskCompletionSource();
        await using var host = new ChatHost(core: o => o.AddExtension(new GateToolExtension(started, gate)));
        host.Faux.Enqueue(FauxTurn.Tool(GateToolExtension.ToolName, new { }), FauxTurn.Say("Done."));

        await host.SendAsync("start the long job");
        await started.Task.WaitAsync(TimeSpan.FromSeconds(5));
        var steer = Updates.Text("also check the logs");
        await host.RouteAsync(steer);
        await Wait.UntilAsync(() => host.Bot.Reactions.Count == 1, "the 👀 reaction");
        gate.SetResult();
        await host.SettleAsync();

        var reaction = Assert.Single(host.Bot.Reactions); // the first message started the run: no reaction
        Assert.Equal(steer.Message!.Id, reaction.MessageId);
        Assert.Equal("👀", Assert.IsType<ReactionTypeEmoji>(Assert.Single(reaction.Reaction!)).Emoji);
        var messages = host.Faux.Requests[1].Messages.ToList(); // the model call after the tool
        var toolResult = messages.FindLastIndex(m => m is ToolResultMessage);
        var steered = messages.FindIndex(m => m is UserMessage && m.Text.EndsWith("]\nalso check the logs", StringComparison.Ordinal));
        Assert.True(steered > toolResult, "the message is delivered as steering after the tool result");
    }

    [Fact]
    public async Task A_message_sent_while_idle_gets_no_reaction()
    {
        await using var host = new ChatHost();
        host.Faux.Enqueue(FauxTurn.Say("Hi."));

        await host.SendAsync("hello");
        await host.SettleAsync();

        Assert.Empty(host.Bot.Reactions);
    }

    [Fact]
    public async Task New_during_the_window_drops_the_buffered_text()
    {
        var time = new FakeTimeProvider(DateTimeOffset.UtcNow);
        await using var host = new ChatHost(time: time);

        await host.SendAsync("forget me");
        var chat = await host.ChatAsync();
        await Wait.UntilAsync(() => chat.Inbox.Buffered == 1 && chat.Inbox.OpenLeases == 0, "the buffered text");
        await host.SendAsync("/new");
        Assert.Equal(0, chat.Inbox.Buffered);

        time.Advance(host.Options.InputDebounce * 2);
        await host.SettleAsync();
        Assert.Empty(host.Faux.Requests);
        Assert.Contains("🆕 New conversation.", host.Texts);
    }

    [Fact]
    public async Task Resume_during_the_window_drops_the_buffered_text()
    {
        var time = new FakeTimeProvider(DateTimeOffset.UtcNow);
        await using var host = new ChatHost(time: time);
        host.Faux.Enqueue(FauxTurn.Say("Hi."));
        await host.SendAsync("hello");
        await FlushAsync(host, time, parts: 1);
        await host.SettleAsync();
        var chat = await host.ChatAsync();
        var earlier = chat.Session.SessionId;
        await host.SendAsync("/new");

        await host.SendAsync("forget me");
        await Wait.UntilAsync(() => chat.Inbox.Buffered == 1 && chat.Inbox.OpenLeases == 0, "the buffered text");
        await host.SendAsync($"/resume {earlier}");
        Assert.Equal(0, chat.Inbox.Buffered);

        time.Advance(host.Options.InputDebounce * 2);
        await host.SettleAsync();
        Assert.Single(host.Faux.Requests);
        Assert.Equal(earlier, chat.Session.SessionId);
        Assert.Contains(host.Texts, t => t.StartsWith("Resumed", StringComparison.Ordinal) && t.Contains(earlier, StringComparison.Ordinal));
    }

    [Fact]
    public async Task Stop_during_the_window_drops_the_buffered_text()
    {
        var time = new FakeTimeProvider(DateTimeOffset.UtcNow);
        await using var host = new ChatHost(time: time);

        await host.SendAsync("never mind");
        var chat = await host.ChatAsync();
        await Wait.UntilAsync(() => chat.Inbox.Buffered == 1 && chat.Inbox.OpenLeases == 0, "the buffered text");
        await host.SendAsync("/stop");
        Assert.Equal(0, chat.Inbox.Buffered);

        time.Advance(host.Options.InputDebounce * 2);
        await host.SettleAsync();
        Assert.Empty(host.Faux.Requests);
    }

    [Fact]
    public async Task A_forwarded_command_is_content_not_a_command()
    {
        await using var host = new ChatHost();
        host.Faux.Enqueue(FauxTurn.Say("Noted."));
        await host.SendAsync("/status"); // creates the chat
        var chat = await host.ChatAsync();
        var session = chat.Session.SessionId;

        await host.RouteAsync(Updates.Forward("/new beginnings", Updates.Channel("Daily News")));
        await host.SettleAsync();

        Assert.Equal(session, chat.Session.SessionId);
        Assert.DoesNotContain("🆕 New conversation.", host.Texts);
        Assert.EndsWith("\n[Forwarded from Daily News]\n/new beginnings", Prompt(host), StringComparison.Ordinal);
    }

    [Fact]
    public async Task A_forwarded_text_while_input_is_pending_goes_to_the_agent()
    {
        await using var host = new ChatHost();
        host.Faux.Enqueue(FauxTurn.Say("Interesting."));
        await host.SendAsync("/status");
        var chat = await host.ChatAsync();
        var answer = chat.Ui.InputAsync("Your name?", "name", null, CancellationToken.None);
        await Wait.UntilAsync(() => chat.Ui.IsWaitingForInput, "the question");

        await host.RouteAsync(Updates.Forward("Ada Lovelace", Updates.UserOrigin("ada", "Ada", "Lovelace")));
        await host.SettleAsync();

        Assert.True(chat.Ui.IsWaitingForInput);
        Assert.False(answer.IsCompleted);
        Assert.EndsWith("\n[Forwarded from @ada]\nAda Lovelace", Prompt(host));
        await host.SendAsync("Ada");
        Assert.Equal("Ada", await answer.WaitAsync(TimeSpan.FromSeconds(5)));
    }

    [Fact]
    public async Task Unsupported_messages_get_the_reply_and_take_no_lease()
    {
        await using var host = new ChatHost();

        await host.RouteAsync(Updates.StickerMessage());

        Assert.Equal("I can read text, photos, files and voice messages.", Assert.Single(host.Bot.Sent).Text);
        Assert.Null(host.Chats.Find(Updates.AllowedUser));
        Assert.Empty(host.Faux.Requests);
    }
}