diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs index 4e7574af698..bceca5ca726 100644 --- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs @@ -215,7 +215,26 @@ private static PartitionKey BuildPartitionKey(State state) } /// - protected override async ValueTask> ProvideChatHistoryAsync(InvokingContext context, CancellationToken cancellationToken = default) + protected override ValueTask> ProvideChatHistoryAsync(InvokingContext context, CancellationToken cancellationToken = default) => + new(this.GetMessagesAsync(context.Session, cancellationToken)); + + /// + /// Gets the messages stored for the specified session. + /// + /// The agent session to get state from. + /// The cancellation token. + /// + /// The messages in the timestamp-based order used by the agent invocation pipeline. Messages with equal stored + /// timestamps have no guaranteed relative order. When is set, messages with + /// the latest timestamps are selected; if the limit intersects a timestamp tie, which tied messages are included + /// is unspecified. + /// + /// + /// This method returns messages as stored and does not apply the output filter or chat-history source attribution + /// used by the agent invocation pipeline. Use when that processing + /// is required. controls the query page size. + /// + public async Task> GetMessagesAsync(AgentSession? session, CancellationToken cancellationToken = default) { #pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks if (this._disposed) @@ -224,7 +243,7 @@ protected override async ValueTask> ProvideChatHistoryA } #pragma warning restore CA1513 - var state = this._sessionState.GetOrInitializeState(context.Session); + var state = this._sessionState.GetOrInitializeState(session); var partitionKey = BuildPartitionKey(state); // Fetch most recent messages in descending order when limit is set, then reverse to ascending @@ -233,7 +252,7 @@ protected override async ValueTask> ProvideChatHistoryA .WithParameter("@conversationId", state.ConversationId) .WithParameter("@type", "ChatMessage"); - var iterator = this._container.GetItemQueryIterator(query, requestOptions: new QueryRequestOptions + using var iterator = this._container.GetItemQueryIterator(query, requestOptions: new QueryRequestOptions { PartitionKey = partitionKey, MaxItemCount = this.MaxItemCount // Configurable query performance diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatHistoryProviderTests.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatHistoryProviderTests.cs index 23a3a81014d..b4fc013523f 100644 --- a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatHistoryProviderTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatHistoryProviderTests.cs @@ -3,6 +3,7 @@ using System; using System.Collections.Generic; using System.Linq; +using System.Threading; using System.Threading.Tasks; using Azure.Core; using Azure.Identity; @@ -807,6 +808,137 @@ public async Task HierarchicalAndSimplePartitioning_ShouldCoexistAsync() Assert.Equal("Hierarchical partitioning message", hierarchicalMessageList[0].Text); } + [Fact] + [Trait("Category", "CosmosDB")] + public async Task GetMessagesAsync_WithMessages_ShouldReturnAllMessagesAcrossPagesAsync() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + var session = CreateMockSession(); + const string ConversationId = "get-messages-test"; + + using var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, + _ => new CosmosChatHistoryProvider.State(ConversationId)) + { + MaxItemCount = 2, + }; + + List messages = + [ + new(ChatRole.User, "Message 1"), + new(ChatRole.Assistant, "Message 2"), + new(ChatRole.User, "Message 3"), + new(ChatRole.Assistant, "Message 4"), + new(ChatRole.User, "Message 5"), + ]; + + var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, session, messages, []); + await provider.InvokedAsync(context); + + // Wait a moment for eventual consistency + await Task.Delay(100); + + // Act + var retrievedMessages = (await provider.GetMessagesAsync(session)).ToList(); + + // Assert + Assert.Equal(messages.Count, retrievedMessages.Count); + // Batch writes share a Unix-seconds timestamp, so this test intentionally verifies page completeness + // without asserting relative order among tied messages. + Assert.Equal( + messages.Select(message => message.Text).Order(), + retrievedMessages.Select(message => message.Text).Order()); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public async Task GetMessagesAsync_WithNoMessages_ShouldReturnEmptyAsync() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + var session = CreateMockSession(); + + using var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, + _ => new CosmosChatHistoryProvider.State("get-messages-empty-test")); + + // Act + var messages = await provider.GetMessagesAsync(session); + + // Assert + Assert.Empty(messages); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public async Task GetMessagesAsync_DoesNotApplyInvocationFilterOrSourceAttributionAsync() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + var session = CreateMockSession(); + + using var provider = new CosmosChatHistoryProvider( + this._connectionString, + s_testDatabaseId, + TestContainerId, + _ => new CosmosChatHistoryProvider.State("get-messages-filter-test"), + provideOutputMessageFilter: messages => messages.Where(message => message.Text != "Hidden")); + + List messages = + [ + new(ChatRole.User, "Visible"), + new(ChatRole.Assistant, "Hidden"), + ]; + + var invokedContext = new ChatHistoryProvider.InvokedContext(s_mockAgent, session, messages, []); + await provider.InvokedAsync(invokedContext); + + // Act + var directMessages = (await provider.GetMessagesAsync(session)).ToList(); + var invokingContext = new ChatHistoryProvider.InvokingContext(s_mockAgent, session, []); + var invocationMessages = (await provider.InvokingAsync(invokingContext)).ToList(); + + // Assert + Assert.Equal(2, directMessages.Count); + Assert.All(directMessages, message => Assert.Equal(AgentRequestMessageSourceType.External, message.GetAgentRequestMessageSourceType())); + Assert.Single(invocationMessages); + Assert.Equal("Visible", invocationMessages[0].Text); + Assert.Equal(AgentRequestMessageSourceType.ChatHistory, invocationMessages[0].GetAgentRequestMessageSourceType()); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public async Task GetMessagesAsync_AfterDispose_ShouldThrowObjectDisposedExceptionAsync() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + var session = CreateMockSession(); + var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, + _ => new CosmosChatHistoryProvider.State("get-messages-disposed-test")); + provider.Dispose(); + + // Act & Assert + await Assert.ThrowsAsync(() => provider.GetMessagesAsync(session)); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public async Task GetMessagesAsync_WithCanceledToken_ShouldThrowOperationCanceledExceptionAsync() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + var session = CreateMockSession(); + + using var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, + _ => new CosmosChatHistoryProvider.State("get-messages-cancellation-test")); + + using var cancellationTokenSource = new CancellationTokenSource(); + cancellationTokenSource.Cancel(); + + // Act & Assert + await Assert.ThrowsAnyAsync( + () => provider.GetMessagesAsync(session, cancellationTokenSource.Token)); + } + [Fact] [Trait("Category", "CosmosDB")] public async Task MaxMessagesToRetrieve_ShouldLimitAndReturnMostRecentAsync()