// Copyright (c) Microsoft. All rights reserved. using System.Collections.Generic; using Microsoft.Agents.AI.Compaction; using Microsoft.Extensions.AI; namespace Microsoft.Agents.AI.UnitTests.Compaction; /// /// Contains tests for the class. /// public class MessageIndexTests { [Fact] public void CreateEmptyListReturnsEmptyGroups() { // Arrange List messages = []; // Act MessageIndex groups = MessageIndex.Create(messages); // Assert Assert.Empty(groups.Groups); } [Fact] public void CreateSystemMessageCreatesSystemGroup() { // Arrange List messages = [ new ChatMessage(ChatRole.System, "You are helpful."), ]; // Act MessageIndex groups = MessageIndex.Create(messages); // Assert Assert.Single(groups.Groups); Assert.Equal(MessageGroupKind.System, groups.Groups[0].Kind); Assert.Single(groups.Groups[0].Messages); } [Fact] public void CreateUserMessageCreatesUserGroup() { // Arrange List messages = [ new ChatMessage(ChatRole.User, "Hello"), ]; // Act MessageIndex groups = MessageIndex.Create(messages); // Assert Assert.Single(groups.Groups); Assert.Equal(MessageGroupKind.User, groups.Groups[0].Kind); } [Fact] public void CreateAssistantTextMessageCreatesAssistantTextGroup() { // Arrange List messages = [ new ChatMessage(ChatRole.Assistant, "Hi there!"), ]; // Act MessageIndex groups = MessageIndex.Create(messages); // Assert Assert.Single(groups.Groups); Assert.Equal(MessageGroupKind.AssistantText, groups.Groups[0].Kind); } [Fact] public void CreateToolCallWithResultsCreatesAtomicToolCallGroup() { // Arrange ChatMessage assistantMessage = new(ChatRole.Assistant, [new FunctionCallContent("call1", "get_weather", new Dictionary { ["city"] = "Seattle" })]); ChatMessage toolResult = new(ChatRole.Tool, "Sunny, 72°F"); List messages = [assistantMessage, toolResult]; // Act MessageIndex groups = MessageIndex.Create(messages); // Assert Assert.Single(groups.Groups); Assert.Equal(MessageGroupKind.ToolCall, groups.Groups[0].Kind); Assert.Equal(2, groups.Groups[0].Messages.Count); Assert.Same(assistantMessage, groups.Groups[0].Messages[0]); Assert.Same(toolResult, groups.Groups[0].Messages[1]); } [Fact] public void CreateMixedConversationGroupsCorrectly() { // Arrange ChatMessage systemMsg = new(ChatRole.System, "You are helpful."); ChatMessage userMsg = new(ChatRole.User, "What's the weather?"); ChatMessage assistantToolCall = new(ChatRole.Assistant, [new FunctionCallContent("call1", "get_weather")]); ChatMessage toolResult = new(ChatRole.Tool, "Sunny"); ChatMessage assistantText = new(ChatRole.Assistant, "The weather is sunny!"); List messages = [systemMsg, userMsg, assistantToolCall, toolResult, assistantText]; // Act MessageIndex groups = MessageIndex.Create(messages); // Assert Assert.Equal(4, groups.Groups.Count); Assert.Equal(MessageGroupKind.System, groups.Groups[0].Kind); Assert.Equal(MessageGroupKind.User, groups.Groups[1].Kind); Assert.Equal(MessageGroupKind.ToolCall, groups.Groups[2].Kind); Assert.Equal(2, groups.Groups[2].Messages.Count); Assert.Equal(MessageGroupKind.AssistantText, groups.Groups[3].Kind); } [Fact] public void CreateMultipleToolResultsGroupsAllWithAssistant() { // Arrange ChatMessage assistantToolCall = new(ChatRole.Assistant, [ new FunctionCallContent("call1", "get_weather"), new FunctionCallContent("call2", "get_time"), ]); ChatMessage toolResult1 = new(ChatRole.Tool, "Sunny"); ChatMessage toolResult2 = new(ChatRole.Tool, "3:00 PM"); List messages = [assistantToolCall, toolResult1, toolResult2]; // Act MessageIndex groups = MessageIndex.Create(messages); // Assert Assert.Single(groups.Groups); Assert.Equal(MessageGroupKind.ToolCall, groups.Groups[0].Kind); Assert.Equal(3, groups.Groups[0].Messages.Count); } [Fact] public void GetIncludedMessagesExcludesMarkedGroups() { // Arrange ChatMessage msg1 = new(ChatRole.User, "First"); ChatMessage msg2 = new(ChatRole.Assistant, "Response"); ChatMessage msg3 = new(ChatRole.User, "Second"); MessageIndex groups = MessageIndex.Create([msg1, msg2, msg3]); groups.Groups[1].IsExcluded = true; // Act List included = [.. groups.GetIncludedMessages()]; // Assert Assert.Equal(2, included.Count); Assert.Same(msg1, included[0]); Assert.Same(msg3, included[1]); } [Fact] public void GetAllMessagesIncludesExcludedGroups() { // Arrange ChatMessage msg1 = new(ChatRole.User, "First"); ChatMessage msg2 = new(ChatRole.Assistant, "Response"); MessageIndex groups = MessageIndex.Create([msg1, msg2]); groups.Groups[0].IsExcluded = true; // Act List all = [.. groups.GetAllMessages()]; // Assert Assert.Equal(2, all.Count); } [Fact] public void IncludedGroupCountReflectsExclusions() { // Arrange MessageIndex groups = MessageIndex.Create( [ new ChatMessage(ChatRole.User, "A"), new ChatMessage(ChatRole.Assistant, "B"), new ChatMessage(ChatRole.User, "C"), ]); groups.Groups[1].IsExcluded = true; // Act & Assert Assert.Equal(2, groups.IncludedGroupCount); Assert.Equal(2, groups.IncludedMessageCount); } [Fact] public void CreateSummaryMessageCreatesSummaryGroup() { // Arrange ChatMessage summaryMessage = new(ChatRole.Assistant, "[Summary of earlier conversation]: key facts..."); (summaryMessage.AdditionalProperties ??= [])[MessageGroup.SummaryPropertyKey] = true; List messages = [summaryMessage]; // Act MessageIndex groups = MessageIndex.Create(messages); // Assert Assert.Single(groups.Groups); Assert.Equal(MessageGroupKind.Summary, groups.Groups[0].Kind); Assert.Same(summaryMessage, groups.Groups[0].Messages[0]); } [Fact] public void CreateSummaryAmongOtherMessagesGroupsCorrectly() { // Arrange ChatMessage systemMsg = new(ChatRole.System, "You are helpful."); ChatMessage summaryMsg = new(ChatRole.Assistant, "[Summary]: previous context"); (summaryMsg.AdditionalProperties ??= [])[MessageGroup.SummaryPropertyKey] = true; ChatMessage userMsg = new(ChatRole.User, "Continue..."); List messages = [systemMsg, summaryMsg, userMsg]; // Act MessageIndex groups = MessageIndex.Create(messages); // Assert Assert.Equal(3, groups.Groups.Count); Assert.Equal(MessageGroupKind.System, groups.Groups[0].Kind); Assert.Equal(MessageGroupKind.Summary, groups.Groups[1].Kind); Assert.Equal(MessageGroupKind.User, groups.Groups[2].Kind); } [Fact] public void MessageGroupStoresPassedCounts() { // Arrange & Act MessageGroup group = new(MessageGroupKind.User, [new ChatMessage(ChatRole.User, "Hello")], byteCount: 5, tokenCount: 2); // Assert Assert.Equal(1, group.MessageCount); Assert.Equal(5, group.ByteCount); Assert.Equal(2, group.TokenCount); } [Fact] public void MessageGroupMessagesAreImmutable() { // Arrange IReadOnlyList messages = [new ChatMessage(ChatRole.User, "Hello")]; MessageGroup group = new(MessageGroupKind.User, messages, byteCount: 5, tokenCount: 1); // Assert — Messages is IReadOnlyList, not IList Assert.IsAssignableFrom>(group.Messages); Assert.Same(messages, group.Messages); } [Fact] public void CreateComputesByteCountUtf8() { // Arrange — "Hello" is 5 UTF-8 bytes MessageIndex groups = MessageIndex.Create([new ChatMessage(ChatRole.User, "Hello")]); // Assert Assert.Equal(5, groups.Groups[0].ByteCount); } [Fact] public void CreateComputesByteCountMultiByteChars() { // Arrange — "café" has a multi-byte 'é' (2 bytes in UTF-8) → 5 bytes total MessageIndex groups = MessageIndex.Create([new ChatMessage(ChatRole.User, "café")]); // Assert Assert.Equal(5, groups.Groups[0].ByteCount); } [Fact] public void CreateComputesByteCountMultipleMessagesInGroup() { // Arrange — ToolCall group: assistant (tool call, null text) + tool result "OK" (2 bytes) ChatMessage assistantMsg = new(ChatRole.Assistant, [new FunctionCallContent("call1", "fn")]); ChatMessage toolResult = new(ChatRole.Tool, "OK"); MessageIndex groups = MessageIndex.Create([assistantMsg, toolResult]); // Assert — single ToolCall group with 2 messages Assert.Single(groups.Groups); Assert.Equal(2, groups.Groups[0].MessageCount); Assert.Equal(2, groups.Groups[0].ByteCount); // "OK" = 2 bytes, assistant text is null } [Fact] public void CreateDefaultTokenCountIsHeuristic() { // Arrange — "Hello world test data!" = 22 UTF-8 bytes → 22 / 4 = 5 estimated tokens MessageIndex groups = MessageIndex.Create([new ChatMessage(ChatRole.User, "Hello world test data!")]); // Assert Assert.Equal(22, groups.Groups[0].ByteCount); Assert.Equal(22 / 4, groups.Groups[0].TokenCount); } [Fact] public void CreateNullTextHasZeroCounts() { // Arrange — message with no text (e.g., pure function call) ChatMessage msg = new(ChatRole.Assistant, [new FunctionCallContent("call1", "get_weather")]); ChatMessage tool = new(ChatRole.Tool, string.Empty); MessageIndex groups = MessageIndex.Create([msg, tool]); // Assert Assert.Equal(2, groups.Groups[0].MessageCount); Assert.Equal(0, groups.Groups[0].ByteCount); Assert.Equal(0, groups.Groups[0].TokenCount); } [Fact] public void TotalAggregatesSumAllGroups() { // Arrange MessageIndex groups = MessageIndex.Create( [ new ChatMessage(ChatRole.User, "AAAA"), // 4 bytes new ChatMessage(ChatRole.Assistant, "BBBB"), // 4 bytes ]); groups.Groups[0].IsExcluded = true; // Act & Assert — totals include excluded groups Assert.Equal(2, groups.TotalGroupCount); Assert.Equal(2, groups.TotalMessageCount); Assert.Equal(8, groups.TotalByteCount); Assert.Equal(2, groups.TotalTokenCount); // Each group: 4 bytes / 4 = 1 token, 2 groups = 2 } [Fact] public void IncludedAggregatesExcludeMarkedGroups() { // Arrange MessageIndex groups = MessageIndex.Create( [ new ChatMessage(ChatRole.User, "AAAA"), // 4 bytes new ChatMessage(ChatRole.Assistant, "BBBB"), // 4 bytes new ChatMessage(ChatRole.User, "CCCC"), // 4 bytes ]); groups.Groups[0].IsExcluded = true; // Act & Assert Assert.Equal(3, groups.TotalGroupCount); Assert.Equal(2, groups.IncludedGroupCount); Assert.Equal(3, groups.TotalMessageCount); Assert.Equal(2, groups.IncludedMessageCount); Assert.Equal(12, groups.TotalByteCount); Assert.Equal(8, groups.IncludedByteCount); Assert.Equal(3, groups.TotalTokenCount); // 12 / 4 = 3 (across 3 groups of 4 bytes each = 1+1+1) Assert.Equal(2, groups.IncludedTokenCount); // 8 / 4 = 2 (2 included groups of 4 bytes = 1+1) } [Fact] public void ToolCallGroupAggregatesAcrossMessages() { // Arrange — tool call group with assistant "Ask" (3 bytes) + tool result "OK" (2 bytes) ChatMessage assistantMsg = new(ChatRole.Assistant, [new FunctionCallContent("call1", "fn")]); ChatMessage toolResult = new(ChatRole.Tool, "OK"); MessageIndex groups = MessageIndex.Create([assistantMsg, toolResult]); // Assert — single group with 2 messages Assert.Single(groups.Groups); Assert.Equal(2, groups.Groups[0].MessageCount); Assert.Equal(2, groups.Groups[0].ByteCount); // assistant text is null (function call), tool result is "OK" = 2 bytes Assert.Equal(1, groups.TotalGroupCount); Assert.Equal(2, groups.TotalMessageCount); } [Fact] public void CreateAssignsTurnIndicesSingleTurn() { // Arrange — System (no turn), User + Assistant = turn 1 MessageIndex groups = MessageIndex.Create( [ new ChatMessage(ChatRole.System, "You are helpful."), new ChatMessage(ChatRole.User, "Hello"), new ChatMessage(ChatRole.Assistant, "Hi!"), ]); // Assert Assert.Null(groups.Groups[0].TurnIndex); // System Assert.Equal(1, groups.Groups[1].TurnIndex); // User Assert.Equal(1, groups.Groups[2].TurnIndex); // Assistant Assert.Equal(1, groups.TotalTurnCount); Assert.Equal(1, groups.IncludedTurnCount); } [Fact] public void CreateAssignsTurnIndicesMultiTurn() { // Arrange — 3 user turns MessageIndex groups = MessageIndex.Create( [ new ChatMessage(ChatRole.System, "System prompt."), new ChatMessage(ChatRole.User, "Q1"), new ChatMessage(ChatRole.Assistant, "A1"), new ChatMessage(ChatRole.User, "Q2"), new ChatMessage(ChatRole.Assistant, "A2"), new ChatMessage(ChatRole.User, "Q3"), ]); // Assert — 6 groups: System(null), User(1), Assistant(1), User(2), Assistant(2), User(3) Assert.Null(groups.Groups[0].TurnIndex); Assert.Equal(1, groups.Groups[1].TurnIndex); Assert.Equal(1, groups.Groups[2].TurnIndex); Assert.Equal(2, groups.Groups[3].TurnIndex); Assert.Equal(2, groups.Groups[4].TurnIndex); Assert.Equal(3, groups.Groups[5].TurnIndex); Assert.Equal(3, groups.TotalTurnCount); } [Fact] public void CreateTurnSpansToolCallGroups() { // Arrange — turn 1 includes User, ToolCall, AssistantText ChatMessage assistantToolCall = new(ChatRole.Assistant, [new FunctionCallContent("call1", "get_weather")]); ChatMessage toolResult = new(ChatRole.Tool, "Sunny"); MessageIndex groups = MessageIndex.Create( [ new ChatMessage(ChatRole.User, "What's the weather?"), assistantToolCall, toolResult, new ChatMessage(ChatRole.Assistant, "The weather is sunny!"), ]); // Assert — all 3 groups belong to turn 1 Assert.Equal(3, groups.Groups.Count); Assert.Equal(1, groups.Groups[0].TurnIndex); // User Assert.Equal(1, groups.Groups[1].TurnIndex); // ToolCall Assert.Equal(1, groups.Groups[2].TurnIndex); // AssistantText Assert.Equal(1, groups.TotalTurnCount); } [Fact] public void GetTurnGroupsReturnsGroupsForSpecificTurn() { // Arrange MessageIndex groups = MessageIndex.Create( [ new ChatMessage(ChatRole.System, "System."), new ChatMessage(ChatRole.User, "Q1"), new ChatMessage(ChatRole.Assistant, "A1"), new ChatMessage(ChatRole.User, "Q2"), new ChatMessage(ChatRole.Assistant, "A2"), ]); // Act List turn1 = [.. groups.GetTurnGroups(1)]; List turn2 = [.. groups.GetTurnGroups(2)]; // Assert Assert.Equal(2, turn1.Count); Assert.Equal(MessageGroupKind.User, turn1[0].Kind); Assert.Equal(MessageGroupKind.AssistantText, turn1[1].Kind); Assert.Equal(2, turn2.Count); Assert.Equal(MessageGroupKind.User, turn2[0].Kind); Assert.Equal(MessageGroupKind.AssistantText, turn2[1].Kind); } [Fact] public void IncludedTurnCountReflectsExclusions() { // Arrange — 2 turns, exclude all groups in turn 1 MessageIndex groups = MessageIndex.Create( [ new ChatMessage(ChatRole.User, "Q1"), new ChatMessage(ChatRole.Assistant, "A1"), new ChatMessage(ChatRole.User, "Q2"), new ChatMessage(ChatRole.Assistant, "A2"), ]); groups.Groups[0].IsExcluded = true; // User Q1 (turn 1) groups.Groups[1].IsExcluded = true; // Assistant A1 (turn 1) // Assert Assert.Equal(2, groups.TotalTurnCount); Assert.Equal(1, groups.IncludedTurnCount); // Only turn 2 has included groups } [Fact] public void TotalTurnCountZeroWhenNoUserMessages() { // Arrange — only system messages MessageIndex groups = MessageIndex.Create( [ new ChatMessage(ChatRole.System, "System."), ]); // Assert Assert.Equal(0, groups.TotalTurnCount); Assert.Equal(0, groups.IncludedTurnCount); } [Fact] public void IncludedTurnCountPartialExclusionStillCountsTurn() { // Arrange — turn 1 has 2 groups, only one excluded MessageIndex groups = MessageIndex.Create( [ new ChatMessage(ChatRole.User, "Q1"), new ChatMessage(ChatRole.Assistant, "A1"), ]); groups.Groups[1].IsExcluded = true; // Exclude assistant but user is still included // Assert — turn 1 still has one included group Assert.Equal(1, groups.TotalTurnCount); Assert.Equal(1, groups.IncludedTurnCount); } [Fact] public void UpdateAppendsNewMessagesIncrementally() { // Arrange — create with 2 messages List messages = [ new ChatMessage(ChatRole.User, "Q1"), new ChatMessage(ChatRole.Assistant, "A1"), ]; MessageIndex index = MessageIndex.Create(messages); Assert.Equal(2, index.Groups.Count); Assert.Equal(2, index.ProcessedMessageCount); // Act — add 2 more messages and update messages.Add(new ChatMessage(ChatRole.User, "Q2")); messages.Add(new ChatMessage(ChatRole.Assistant, "A2")); index.Update(messages); // Assert — should have 4 groups total, processed count updated Assert.Equal(4, index.Groups.Count); Assert.Equal(4, index.ProcessedMessageCount); Assert.Equal(MessageGroupKind.User, index.Groups[2].Kind); Assert.Equal(MessageGroupKind.AssistantText, index.Groups[3].Kind); } [Fact] public void UpdateNoOpWhenNoNewMessages() { // Arrange List messages = [ new ChatMessage(ChatRole.User, "Q1"), ]; MessageIndex index = MessageIndex.Create(messages); int originalCount = index.Groups.Count; // Act — update with same count index.Update(messages); // Assert — nothing changed Assert.Equal(originalCount, index.Groups.Count); } [Fact] public void UpdateRebuildsWhenMessagesShrink() { // Arrange — create with 3 messages List messages = [ new ChatMessage(ChatRole.User, "Q1"), new ChatMessage(ChatRole.Assistant, "A1"), new ChatMessage(ChatRole.User, "Q2"), ]; MessageIndex index = MessageIndex.Create(messages); Assert.Equal(3, index.Groups.Count); // Exclude a group to verify rebuild clears state index.Groups[0].IsExcluded = true; // Act — update with fewer messages (simulates storage compaction) List shortened = [ new ChatMessage(ChatRole.User, "Q2"), ]; index.Update(shortened); // Assert — rebuilt from scratch Assert.Single(index.Groups); Assert.False(index.Groups[0].IsExcluded); Assert.Equal(1, index.ProcessedMessageCount); } [Fact] public void UpdatePreservesExistingGroupExclusionState() { // Arrange List messages = [ new ChatMessage(ChatRole.User, "Q1"), new ChatMessage(ChatRole.Assistant, "A1"), ]; MessageIndex index = MessageIndex.Create(messages); index.Groups[0].IsExcluded = true; index.Groups[0].ExcludeReason = "Test exclusion"; // Act — append new messages messages.Add(new ChatMessage(ChatRole.User, "Q2")); index.Update(messages); // Assert — original exclusion state preserved Assert.True(index.Groups[0].IsExcluded); Assert.Equal("Test exclusion", index.Groups[0].ExcludeReason); Assert.Equal(3, index.Groups.Count); } [Fact] public void InsertGroupInsertsAtSpecifiedIndex() { // Arrange MessageIndex index = MessageIndex.Create( [ new ChatMessage(ChatRole.User, "Q1"), new ChatMessage(ChatRole.User, "Q2"), ]); // Act — insert between Q1 and Q2 ChatMessage summaryMsg = new(ChatRole.Assistant, "[Summary]"); MessageGroup inserted = index.InsertGroup(1, MessageGroupKind.Summary, [summaryMsg], turnIndex: 1); // Assert Assert.Equal(3, index.Groups.Count); Assert.Same(inserted, index.Groups[1]); Assert.Equal(MessageGroupKind.Summary, index.Groups[1].Kind); Assert.Equal("[Summary]", index.Groups[1].Messages[0].Text); Assert.Equal(1, inserted.TurnIndex); } [Fact] public void AddGroupAppendsToEnd() { // Arrange MessageIndex index = MessageIndex.Create( [ new ChatMessage(ChatRole.User, "Q1"), ]); // Act ChatMessage msg = new(ChatRole.Assistant, "Appended"); MessageGroup added = index.AddGroup(MessageGroupKind.AssistantText, [msg], turnIndex: 1); // Assert Assert.Equal(2, index.Groups.Count); Assert.Same(added, index.Groups[1]); Assert.Equal("Appended", index.Groups[1].Messages[0].Text); } [Fact] public void InsertGroupComputesByteAndTokenCounts() { // Arrange MessageIndex index = MessageIndex.Create( [ new ChatMessage(ChatRole.User, "Q1"), ]); // Act — insert a group with known text ChatMessage msg = new(ChatRole.Assistant, "Hello"); // 5 bytes, ~1 token (5/4) MessageGroup inserted = index.InsertGroup(0, MessageGroupKind.AssistantText, [msg]); // Assert Assert.Equal(5, inserted.ByteCount); Assert.Equal(1, inserted.TokenCount); // 5 / 4 = 1 (integer division) } }