diff --git a/dotnet/src/Microsoft.Agents.AI/Compaction/CompactionMessageGroup.cs b/dotnet/src/Microsoft.Agents.AI/Compaction/CompactionMessageGroup.cs index 049fa3013f..ddf4c20b24 100644 --- a/dotnet/src/Microsoft.Agents.AI/Compaction/CompactionMessageGroup.cs +++ b/dotnet/src/Microsoft.Agents.AI/Compaction/CompactionMessageGroup.cs @@ -1,5 +1,6 @@ // Copyright (c) Microsoft. All rights reserved. +using System; using System.Collections.Generic; using System.Diagnostics.CodeAnalysis; using System.Text.Json.Serialization; @@ -100,6 +101,14 @@ public sealed class CompactionMessageGroup /// public int? TurnIndex { get; } + private bool _isExcluded; + + /// + /// An optional callback invoked when changes value. + /// Used internally by to invalidate cached aggregates. + /// + internal Action? ExclusionChanged; + /// /// Gets or sets a value indicating whether this group is excluded from the projected message list. /// @@ -107,7 +116,18 @@ public sealed class CompactionMessageGroup /// Excluded groups are preserved in the collection for diagnostics or storage purposes /// but are not included when calling . /// - public bool IsExcluded { get; set; } + public bool IsExcluded + { + get => _isExcluded; + set + { + if (_isExcluded != value) + { + _isExcluded = value; + ExclusionChanged?.Invoke(); + } + } + } /// /// Gets or sets an optional reason explaining why this group was excluded. diff --git a/dotnet/src/Microsoft.Agents.AI/Compaction/CompactionMessageIndex.cs b/dotnet/src/Microsoft.Agents.AI/Compaction/CompactionMessageIndex.cs index 003a70f2b3..ca89277da5 100644 --- a/dotnet/src/Microsoft.Agents.AI/Compaction/CompactionMessageIndex.cs +++ b/dotnet/src/Microsoft.Agents.AI/Compaction/CompactionMessageIndex.cs @@ -27,6 +27,20 @@ public sealed class CompactionMessageIndex private int _currentTurn; private ChatMessage? _lastProcessedMessage; + // Cached values for derived properties — invalidated whenever groups are added/removed + // or a group's IsExcluded state changes. + private int? _cachedTotalMessageCount; + private int? _cachedTotalByteCount; + private int? _cachedTotalTokenCount; + private int? _cachedIncludedGroupCount; + private int? _cachedIncludedMessageCount; + private int? _cachedIncludedByteCount; + private int? _cachedIncludedTokenCount; + private int? _cachedTotalTurnCount; + private int? _cachedIncludedTurnCount; + private int? _cachedIncludedNonSystemGroupCount; + private int? _cachedRawMessageCount; + /// /// Gets the list of message groups in this collection. /// @@ -47,6 +61,12 @@ public sealed class CompactionMessageIndex this.Groups = Throw.IfNull(groups, nameof(groups)); this.Tokenizer = tokenizer; + // Register all pre-existing groups so that IsExcluded changes invalidate the cache. + for (int i = 0; i < groups.Count; i++) + { + this.RegisterGroup(groups[i]); + } + // Restore turn counter and last processed message from the groups for (int index = groups.Count - 1; index >= 0; --index) { @@ -123,6 +143,7 @@ public sealed class CompactionMessageIndex this.Groups.Clear(); this._currentTurn = 0; this._lastProcessedMessage = null; + this.InvalidateCache(); return; } @@ -184,13 +205,13 @@ public sealed class CompactionMessageIndex if (message.Role == ChatRole.System) { // System messages are not part of any turn - this.Groups.Add(CreateGroup(CompactionGroupKind.System, [message], this.Tokenizer, turnIndex: null)); + this.AddAndRegisterGroup(CreateGroup(CompactionGroupKind.System, [message], this.Tokenizer, turnIndex: null)); index++; } else if (message.Role == ChatRole.User) { this._currentTurn++; - this.Groups.Add(CreateGroup(CompactionGroupKind.User, [message], this.Tokenizer, this._currentTurn)); + this.AddAndRegisterGroup(CreateGroup(CompactionGroupKind.User, [message], this.Tokenizer, this._currentTurn)); index++; } else if (message.Role == ChatRole.Assistant && HasToolCalls(message)) @@ -207,11 +228,11 @@ public sealed class CompactionMessageIndex index++; } - this.Groups.Add(CreateGroup(CompactionGroupKind.ToolCall, groupMessages, this.Tokenizer, this._currentTurn)); + this.AddAndRegisterGroup(CreateGroup(CompactionGroupKind.ToolCall, groupMessages, this.Tokenizer, this._currentTurn)); } else if (message.Role == ChatRole.Assistant && IsSummaryMessage(message)) { - this.Groups.Add(CreateGroup(CompactionGroupKind.Summary, [message], this.Tokenizer, this._currentTurn)); + this.AddAndRegisterGroup(CreateGroup(CompactionGroupKind.Summary, [message], this.Tokenizer, this._currentTurn)); index++; } else if (message.Role == ChatRole.Assistant && HasOnlyReasoning(message)) @@ -247,17 +268,17 @@ public sealed class CompactionMessageIndex index++; } - this.Groups.Add(CreateGroup(CompactionGroupKind.ToolCall, groupMessages, this.Tokenizer, this._currentTurn)); + this.AddAndRegisterGroup(CreateGroup(CompactionGroupKind.ToolCall, groupMessages, this.Tokenizer, this._currentTurn)); } else { - this.Groups.Add(CreateGroup(CompactionGroupKind.AssistantText, [message], this.Tokenizer, this._currentTurn)); + this.AddAndRegisterGroup(CreateGroup(CompactionGroupKind.AssistantText, [message], this.Tokenizer, this._currentTurn)); index++; } } else { - this.Groups.Add(CreateGroup(CompactionGroupKind.AssistantText, [message], this.Tokenizer, this._currentTurn)); + this.AddAndRegisterGroup(CreateGroup(CompactionGroupKind.AssistantText, [message], this.Tokenizer, this._currentTurn)); index++; } } @@ -266,6 +287,8 @@ public sealed class CompactionMessageIndex { this._lastProcessedMessage = messages[^1]; } + + this.InvalidateCache(); } /// @@ -281,6 +304,8 @@ public sealed class CompactionMessageIndex { CompactionMessageGroup group = CreateGroup(kind, messages, this.Tokenizer, turnIndex); this.Groups.Insert(index, group); + this.RegisterGroup(group); + this.InvalidateCache(); return group; } @@ -296,6 +321,8 @@ public sealed class CompactionMessageIndex { CompactionMessageGroup group = CreateGroup(kind, messages, this.Tokenizer, turnIndex); this.Groups.Add(group); + this.RegisterGroup(group); + this.InvalidateCache(); return group; } @@ -320,57 +347,57 @@ public sealed class CompactionMessageIndex /// /// Gets the total number of messages across all groups, including excluded ones. /// - public int TotalMessageCount => this.Groups.Sum(group => group.MessageCount); + public int TotalMessageCount => _cachedTotalMessageCount ??= this.Groups.Sum(group => group.MessageCount); /// /// Gets the total UTF-8 byte count across all groups, including excluded ones. /// - public int TotalByteCount => this.Groups.Sum(group => group.ByteCount); + public int TotalByteCount => _cachedTotalByteCount ??= this.Groups.Sum(group => group.ByteCount); /// /// Gets the total token count across all groups, including excluded ones. /// - public int TotalTokenCount => this.Groups.Sum(group => group.TokenCount); + public int TotalTokenCount => _cachedTotalTokenCount ??= this.Groups.Sum(group => group.TokenCount); /// /// Gets the total number of groups that are not excluded. /// - public int IncludedGroupCount => this.Groups.Count(group => !group.IsExcluded); + public int IncludedGroupCount => _cachedIncludedGroupCount ??= this.Groups.Count(group => !group.IsExcluded); /// /// Gets the total number of messages across all included (non-excluded) groups. /// - public int IncludedMessageCount => this.Groups.Where(group => !group.IsExcluded).Sum(group => group.MessageCount); + public int IncludedMessageCount => _cachedIncludedMessageCount ??= this.Groups.Where(group => !group.IsExcluded).Sum(group => group.MessageCount); /// /// Gets the total UTF-8 byte count across all included (non-excluded) groups. /// - public int IncludedByteCount => this.Groups.Where(group => !group.IsExcluded).Sum(group => group.ByteCount); + public int IncludedByteCount => _cachedIncludedByteCount ??= this.Groups.Where(group => !group.IsExcluded).Sum(group => group.ByteCount); /// /// Gets the total token count across all included (non-excluded) groups. /// - public int IncludedTokenCount => this.Groups.Where(group => !group.IsExcluded).Sum(group => group.TokenCount); + public int IncludedTokenCount => _cachedIncludedTokenCount ??= this.Groups.Where(group => !group.IsExcluded).Sum(group => group.TokenCount); /// /// Gets the total number of user turns across all groups (including those with excluded groups). /// - public int TotalTurnCount => this.Groups.Select(group => group.TurnIndex).Distinct().Count(turnIndex => turnIndex is not null && turnIndex > 0); + public int TotalTurnCount => _cachedTotalTurnCount ??= this.Groups.Select(group => group.TurnIndex).Distinct().Count(turnIndex => turnIndex is not null && turnIndex > 0); /// /// Gets the number of user turns that have at least one non-excluded group. /// - public int IncludedTurnCount => this.Groups.Where(group => !group.IsExcluded && group.TurnIndex is not null && group.TurnIndex > 0).Select(group => group.TurnIndex).Distinct().Count(); + public int IncludedTurnCount => _cachedIncludedTurnCount ??= this.Groups.Where(group => !group.IsExcluded && group.TurnIndex is not null && group.TurnIndex > 0).Select(group => group.TurnIndex).Distinct().Count(); /// /// Gets the total number of groups across all included (non-excluded) groups that are not . /// - public int IncludedNonSystemGroupCount => this.Groups.Count(group => !group.IsExcluded && group.Kind != CompactionGroupKind.System); + public int IncludedNonSystemGroupCount => _cachedIncludedNonSystemGroupCount ??= this.Groups.Count(group => !group.IsExcluded && group.Kind != CompactionGroupKind.System); /// /// Gets the total number of original messages (that are not summaries). /// - public int RawMessageCount => this.Groups.Where(group => group.Kind != CompactionGroupKind.Summary).Sum(group => group.MessageCount); + public int RawMessageCount => _cachedRawMessageCount ??= this.Groups.Where(group => group.Kind != CompactionGroupKind.Summary).Sum(group => group.MessageCount); /// /// Returns all groups that belong to the specified user turn. @@ -379,6 +406,37 @@ public sealed class CompactionMessageIndex /// The groups belonging to the turn, in order. public IEnumerable GetTurnGroups(int turnIndex) => this.Groups.Where(group => group.TurnIndex == turnIndex); + private void InvalidateCache() + { + _cachedTotalMessageCount = null; + _cachedTotalByteCount = null; + _cachedTotalTokenCount = null; + _cachedIncludedGroupCount = null; + _cachedIncludedMessageCount = null; + _cachedIncludedByteCount = null; + _cachedIncludedTokenCount = null; + _cachedTotalTurnCount = null; + _cachedIncludedTurnCount = null; + _cachedIncludedNonSystemGroupCount = null; + _cachedRawMessageCount = null; + } + + private void RegisterGroup(CompactionMessageGroup group) + { + // Each group is owned by exactly one index, so assignment rather than + // += is intentional — no need to chain callbacks. + group.ExclusionChanged = this.InvalidateCache; + } + + // Adds the group to the list and registers it for cache invalidation. + // Callers that add many groups in a loop (e.g. AppendFromMessages) call + // InvalidateCache() once at the end rather than per-group for efficiency. + private void AddAndRegisterGroup(CompactionMessageGroup group) + { + this.Groups.Add(group); + this.RegisterGroup(group); + } + /// /// Computes the UTF-8 byte count for a set of messages across all content types. /// diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Compaction/CompactionMessageIndexTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Compaction/CompactionMessageIndexTests.cs index ea0ecd0d44..f8b4ebdc5b 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Compaction/CompactionMessageIndexTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Compaction/CompactionMessageIndexTests.cs @@ -1474,4 +1474,197 @@ public class CompactionMessageIndexTests Assert.Equal(CompactionGroupKind.ToolCall, index.Groups[2].Kind); Assert.Equal(3, index.Groups[2].MessageCount); // reasoning + toolCall + toolResult } + + // ----------------------------------------------------------------------- + // Cache invalidation tests + // ----------------------------------------------------------------------- + + [Fact] + public void CachedMetricsAreInvalidatedWhenIsExcludedChanges() + { + // Arrange — two groups, read Included* properties (populates cache) + CompactionMessageIndex index = CompactionMessageIndex.Create( + [ + new ChatMessage(ChatRole.User, "AAAA"), // 4 bytes, 1 token + new ChatMessage(ChatRole.Assistant, "BBBB"), // 4 bytes, 1 token + ]); + + // Prime the cache by reading properties + Assert.Equal(2, index.IncludedGroupCount); + Assert.Equal(2, index.IncludedMessageCount); + Assert.Equal(8, index.IncludedByteCount); + Assert.Equal(2, index.IncludedTokenCount); + Assert.Equal(1, index.IncludedTurnCount); + Assert.Equal(2, index.IncludedNonSystemGroupCount); + + // Act — exclude both groups in turn 1; cache must be invalidated + index.Groups[0].IsExcluded = true; + index.Groups[1].IsExcluded = true; + + // Assert — Included* properties now reflect the exclusions + Assert.Equal(0, index.IncludedGroupCount); + Assert.Equal(0, index.IncludedMessageCount); + Assert.Equal(0, index.IncludedByteCount); + Assert.Equal(0, index.IncludedTokenCount); + Assert.Equal(0, index.IncludedTurnCount); // turn 1 is fully excluded + Assert.Equal(0, index.IncludedNonSystemGroupCount); + } + + [Fact] + public void CachedMetricsAreInvalidatedWhenGroupAddedViaUpdate() + { + // Arrange — start with two messages, read RawMessageCount (populates cache) + List messages = + [ + new ChatMessage(ChatRole.User, "Q1"), + new ChatMessage(ChatRole.Assistant, "A1"), + ]; + CompactionMessageIndex index = CompactionMessageIndex.Create(messages); + + // Prime the cache + Assert.Equal(2, index.RawMessageCount); + Assert.Equal(2, index.IncludedGroupCount); + + // Act — add two more messages + messages.Add(new ChatMessage(ChatRole.User, "Q2")); + messages.Add(new ChatMessage(ChatRole.Assistant, "A2")); + index.Update(messages); + + // Assert — cached values updated + Assert.Equal(4, index.RawMessageCount); + Assert.Equal(4, index.IncludedGroupCount); + } + + [Fact] + public void CachedMetricsAreInvalidatedWhenGroupAddedViaAddGroup() + { + // Arrange — prime the cache + CompactionMessageIndex index = CompactionMessageIndex.Create( + [ + new ChatMessage(ChatRole.User, "Hello"), + ]); + Assert.Equal(1, index.IncludedGroupCount); + Assert.Equal(1, index.TotalMessageCount); + + // Act — add a group manually + index.AddGroup(CompactionGroupKind.AssistantText, [new ChatMessage(ChatRole.Assistant, "Hi")], turnIndex: 1); + + // Assert — cache invalidated and recomputed + Assert.Equal(2, index.IncludedGroupCount); + Assert.Equal(2, index.TotalMessageCount); + } + + [Fact] + public void CachedMetricsAreInvalidatedWhenGroupInsertedViaInsertGroup() + { + // Arrange — prime the cache + CompactionMessageIndex index = CompactionMessageIndex.Create( + [ + new ChatMessage(ChatRole.User, "Hello"), + ]); + Assert.Equal(1, index.IncludedGroupCount); + Assert.Equal(1, index.TotalMessageCount); + + // Act — insert a group + ChatMessage summaryMsg = new(ChatRole.Assistant, "[Summary]"); + (summaryMsg.AdditionalProperties ??= [])[CompactionMessageGroup.SummaryPropertyKey] = true; + index.InsertGroup(0, CompactionGroupKind.Summary, [summaryMsg]); + + // Assert — cache invalidated and recomputed + Assert.Equal(2, index.IncludedGroupCount); + Assert.Equal(2, index.TotalMessageCount); + Assert.Equal(1, index.RawMessageCount); // Summary excluded from RawMessageCount + } + + [Fact] + public void CachedMetricsAreInvalidatedWhenIndexRebuiltByUpdate() + { + // Arrange — populate and prime cache + List messages = + [ + new ChatMessage(ChatRole.User, "Q1"), + new ChatMessage(ChatRole.Assistant, "A1"), + ]; + CompactionMessageIndex index = CompactionMessageIndex.Create(messages); + index.Groups[0].IsExcluded = true; + + // Prime the cache after exclusion + Assert.Equal(1, index.IncludedGroupCount); + + // Act — update with a completely different (shorter) list — forces full rebuild + List newMessages = + [ + new ChatMessage(ChatRole.User, "NewQ"), + ]; + index.Update(newMessages); + + // Assert — full rebuild, previously excluded group is gone, cache is correct + Assert.Equal(1, index.IncludedGroupCount); + Assert.Equal(1, index.RawMessageCount); + } + + [Fact] + public void CachedMetricsAreInvalidatedWhenUpdateClearsAllMessages() + { + // Arrange — prime the cache + CompactionMessageIndex index = CompactionMessageIndex.Create( + [ + new ChatMessage(ChatRole.User, "Q1"), + new ChatMessage(ChatRole.Assistant, "A1"), + ]); + Assert.Equal(2, index.IncludedGroupCount); + Assert.Equal(2, index.RawMessageCount); + + // Act — update with empty list + index.Update([]); + + // Assert — all cached metrics reset + Assert.Equal(0, index.IncludedGroupCount); + Assert.Equal(0, index.TotalGroupCount); + Assert.Equal(0, index.RawMessageCount); + Assert.Equal(0, index.IncludedTokenCount); + } + + [Fact] + public void CachedTotalMetricsUnchangedWhenOnlyExcludedStateChanges() + { + // Arrange — prime the cache + CompactionMessageIndex index = CompactionMessageIndex.Create( + [ + new ChatMessage(ChatRole.User, "AAAA"), // 4 bytes + new ChatMessage(ChatRole.Assistant, "BBBB"), // 4 bytes + ]); + Assert.Equal(2, index.TotalGroupCount); + Assert.Equal(2, index.TotalMessageCount); + Assert.Equal(8, index.TotalByteCount); + + // Act — exclude a group + index.Groups[0].IsExcluded = true; + + // Assert — Total* properties still include both groups + Assert.Equal(2, index.TotalGroupCount); + Assert.Equal(2, index.TotalMessageCount); + Assert.Equal(8, index.TotalByteCount); + Assert.Equal(2, index.TotalTokenCount); + } + + [Fact] + public void ConstructorWithPreExistingGroupsRegistersThemForCacheInvalidation() + { + // Arrange — build groups externally, pass them into constructor + CompactionMessageGroup group1 = new(CompactionGroupKind.User, [new ChatMessage(ChatRole.User, "Q1")], 2, 1, turnIndex: 1); + CompactionMessageGroup group2 = new(CompactionGroupKind.AssistantText, [new ChatMessage(ChatRole.Assistant, "A1")], 2, 1, turnIndex: 1); + List groups = [group1, group2]; + + CompactionMessageIndex index = new(groups); + + // Prime the cache + Assert.Equal(2, index.IncludedGroupCount); + + // Act — exclude a pre-existing group (ExclusionChanged should fire and invalidate) + group1.IsExcluded = true; + + // Assert — cache reflects the change + Assert.Equal(1, index.IncludedGroupCount); + } }