build-flow.api/tests/BuildFlow.Rag.Tests/Chat/ChatContextWindowBuilderTests.cs

using System.Text;
using BuildFlow.Application.Chat.Commands;
using BuildFlow.Application.Chat.Handlers;
using BuildFlow.Application.Chat.Services;
using BuildFlow.Domain.Entities.ChatMessage;
using BuildFlow.Domain.Entities.Conversation;
using BuildFlow.Domain.Enums;
using BuildFlow.Infrastructure.Cache;
using BuildFlow.Infrastructure.Persistence;
using BuildFlow.Rag.Tests.Stub;
using BuildFlow.Shared.Infrastructure;
using FluentAssertions;
using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.Options;
using Xunit;
using Xunit.Abstractions;

namespace BuildFlow.Rag.Tests.Chat;

/// <summary>聊天生成上下文的消息数、Token 和三条生成路径复用测试。</summary>
public sealed class ChatContextWindowBuilderTests(ITestOutputHelper output)
{
    [Fact]
    public async Task ReturnsNewestMessagesWithinMessageLimitInChronologicalOrder()
    {
        await using var db = CreateSeedDb();
        var conversationId = Guid.NewGuid();
        SeedMessages(db, conversationId, 100);
        await db.SaveChangesAsync();
        var builder = CreateBuilder(db, maxMessages: 6, maxTokens: 10_000);

        var result = await builder.BuildAsync(conversationId, CancellationToken.None);

        result.Select(x => x.Text).Should().Equal(
            "message-095",
            "message-096",
            "message-097",
            "message-098",
            "message-099",
            "message-100");
        var estimatedTokens = result.Sum(message => Math.Max(
            1,
            (Encoding.UTF8.GetByteCount(message.Text) + 3) / 4));
        output.WriteLine($"SOP_METRIC rag_context_message_count={result.Count}");
        output.WriteLine($"SOP_METRIC rag_context_estimated_tokens={estimatedTokens}");
    }

    [Fact]
    public async Task StopsBeforeAddingAMessageThatWouldExceedTokenLimit()
    {
        await using var db = CreateSeedDb();
        var conversationId = Guid.NewGuid();
        AddMessage(db, conversationId, "1234", 1);
        AddMessage(db, conversationId, "abcdefghij", 2);
        AddMessage(db, conversationId, "12345678", 3);
        await db.SaveChangesAsync();
        var builder = CreateBuilder(db, maxMessages: 10, maxTokens: 5);

        var result = await builder.BuildAsync(conversationId, CancellationToken.None);

        result.Select(x => x.Text).Should().Equal("abcdefghij", "12345678");
    }

    [Fact]
    public async Task SendHandlerUsesBoundedContextBuilder()
    {
        var setup = await CreateHandlerDbAsync("send");
        await using var db = setup.Db;
        var agent = new RecordingChatAgent();
        var builder = CreateBuilder(db, maxMessages: 3, maxTokens: 10_000);
        var handler = new SendMessageCommandHandler(
            db,
            agent,
            new StubChatMessageCache(),
            builder,
            setup.UserContext);

        await handler.Handle(
            new SendMessageCommand(setup.ConversationId, "new-user"),
            CancellationToken.None);

        agent.Calls.Should().ContainSingle();
        agent.Calls[0].Select(x => x.Text).Should().Equal(
            "message-003",
            "message-004",
            "new-user");
    }

    [Fact]
    public async Task StreamHandlerUsesBoundedContextBuilder()
    {
        var setup = await CreateHandlerDbAsync("stream");
        await using var db = setup.Db;
        var agent = new RecordingChatAgent();
        var builder = CreateBuilder(db, maxMessages: 2, maxTokens: 10_000);
        var handler = new StreamMessageCommandHandler(
            db,
            agent,
            new StubChatMessageCache(),
            builder,
            setup.UserContext);

        await CollectAsync(handler.Handle(
            new StreamMessageCommand(setup.ConversationId, "stream-user"),
            CancellationToken.None));

        agent.Calls.Should().ContainSingle();
        agent.Calls[0].Select(x => x.Text).Should().Equal("message-004", "stream-user");
    }

    [Fact]
    public async Task EditHandlerUsesBoundedContextBuilderAfterTruncation()
    {
        var setup = await CreateHandlerDbAsync("edit");
        await using var db = setup.Db;
        var target = await db.ChatMessages.SingleAsync(x => x.Content == "message-003");
        var agent = new RecordingChatAgent();
        var builder = CreateBuilder(db, maxMessages: 2, maxTokens: 10_000);
        var handler = new EditMessageStreamCommandHandler(
            db,
            agent,
            new StubChatMessageCache(),
            builder,
            setup.UserContext);

        await CollectAsync(handler.Handle(
            new EditMessageStreamCommand(setup.ConversationId, target.Id, "edited-user"),
            CancellationToken.None));

        agent.Calls.Should().ContainSingle();
        agent.Calls[0].Select(x => x.Text).Should().Equal("message-002", "edited-user");
    }

    private static ChatContextWindowBuilder CreateBuilder(
        RagDbContext db,
        int maxMessages,
        int maxTokens)
        => new(
            db,
            Options.Create(new ChatOptions
            {
                MaxContextMessages = maxMessages,
                MaxContextTokens = maxTokens,
                CacheTtlMinutes = 60,
            }));

    private static RagDbContext CreateSeedDb()
    {
        var options = new DbContextOptionsBuilder<RagDbContext>()
            .UseInMemoryDatabase($"chat-window-seed-{Guid.NewGuid():N}")
            .Options;
        return new RagDbContext(options);
    }

    private static async Task<(
        RagDbContext Db,
        Guid ConversationId,
        TestCurrentUserContext UserContext)> CreateHandlerDbAsync(string name)
    {
        var databaseName = $"chat-window-{name}-{Guid.NewGuid():N}";
        var userContext = new TestCurrentUserContext();
        var conversationId = Guid.NewGuid();
        await using (var seedDb = new RagDbContext(
                         new DbContextOptionsBuilder<RagDbContext>()
                             .UseInMemoryDatabase(databaseName)
                             .Options))
        {
            seedDb.Conversations.Add(new ConversationEntity
            {
                Id = conversationId,
                Title = "window",
                UserId = Guid.Parse(userContext.UserId),
                CreatedBy = userContext.UserId,
                UpdatedBy = userContext.UserId,
                OperatorIP = userContext.IPAddress,
            });
            SeedMessages(seedDb, conversationId, 4);
            await seedDb.SaveChangesAsync();
        }

        var options = new DbContextOptionsBuilder<RagDbContext>()
            .UseInMemoryDatabase(databaseName)
            .AddInterceptors(new AuditInterceptor(userContext))
            .Options;
        return (
            new RagDbContext(options),
            conversationId,
            userContext);
    }

    private static void SeedMessages(RagDbContext db, Guid conversationId, int count)
    {
        for (var index = 1; index <= count; index++)
        {
            AddMessage(db, conversationId, $"message-{index:000}", index);
        }
    }

    private static void AddMessage(
        RagDbContext db,
        Guid conversationId,
        string content,
        int minute)
    {
        db.ChatMessages.Add(new ChatMessageEntity
        {
            ConversationId = conversationId,
            Role = minute % 2 == 0 ? ChatRoleEnum.Assistant : ChatRoleEnum.User,
            Content = content,
            CreatedBy = "test",
            UpdatedBy = "test",
            CreatedAt = new DateTime(2026, 1, 1, 0, 0, 0, DateTimeKind.Utc).AddMinutes(minute),
            UpdatedAt = new DateTime(2026, 1, 1, 0, 0, 0, DateTimeKind.Utc).AddMinutes(minute),
        });
    }

    private static async Task<List<T>> CollectAsync<T>(IAsyncEnumerable<T> source)
    {
        var result = new List<T>();
        await foreach (var item in source)
        {
            result.Add(item);
        }

        return result;
    }
}