From b778a521c9e0c5e685f74fa2dfb7412b327a43df Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Wed, 29 Oct 2025 21:19:34 +0000 Subject: [PATCH 01/40] draft commit --- dotnet/Directory.Packages.props | 4 + .../CosmosChatMessageStore.cs | 556 ++++++++++++++++++ .../Microsoft.Agents.AI.Abstractions.csproj | 3 + .../Extensions/AgentProviderExtensions.cs | 18 +- ...oft.Agents.AI.Workflows.Declarative.csproj | 1 + .../Checkpointing/CosmosCheckpointStore.cs | 284 +++++++++ .../CosmosDBWorkflowExtensions.cs | 157 +++++ .../Microsoft.Agents.AI.Workflows.csproj | 3 + .../CosmosDBChatExtensions.cs | 99 ++++ .../Microsoft.Agents.AI.csproj | 3 + .../CosmosChatMessageStoreTests.cs | 387 ++++++++++++ ...ft.Agents.AI.Abstractions.UnitTests.csproj | 2 + .../Microsoft.Agents.AI.UnitTests.csproj | 3 + .../CosmosCheckpointStoreTests.cs | 436 ++++++++++++++ ...osoft.Agents.AI.Workflows.UnitTests.csproj | 2 + 15 files changed, 1952 insertions(+), 6 deletions(-) create mode 100644 dotnet/src/Microsoft.Agents.AI.Abstractions/CosmosChatMessageStore.cs create mode 100644 dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/CosmosCheckpointStore.cs create mode 100644 dotnet/src/Microsoft.Agents.AI.Workflows/CosmosDBWorkflowExtensions.cs create mode 100644 dotnet/src/Microsoft.Agents.AI/CosmosDBChatExtensions.cs create mode 100644 dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/CosmosChatMessageStoreTests.cs create mode 100644 dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/CosmosCheckpointStoreTests.cs diff --git a/dotnet/Directory.Packages.props b/dotnet/Directory.Packages.props index afa5bc77f0..83ef3b5c1a 100644 --- a/dotnet/Directory.Packages.props +++ b/dotnet/Directory.Packages.props @@ -21,6 +21,10 @@ + + + + diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/CosmosChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/CosmosChatMessageStore.cs new file mode 100644 index 0000000000..0f00885c35 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/CosmosChatMessageStore.cs @@ -0,0 +1,556 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Diagnostics.CodeAnalysis; +using System.Linq; +using System.Text.Json; +using System.Text.Json.Serialization; +using System.Threading; +using System.Threading.Tasks; +using Azure.Identity; +using Microsoft.Azure.Cosmos; +using Microsoft.Extensions.AI; + +namespace Microsoft.Agents.AI; + +/// +/// Provides a Cosmos DB implementation of the abstract class. +/// +public sealed class CosmosChatMessageStore : ChatMessageStore, IDisposable +{ + private readonly CosmosClient _cosmosClient; + private readonly Container _container; + private readonly string _conversationId; + private readonly bool _ownsClient; + private bool _disposed; + + /// + /// Cached JSON serializer options for .NET 9.0 compatibility. + /// + private static readonly JsonSerializerOptions s_defaultJsonOptions = CreateDefaultJsonOptions(); + + [UnconditionalSuppressMessage("Trimming", "IL2026:Members annotated with 'RequiresUnreferencedCodeAttribute' require dynamic access", Justification = "JSON serialization is controlled")] + [UnconditionalSuppressMessage("AOT", "IL3050:Calling members annotated with 'RequiresDynamicCodeAttribute' may break functionality when AOT compiling", Justification = "JSON serialization is controlled")] + private static JsonSerializerOptions CreateDefaultJsonOptions() + { + var options = new JsonSerializerOptions(); +#if NET9_0_OR_GREATER + // Configure TypeInfoResolver for .NET 9.0 to enable JSON serialization + options.TypeInfoResolver = new System.Text.Json.Serialization.Metadata.DefaultJsonTypeInfoResolver(); +#endif + return options; + } + + /// + /// Gets or sets the maximum number of messages to return in a single query batch. + /// Default is 100 for optimal performance. + /// + public int MaxItemCount { get; set; } = 100; + + /// + /// Gets or sets the Time-To-Live (TTL) in seconds for messages. + /// Default is 86400 seconds (24 hours). Set to null to disable TTL. + /// + public int? MessageTtlSeconds { get; set; } = 86400; + + /// + /// Initializes a new instance of the class using a connection string. + /// + /// The Cosmos DB connection string. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// Thrown when any required parameter is null. + /// Thrown when any string parameter is null or whitespace. + public CosmosChatMessageStore(string connectionString, string databaseId, string containerId) + : this(connectionString, databaseId, containerId, Guid.NewGuid().ToString("N")) + { + } + + /// + /// Initializes a new instance of the class using a connection string. + /// + /// The Cosmos DB connection string. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// The unique identifier for this conversation thread. + /// Thrown when any required parameter is null. + /// Thrown when any string parameter is null or whitespace. + public CosmosChatMessageStore(string connectionString, string databaseId, string containerId, string conversationId) + { + if (string.IsNullOrWhiteSpace(connectionString)) + { + throw new ArgumentException("Connection string cannot be null or whitespace.", nameof(connectionString)); + } + + if (string.IsNullOrWhiteSpace(databaseId)) + { + throw new ArgumentException("Database ID cannot be null or whitespace.", nameof(databaseId)); + } + + if (string.IsNullOrWhiteSpace(containerId)) + { + throw new ArgumentException("Container ID cannot be null or whitespace.", nameof(containerId)); + } + + if (string.IsNullOrWhiteSpace(conversationId)) + { + throw new ArgumentException("Conversation ID cannot be null or whitespace.", nameof(conversationId)); + } + + this._cosmosClient = new CosmosClient(connectionString); + this._container = this._cosmosClient.GetContainer(databaseId, containerId); + this._conversationId = conversationId; + this._ownsClient = true; + } + + /// + /// Initializes a new instance of the class using DefaultAzureCredential. + /// + /// The Cosmos DB account endpoint URI. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// This parameter is used to distinguish this constructor from the connection string constructor. Always pass true. + /// Thrown when any required parameter is null. + /// Thrown when any string parameter is null or whitespace. + public CosmosChatMessageStore(string accountEndpoint, string databaseId, string containerId, bool useManagedIdentity) + : this(accountEndpoint, databaseId, containerId, Guid.NewGuid().ToString("N"), useManagedIdentity) + { + } + + /// + /// Initializes a new instance of the class using DefaultAzureCredential. + /// + /// The Cosmos DB account endpoint URI. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// The unique identifier for this conversation thread. + /// This parameter is used to distinguish this constructor from the connection string constructor. Always pass true. + /// Thrown when any required parameter is null. + /// Thrown when any string parameter is null or whitespace. + public CosmosChatMessageStore(string accountEndpoint, string databaseId, string containerId, string conversationId, bool useManagedIdentity) + { + if (string.IsNullOrWhiteSpace(accountEndpoint)) + { + throw new ArgumentException("Cannot be null or whitespace", nameof(accountEndpoint)); + } + + if (string.IsNullOrWhiteSpace(databaseId)) + { + throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); + } + + if (string.IsNullOrWhiteSpace(containerId)) + { + throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); + } + + if (string.IsNullOrWhiteSpace(conversationId)) + { + throw new ArgumentException("Cannot be null or whitespace", nameof(conversationId)); + } + + if (!useManagedIdentity) + { + throw new ArgumentException("This constructor requires useManagedIdentity to be true. Use the connection string constructor for key-based authentication.", nameof(useManagedIdentity)); + } + + this._cosmosClient = new CosmosClient(accountEndpoint, new DefaultAzureCredential()); + this._container = this._cosmosClient.GetContainer(databaseId, containerId); + this._conversationId = conversationId; + this._ownsClient = true; + } + + /// + /// Initializes a new instance of the class using an existing . + /// + /// The instance to use for Cosmos DB operations. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// Thrown when is null. + /// Thrown when any string parameter is null or whitespace. + public CosmosChatMessageStore(CosmosClient cosmosClient, string databaseId, string containerId) + : this(cosmosClient, databaseId, containerId, Guid.NewGuid().ToString("N")) + { + } + + /// + /// Initializes a new instance of the class using an existing . + /// + /// The instance to use for Cosmos DB operations. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// The unique identifier for this conversation thread. + /// Thrown when is null. + /// Thrown when any string parameter is null or whitespace. + public CosmosChatMessageStore(CosmosClient cosmosClient, string databaseId, string containerId, string conversationId) + { + this._cosmosClient = cosmosClient ?? throw new ArgumentNullException(nameof(cosmosClient)); + this._ownsClient = false; + + if (string.IsNullOrWhiteSpace(databaseId)) + { + throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); + } + + if (string.IsNullOrWhiteSpace(containerId)) + { + throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); + } + + if (string.IsNullOrWhiteSpace(conversationId)) + { + throw new ArgumentException("Cannot be null or whitespace", nameof(conversationId)); + } + + this._container = this._cosmosClient.GetContainer(databaseId, containerId); + this._conversationId = conversationId; + } + + /// + /// Initializes a new instance of the class from previously serialized state. + /// + /// A representing the serialized state of the message store. + /// The instance to use for Cosmos DB operations. + /// Optional settings for customizing the JSON deserialization process. + /// Thrown when is null. + /// Thrown when the serialized state cannot be deserialized. + [UnconditionalSuppressMessage("Trimming", "IL2026:Members annotated with 'RequiresUnreferencedCodeAttribute' require dynamic access", Justification = "StoreState type is controlled and used for serialization")] + [UnconditionalSuppressMessage("AOT", "IL3050:Calling members annotated with 'RequiresDynamicCodeAttribute' may break functionality when AOT compiling", Justification = "StoreState type is controlled and used for serialization")] + public CosmosChatMessageStore(JsonElement serializedStoreState, CosmosClient cosmosClient, JsonSerializerOptions? jsonSerializerOptions = null) + { + this._cosmosClient = cosmosClient ?? throw new ArgumentNullException(nameof(cosmosClient)); + this._ownsClient = false; + + if (serializedStoreState.ValueKind is JsonValueKind.Object) + { + var state = JsonSerializer.Deserialize(serializedStoreState, jsonSerializerOptions); + if (state?.ConversationId is { } conversationId && state.DatabaseId is { } databaseId && state.ContainerId is { } containerId) + { + this._conversationId = conversationId; + this._container = this._cosmosClient.GetContainer(databaseId, containerId); + return; + } + } + + throw new ArgumentException("Invalid serialized state", nameof(serializedStoreState)); + } + + /// + /// Gets the unique identifier for this conversation thread. + /// + public string ConversationId => this._conversationId; + + /// + /// Gets the identifier of the Cosmos DB database. + /// + public string DatabaseId => this._container.Database.Id; + + /// + /// Gets the identifier of the Cosmos DB container. + /// + public string ContainerId => this._container.Id; + + /// + [UnconditionalSuppressMessage("Trimming", "IL2026:Members annotated with 'RequiresUnreferencedCodeAttribute' require dynamic access", Justification = "ChatMessage deserialization is controlled")] + [UnconditionalSuppressMessage("AOT", "IL3050:Calling members annotated with 'RequiresDynamicCodeAttribute' may break functionality when AOT compiling", Justification = "ChatMessage deserialization is controlled")] + public override async Task> GetMessagesAsync(CancellationToken cancellationToken = default) + { +#pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks + if (this._disposed) + { + throw new ObjectDisposedException(GetType().FullName); + } +#pragma warning restore CA1513 + + // Use type discriminator for efficient queries + var query = new QueryDefinition("SELECT * FROM c WHERE c.conversationId = @conversationId AND c.Type = @type ORDER BY c.Timestamp ASC") + .WithParameter("@conversationId", this._conversationId) + .WithParameter("@type", "ChatMessage"); + + var iterator = this._container.GetItemQueryIterator(query, requestOptions: new QueryRequestOptions + { + PartitionKey = new PartitionKey(this._conversationId), + MaxItemCount = this.MaxItemCount // Configurable query performance + }); + + var messages = new List(); + + while (iterator.HasMoreResults) + { + var response = await iterator.ReadNextAsync(cancellationToken).ConfigureAwait(false); + + foreach (var document in response) + { + if (!string.IsNullOrEmpty(document.Message)) + { + var message = JsonSerializer.Deserialize(document.Message, s_defaultJsonOptions); + if (message != null) + { + messages.Add(message); + } + } + } + } + + return messages; + } + + /// + [UnconditionalSuppressMessage("Trimming", "IL2026:Members annotated with 'RequiresUnreferencedCodeAttribute' require dynamic access", Justification = "ChatMessage serialization is controlled")] + [UnconditionalSuppressMessage("AOT", "IL3050:Calling members annotated with 'RequiresDynamicCodeAttribute' may break functionality when AOT compiling", Justification = "ChatMessage serialization is controlled")] + public override async Task AddMessagesAsync(IEnumerable messages, CancellationToken cancellationToken = default) + { + if (messages is null) + { + throw new ArgumentNullException(nameof(messages)); + } + +#pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks + if (this._disposed) + { + throw new ObjectDisposedException(GetType().FullName); + } +#pragma warning restore CA1513 + + var messageList = messages.ToList(); + if (messageList.Count == 0) + { + return; + } + + // Use transactional batch for atomic operations + if (messageList.Count > 1) + { + await this.AddMessagesInBatchAsync(messageList, cancellationToken).ConfigureAwait(false); + } + else + { + await this.AddSingleMessageAsync(messageList[0], cancellationToken).ConfigureAwait(false); + } + } + + /// + /// Adds multiple messages using transactional batch operations for atomicity. + /// + private async Task AddMessagesInBatchAsync(IList messages, CancellationToken cancellationToken) + { + var partitionKey = new PartitionKey(this._conversationId); + var batch = this._container.CreateTransactionalBatch(partitionKey); + var currentTimestamp = DateTimeOffset.UtcNow.ToUnixTimeSeconds(); + + foreach (var message in messages) + { + var document = this.CreateMessageDocument(message, currentTimestamp); + batch.CreateItem(document); + } + + try + { + var response = await batch.ExecuteAsync(cancellationToken).ConfigureAwait(false); + if (!response.IsSuccessStatusCode) + { + throw new InvalidOperationException($"Batch operation failed with status: {response.StatusCode}"); + } + } + catch (CosmosException ex) when (ex.StatusCode == System.Net.HttpStatusCode.RequestEntityTooLarge) + { + // Fallback to individual operations if batch is too large + foreach (var message in messages) + { + await this.AddSingleMessageAsync(message, cancellationToken).ConfigureAwait(false); + } + } + } + + /// + /// Adds a single message to the store. + /// + private async Task AddSingleMessageAsync(ChatMessage message, CancellationToken cancellationToken) + { + var document = this.CreateMessageDocument(message, DateTimeOffset.UtcNow.ToUnixTimeSeconds()); + await this._container.CreateItemAsync(document, new PartitionKey(this._conversationId), cancellationToken: cancellationToken).ConfigureAwait(false); + } + + /// + /// Creates a message document with enhanced metadata. + /// + [UnconditionalSuppressMessage("Trimming", "IL2026:Members annotated with 'RequiresUnreferencedCodeAttribute' require dynamic access", Justification = "ChatMessage serialization is controlled")] + [UnconditionalSuppressMessage("AOT", "IL3050:Calling members annotated with 'RequiresDynamicCodeAttribute' may break functionality when AOT compiling", Justification = "ChatMessage serialization is controlled")] + private CosmosMessageDocument CreateMessageDocument(ChatMessage message, long timestamp) + { + return new CosmosMessageDocument + { + Id = Guid.NewGuid().ToString(), + ConversationId = this._conversationId, + Timestamp = timestamp, + MessageId = message.MessageId ?? Guid.NewGuid().ToString(), + Role = message.Role.Value ?? "unknown", + Message = JsonSerializer.Serialize(message, s_defaultJsonOptions), + Type = "ChatMessage", // Type discriminator + Ttl = this.MessageTtlSeconds // Configurable TTL + }; + } + + /// + [UnconditionalSuppressMessage("Trimming", "IL2026:Members annotated with 'RequiresUnreferencedCodeAttribute' require dynamic access", Justification = "StoreState serialization is controlled")] + [UnconditionalSuppressMessage("AOT", "IL3050:Calling members annotated with 'RequiresDynamicCodeAttribute' may break functionality when AOT compiling", Justification = "StoreState serialization is controlled")] + public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) + { +#pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks + if (this._disposed) + { + throw new ObjectDisposedException(GetType().FullName); + } +#pragma warning restore CA1513 + + var state = new StoreState + { + ConversationId = this._conversationId, + DatabaseId = this.DatabaseId, + ContainerId = this.ContainerId + }; + + var options = jsonSerializerOptions ?? s_defaultJsonOptions; + return JsonSerializer.SerializeToElement(state, options); + } + + /// + /// Gets the count of messages in this conversation. + /// This is an additional utility method beyond the base contract. + /// + /// The cancellation token. + /// The number of messages in the conversation. + public async Task GetMessageCountAsync(CancellationToken cancellationToken = default) + { +#pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks + if (this._disposed) + { + throw new ObjectDisposedException(GetType().FullName); + } +#pragma warning restore CA1513 + + // Efficient count query + var query = new QueryDefinition("SELECT VALUE COUNT(1) FROM c WHERE c.conversationId = @conversationId AND c.type = @type") + .WithParameter("@conversationId", this._conversationId) + .WithParameter("@type", "ChatMessage"); + + var iterator = this._container.GetItemQueryIterator(query, requestOptions: new QueryRequestOptions + { + PartitionKey = new PartitionKey(this._conversationId) + }); + + if (iterator.HasMoreResults) + { + var response = await iterator.ReadNextAsync(cancellationToken).ConfigureAwait(false); + return response.FirstOrDefault(); + } + + return 0; + } + + /// + /// Deletes all messages in this conversation. + /// This is an additional utility method beyond the base contract. + /// + /// The cancellation token. + /// The number of messages deleted. + public async Task ClearMessagesAsync(CancellationToken cancellationToken = default) + { +#pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks + if (this._disposed) + { + throw new ObjectDisposedException(GetType().FullName); + } +#pragma warning restore CA1513 + + // Batch delete for efficiency + var query = new QueryDefinition("SELECT c.id FROM c WHERE c.conversationId = @conversationId AND c.type = @type") + .WithParameter("@conversationId", this._conversationId) + .WithParameter("@type", "ChatMessage"); + + var iterator = this._container.GetItemQueryIterator(query, requestOptions: new QueryRequestOptions + { + PartitionKey = new PartitionKey(this._conversationId), + MaxItemCount = this.MaxItemCount + }); + + var deletedCount = 0; + var partitionKey = new PartitionKey(this._conversationId); + + while (iterator.HasMoreResults) + { + var response = await iterator.ReadNextAsync(cancellationToken).ConfigureAwait(false); + var batch = this._container.CreateTransactionalBatch(partitionKey); + var batchItemCount = 0; + + foreach (var item in response) + { + if (item.TryGetProperty("id", out var idElement) && idElement.ValueKind == JsonValueKind.String) + { + batch.DeleteItem(idElement.GetString()); + batchItemCount++; + deletedCount++; + } + } + + if (batchItemCount > 0) + { + await batch.ExecuteAsync(cancellationToken).ConfigureAwait(false); + } + } + + return deletedCount; + } + + /// + public void Dispose() + { + if (!this._disposed) + { + if (this._ownsClient) + { + this._cosmosClient?.Dispose(); + } + this._disposed = true; + } + } + + private sealed class StoreState + { + public string ConversationId { get; set; } = string.Empty; + public string DatabaseId { get; set; } = string.Empty; + public string ContainerId { get; set; } = string.Empty; + } + + /// + /// Represents a document stored in Cosmos DB for chat messages. + /// + [SuppressMessage("Performance", "CA1812:Avoid uninstantiated internal classes", Justification = "Instantiated by Cosmos DB operations")] + private sealed class CosmosMessageDocument + { + [Newtonsoft.Json.JsonProperty("id")] + public string Id { get; set; } = string.Empty; + + [Newtonsoft.Json.JsonProperty("conversationId")] + public string ConversationId { get; set; } = string.Empty; + + [Newtonsoft.Json.JsonProperty(nameof(Timestamp))] + public long Timestamp { get; set; } + + [Newtonsoft.Json.JsonProperty(nameof(MessageId))] + public string MessageId { get; set; } = string.Empty; + + [Newtonsoft.Json.JsonProperty(nameof(Role))] + public string Role { get; set; } = string.Empty; + + [Newtonsoft.Json.JsonProperty(nameof(Message))] + public string Message { get; set; } = string.Empty; + + [Newtonsoft.Json.JsonProperty(nameof(Type))] + public string Type { get; set; } = string.Empty; + + [Newtonsoft.Json.JsonProperty(nameof(Ttl))] + public int? Ttl { get; set; } + } +} \ No newline at end of file diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/Microsoft.Agents.AI.Abstractions.csproj b/dotnet/src/Microsoft.Agents.AI.Abstractions/Microsoft.Agents.AI.Abstractions.csproj index 4add7f427c..d58a173b3e 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/Microsoft.Agents.AI.Abstractions.csproj +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/Microsoft.Agents.AI.Abstractions.csproj @@ -28,6 +28,9 @@ + + + diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Extensions/AgentProviderExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Extensions/AgentProviderExtensions.cs index 5b6bbbc297..6c09043bb4 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Extensions/AgentProviderExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Extensions/AgentProviderExtensions.cs @@ -4,20 +4,24 @@ using System.Linq; using System.Threading; using System.Threading.Tasks; +#if NET9_0_OR_GREATER using Azure.AI.Agents.Persistent; +#endif using Microsoft.Extensions.AI; namespace Microsoft.Agents.AI.Workflows.Declarative.Extensions; internal static class AgentProviderExtensions { - private static readonly HashSet s_failureStatus = +#if NET9_0_OR_GREATER + private static readonly HashSet s_failureStatus = [ - Azure.AI.Agents.Persistent.RunStatus.Failed, - Azure.AI.Agents.Persistent.RunStatus.Cancelled, - Azure.AI.Agents.Persistent.RunStatus.Cancelling, - Azure.AI.Agents.Persistent.RunStatus.Expired, + global::Azure.AI.Agents.Persistent.RunStatus.Failed, + global::Azure.AI.Agents.Persistent.RunStatus.Cancelled, + global::Azure.AI.Agents.Persistent.RunStatus.Cancelling, + global::Azure.AI.Agents.Persistent.RunStatus.Expired, ]; +#endif public static async ValueTask InvokeAgentAsync( this WorkflowAgentProvider agentProvider, @@ -60,12 +64,14 @@ inputMessages is not null ? updates.Add(update); +#if NET9_0_OR_GREATER if (update.RawRepresentation is ChatResponseUpdate chatUpdate && - chatUpdate.RawRepresentation is RunUpdate runUpdate && + chatUpdate.RawRepresentation is global::Azure.AI.Agents.Persistent.RunUpdate runUpdate && s_failureStatus.Contains(runUpdate.Value.Status)) { throw new DeclarativeActionException($"Unexpected failure invoking agent, run {runUpdate.Value.Status}: {agent.Name ?? agent.Id} [{runUpdate.Value.Id}/{conversationId}]"); } +#endif if (autoSend) { diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Microsoft.Agents.AI.Workflows.Declarative.csproj b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Microsoft.Agents.AI.Workflows.Declarative.csproj index b3b2a86ab4..7b268f6c42 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Microsoft.Agents.AI.Workflows.Declarative.csproj +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Microsoft.Agents.AI.Workflows.Declarative.csproj @@ -22,6 +22,7 @@ + diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/CosmosCheckpointStore.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/CosmosCheckpointStore.cs new file mode 100644 index 0000000000..c9814e0456 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/CosmosCheckpointStore.cs @@ -0,0 +1,284 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Linq; +using System.Text.Json; +using System.Threading.Tasks; +using Azure.Identity; +using Microsoft.Azure.Cosmos; +using Newtonsoft.Json; +using Newtonsoft.Json.Linq; + +namespace Microsoft.Agents.AI.Workflows.Checkpointing; + +/// +/// Provides a Cosmos DB implementation of the abstract class. +/// +/// The type of objects to store as checkpoint values. +public class CosmosCheckpointStore : JsonCheckpointStore, IDisposable +{ + private readonly CosmosClient _cosmosClient; + private readonly Container _container; + private readonly bool _ownsClient; + private bool _disposed; + + /// + /// Initializes a new instance of the class using a connection string. + /// + /// The Cosmos DB connection string. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// Thrown when any required parameter is null. + /// Thrown when any string parameter is null or whitespace. + public CosmosCheckpointStore(string connectionString, string databaseId, string containerId) + { + if (string.IsNullOrWhiteSpace(connectionString)) + throw new ArgumentException("Cannot be null or whitespace", nameof(connectionString)); + if (string.IsNullOrWhiteSpace(databaseId)) + throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); + if (string.IsNullOrWhiteSpace(containerId)) + throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); + + var cosmosClientOptions = new CosmosClientOptions(); + + _cosmosClient = new CosmosClient(connectionString, cosmosClientOptions); + _container = _cosmosClient.GetContainer(databaseId, containerId); + _ownsClient = true; + } + + /// + /// Initializes a new instance of the class using DefaultAzureCredential. + /// + /// The Cosmos DB account endpoint URI. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// This parameter is used to distinguish this constructor from the connection string constructor. Always pass true. + /// Thrown when any required parameter is null. + /// Thrown when any string parameter is null or whitespace. + public CosmosCheckpointStore(string accountEndpoint, string databaseId, string containerId, bool useManagedIdentity) + { + if (string.IsNullOrWhiteSpace(accountEndpoint)) + throw new ArgumentException("Cannot be null or whitespace", nameof(accountEndpoint)); + if (string.IsNullOrWhiteSpace(databaseId)) + throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); + if (string.IsNullOrWhiteSpace(containerId)) + throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); + + if (!useManagedIdentity) + throw new ArgumentException("This constructor requires useManagedIdentity to be true. Use the connection string constructor for key-based authentication.", nameof(useManagedIdentity)); + + var cosmosClientOptions = new CosmosClientOptions + { + SerializerOptions = new CosmosSerializationOptions + { + PropertyNamingPolicy = CosmosPropertyNamingPolicy.CamelCase + } + }; + + _cosmosClient = new CosmosClient(accountEndpoint, new DefaultAzureCredential(), cosmosClientOptions); + _container = _cosmosClient.GetContainer(databaseId, containerId); + _ownsClient = true; + } + + /// + /// Initializes a new instance of the class using an existing . + /// + /// The instance to use for Cosmos DB operations. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// Thrown when is null. + /// Thrown when any string parameter is null or whitespace. + public CosmosCheckpointStore(CosmosClient cosmosClient, string databaseId, string containerId) + { + _cosmosClient = cosmosClient ?? throw new ArgumentNullException(nameof(cosmosClient)); + + if (string.IsNullOrWhiteSpace(databaseId)) + throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); + if (string.IsNullOrWhiteSpace(containerId)) + throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); + + _container = _cosmosClient.GetContainer(databaseId, containerId); + _ownsClient = false; + } + + /// + /// Gets the identifier of the Cosmos DB database. + /// + public string DatabaseId => _container.Database.Id; + + /// + /// Gets the identifier of the Cosmos DB container. + /// + public string ContainerId => _container.Id; + + /// + public override async ValueTask CreateCheckpointAsync(string runId, JsonElement value, CheckpointInfo? parent = null) + { + if (string.IsNullOrWhiteSpace(runId)) + throw new ArgumentException("Cannot be null or whitespace", nameof(runId)); +#pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks + if (_disposed) + throw new ObjectDisposedException(GetType().FullName); +#pragma warning restore CA1513 + + var checkpointId = Guid.NewGuid().ToString("N"); + var checkpointInfo = new CheckpointInfo(runId, checkpointId); + + var document = new CosmosCheckpointDocument + { + Id = $"{runId}_{checkpointId}", + RunId = runId, + CheckpointId = checkpointId, + Value = JToken.Parse(value.GetRawText()), + ParentCheckpointId = parent?.CheckpointId, + Timestamp = DateTimeOffset.UtcNow.ToUnixTimeSeconds() + }; + + await _container.CreateItemAsync(document, new PartitionKey(runId)).ConfigureAwait(false); + return checkpointInfo; + } + + /// + public override async ValueTask RetrieveCheckpointAsync(string runId, CheckpointInfo key) + { + if (string.IsNullOrWhiteSpace(runId)) + throw new ArgumentException("Cannot be null or whitespace", nameof(runId)); + if (key is null) + throw new ArgumentNullException(nameof(key)); +#pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks + if (_disposed) + throw new ObjectDisposedException(GetType().FullName); +#pragma warning restore CA1513 + + var id = $"{runId}_{key.CheckpointId}"; + + try + { + var response = await _container.ReadItemAsync(id, new PartitionKey(runId)).ConfigureAwait(false); + using var document = JsonDocument.Parse(response.Resource.Value.ToString()); + return document.RootElement.Clone(); + } + catch (CosmosException ex) when (ex.StatusCode == System.Net.HttpStatusCode.NotFound) + { + throw new InvalidOperationException($"Checkpoint with ID '{key.CheckpointId}' for run '{runId}' not found."); + } + } + + /// + public override async ValueTask> RetrieveIndexAsync(string runId, CheckpointInfo? withParent = null) + { + if (string.IsNullOrWhiteSpace(runId)) + throw new ArgumentException("Cannot be null or whitespace", nameof(runId)); +#pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks + if (_disposed) + throw new ObjectDisposedException(GetType().FullName); +#pragma warning restore CA1513 + + QueryDefinition query; + + if (withParent == null) + { + query = new QueryDefinition("SELECT c.runId, c.checkpointId FROM c WHERE c.runId = @runId ORDER BY c.timestamp ASC") + .WithParameter("@runId", runId); + } + else + { + query = new QueryDefinition("SELECT c.runId, c.checkpointId FROM c WHERE c.runId = @runId AND c.parentCheckpointId = @parentCheckpointId ORDER BY c.timestamp ASC") + .WithParameter("@runId", runId) + .WithParameter("@parentCheckpointId", withParent.CheckpointId); + } + + var iterator = _container.GetItemQueryIterator(query); + var checkpoints = new List(); + + while (iterator.HasMoreResults) + { + var response = await iterator.ReadNextAsync().ConfigureAwait(false); + checkpoints.AddRange(response.Select(r => new CheckpointInfo(r.RunId, r.CheckpointId))); + } + + return checkpoints; + } + + /// + public void Dispose() + { + this.Dispose(true); + GC.SuppressFinalize(this); + } + + /// + /// Releases the unmanaged resources used by the and optionally releases the managed resources. + /// + /// true to release both managed and unmanaged resources; false to release only unmanaged resources. + protected virtual void Dispose(bool disposing) + { + if (!this._disposed) + { + if (disposing && this._ownsClient) + { + this._cosmosClient?.Dispose(); + } + this._disposed = true; + } + } + + /// + /// Represents a checkpoint document stored in Cosmos DB. + /// + internal sealed class CosmosCheckpointDocument + { + [JsonProperty("id")] + public string Id { get; set; } = string.Empty; + + [JsonProperty("runId")] + public string RunId { get; set; } = string.Empty; + + [JsonProperty("checkpointId")] + public string CheckpointId { get; set; } = string.Empty; + + [JsonProperty("value")] + public JToken Value { get; set; } = JValue.CreateNull(); + + [JsonProperty("parentCheckpointId")] + public string? ParentCheckpointId { get; set; } + + [JsonProperty("timestamp")] + public long Timestamp { get; set; } + } + + /// + /// Represents the result of a checkpoint query. + /// + [System.Diagnostics.CodeAnalysis.SuppressMessage("Performance", "CA1812:Avoid uninstantiated internal classes", Justification = "Instantiated by Cosmos DB query deserialization")] + private sealed class CheckpointQueryResult + { + public string RunId { get; set; } = string.Empty; + public string CheckpointId { get; set; } = string.Empty; + } +} + +/// +/// Provides a non-generic Cosmos DB implementation of the abstract class. +/// +public sealed class CosmosCheckpointStore : CosmosCheckpointStore +{ + /// + public CosmosCheckpointStore(string connectionString, string databaseId, string containerId) + : base(connectionString, databaseId, containerId) + { + } + + /// + public CosmosCheckpointStore(string accountEndpoint, string databaseId, string containerId, bool useManagedIdentity) + : base(accountEndpoint, databaseId, containerId, useManagedIdentity) + { + } + + /// + public CosmosCheckpointStore(CosmosClient cosmosClient, string databaseId, string containerId) + : base(cosmosClient, databaseId, containerId) + { + } +} \ No newline at end of file diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/CosmosDBWorkflowExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/CosmosDBWorkflowExtensions.cs new file mode 100644 index 0000000000..b36651f3f3 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/CosmosDBWorkflowExtensions.cs @@ -0,0 +1,157 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Text.Json; +using Microsoft.Azure.Cosmos; +using Microsoft.Agents.AI.Workflows.Checkpointing; + +namespace Microsoft.Agents.AI.Workflows; + +/// +/// Provides extension methods for integrating Cosmos DB checkpoint storage with the Agent Framework. +/// +public static class CosmosDBWorkflowExtensions +{ + /// + /// Creates a Cosmos DB checkpoint store using connection string authentication. + /// + /// The Cosmos DB connection string. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// A new instance of . + /// Thrown when any string parameter is null or whitespace. + public static CosmosCheckpointStore CreateCheckpointStore( + string connectionString, + string databaseId, + string containerId) + { + if (string.IsNullOrWhiteSpace(connectionString)) + throw new ArgumentException("Cannot be null or whitespace", nameof(connectionString)); + if (string.IsNullOrWhiteSpace(databaseId)) + throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); + if (string.IsNullOrWhiteSpace(containerId)) + throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); + + return new CosmosCheckpointStore(connectionString, databaseId, containerId); + } + + /// + /// Creates a Cosmos DB checkpoint store using managed identity authentication. + /// + /// The Cosmos DB account endpoint URI. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// A new instance of . + /// Thrown when any string parameter is null or whitespace. + public static CosmosCheckpointStore CreateCheckpointStoreUsingManagedIdentity( + string accountEndpoint, + string databaseId, + string containerId) + { + if (string.IsNullOrWhiteSpace(accountEndpoint)) + throw new ArgumentException("Cannot be null or whitespace", nameof(accountEndpoint)); + if (string.IsNullOrWhiteSpace(databaseId)) + throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); + if (string.IsNullOrWhiteSpace(containerId)) + throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); + + return new CosmosCheckpointStore(accountEndpoint, databaseId, containerId, useManagedIdentity: true); + } + + /// + /// Creates a Cosmos DB checkpoint store using an existing . + /// + /// The instance to use for Cosmos DB operations. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// A new instance of . + /// Thrown when any required parameter is null. + /// Thrown when any string parameter is null or whitespace. + public static CosmosCheckpointStore CreateCheckpointStore( + CosmosClient cosmosClient, + string databaseId, + string containerId) + { + if (cosmosClient is null) + throw new ArgumentNullException(nameof(cosmosClient)); + if (string.IsNullOrWhiteSpace(databaseId)) + throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); + if (string.IsNullOrWhiteSpace(containerId)) + throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); + + return new CosmosCheckpointStore(cosmosClient, databaseId, containerId); + } + + /// + /// Creates a generic Cosmos DB checkpoint store using connection string authentication. + /// + /// The type of objects to store as checkpoint values. + /// The Cosmos DB connection string. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// A new instance of . + /// Thrown when any string parameter is null or whitespace. + public static CosmosCheckpointStore CreateCheckpointStore( + string connectionString, + string databaseId, + string containerId) + { + if (string.IsNullOrWhiteSpace(connectionString)) + throw new ArgumentException("Cannot be null or whitespace", nameof(connectionString)); + if (string.IsNullOrWhiteSpace(databaseId)) + throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); + if (string.IsNullOrWhiteSpace(containerId)) + throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); + + return new CosmosCheckpointStore(connectionString, databaseId, containerId); + } + + /// + /// Creates a generic Cosmos DB checkpoint store using managed identity authentication. + /// + /// The type of objects to store as checkpoint values. + /// The Cosmos DB account endpoint URI. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// A new instance of . + /// Thrown when any string parameter is null or whitespace. + public static CosmosCheckpointStore CreateCheckpointStoreUsingManagedIdentity( + string accountEndpoint, + string databaseId, + string containerId) + { + if (string.IsNullOrWhiteSpace(accountEndpoint)) + throw new ArgumentException("Cannot be null or whitespace", nameof(accountEndpoint)); + if (string.IsNullOrWhiteSpace(databaseId)) + throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); + if (string.IsNullOrWhiteSpace(containerId)) + throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); + + return new CosmosCheckpointStore(accountEndpoint, databaseId, containerId, useManagedIdentity: true); + } + + /// + /// Creates a generic Cosmos DB checkpoint store using an existing . + /// + /// The type of objects to store as checkpoint values. + /// The instance to use for Cosmos DB operations. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// A new instance of . + /// Thrown when any required parameter is null. + /// Thrown when any string parameter is null or whitespace. + public static CosmosCheckpointStore CreateCheckpointStore( + CosmosClient cosmosClient, + string databaseId, + string containerId) + { + if (cosmosClient is null) + throw new ArgumentNullException(nameof(cosmosClient)); + if (string.IsNullOrWhiteSpace(databaseId)) + throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); + if (string.IsNullOrWhiteSpace(containerId)) + throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); + + return new CosmosCheckpointStore(cosmosClient, databaseId, containerId); + } +} \ No newline at end of file diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Microsoft.Agents.AI.Workflows.csproj b/dotnet/src/Microsoft.Agents.AI.Workflows/Microsoft.Agents.AI.Workflows.csproj index ff2e9dee64..0ac25f0e1d 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Microsoft.Agents.AI.Workflows.csproj +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Microsoft.Agents.AI.Workflows.csproj @@ -31,6 +31,9 @@ + + + diff --git a/dotnet/src/Microsoft.Agents.AI/CosmosDBChatExtensions.cs b/dotnet/src/Microsoft.Agents.AI/CosmosDBChatExtensions.cs new file mode 100644 index 0000000000..ba42656bd6 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI/CosmosDBChatExtensions.cs @@ -0,0 +1,99 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using Microsoft.Azure.Cosmos; + +namespace Microsoft.Agents.AI; + +/// +/// Provides extension methods for integrating Cosmos DB chat message storage with the Agent Framework. +/// +public static class CosmosDBChatExtensions +{ + /// + /// Configures the agent to use Cosmos DB for message storage with connection string authentication. + /// + /// The chat client agent options to configure. + /// The Cosmos DB connection string. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// The configured . + /// Thrown when is null. + /// Thrown when any string parameter is null or whitespace. + public static ChatClientAgentOptions WithCosmosDBMessageStore( + this ChatClientAgentOptions options, + string connectionString, + string databaseId, + string containerId) + { + if (options is null) + throw new ArgumentNullException(nameof(options)); + if (string.IsNullOrWhiteSpace(connectionString)) + throw new ArgumentException("Cannot be null or whitespace", nameof(connectionString)); + if (string.IsNullOrWhiteSpace(databaseId)) + throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); + if (string.IsNullOrWhiteSpace(containerId)) + throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); + + options.ChatMessageStoreFactory = context => new CosmosChatMessageStore(connectionString, databaseId, containerId); + return options; + } + + /// + /// Configures the agent to use Cosmos DB for message storage with managed identity authentication. + /// + /// The chat client agent options to configure. + /// The Cosmos DB account endpoint URI. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// The configured . + /// Thrown when is null. + /// Thrown when any string parameter is null or whitespace. + public static ChatClientAgentOptions WithCosmosDBMessageStoreUsingManagedIdentity( + this ChatClientAgentOptions options, + string accountEndpoint, + string databaseId, + string containerId) + { + if (options is null) + throw new ArgumentNullException(nameof(options)); + if (string.IsNullOrWhiteSpace(accountEndpoint)) + throw new ArgumentException("Cannot be null or whitespace", nameof(accountEndpoint)); + if (string.IsNullOrWhiteSpace(databaseId)) + throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); + if (string.IsNullOrWhiteSpace(containerId)) + throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); + + options.ChatMessageStoreFactory = context => new CosmosChatMessageStore(accountEndpoint, databaseId, containerId, useManagedIdentity: true); + return options; + } + + /// + /// Configures the agent to use Cosmos DB for message storage with an existing . + /// + /// The chat client agent options to configure. + /// The instance to use for Cosmos DB operations. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// The configured . + /// Thrown when any required parameter is null. + /// Thrown when any string parameter is null or whitespace. + public static ChatClientAgentOptions WithCosmosDBMessageStore( + this ChatClientAgentOptions options, + CosmosClient cosmosClient, + string databaseId, + string containerId) + { + if (options is null) + throw new ArgumentNullException(nameof(options)); + if (cosmosClient is null) + throw new ArgumentNullException(nameof(cosmosClient)); + if (string.IsNullOrWhiteSpace(databaseId)) + throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); + if (string.IsNullOrWhiteSpace(containerId)) + throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); + + options.ChatMessageStoreFactory = context => new CosmosChatMessageStore(cosmosClient, databaseId, containerId); + return options; + } +} \ No newline at end of file diff --git a/dotnet/src/Microsoft.Agents.AI/Microsoft.Agents.AI.csproj b/dotnet/src/Microsoft.Agents.AI/Microsoft.Agents.AI.csproj index d95b8ea52f..b1a6b95c1b 100644 --- a/dotnet/src/Microsoft.Agents.AI/Microsoft.Agents.AI.csproj +++ b/dotnet/src/Microsoft.Agents.AI/Microsoft.Agents.AI.csproj @@ -21,6 +21,9 @@ + + + diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/CosmosChatMessageStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/CosmosChatMessageStoreTests.cs new file mode 100644 index 0000000000..388c753806 --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/CosmosChatMessageStoreTests.cs @@ -0,0 +1,387 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Linq; +using System.Threading.Tasks; +using Microsoft.Azure.Cosmos; +using Microsoft.Extensions.AI; +using Microsoft.Agents.AI; +using Xunit; + +namespace Microsoft.Agents.AI.Abstractions.UnitTests; + +/// +/// Contains tests for . +/// +/// Test Modes: +/// - Default Mode: Cleans up all test data after each test run (deletes database) +/// - Preserve Mode: Keeps containers and data for inspection in Cosmos DB Emulator Data Explorer +/// +/// To enable Preserve Mode, set environment variable: COSMOS_PRESERVE_CONTAINERS=true +/// Example: $env:COSMOS_PRESERVE_CONTAINERS="true"; dotnet test +/// +/// In Preserve Mode, you can view the data in Cosmos DB Emulator Data Explorer at: +/// https://localhost:8081/_explorer/index.html +/// Database: AgentFrameworkTests +/// Container: ChatMessages +/// +/// Environment Variable Reference: +/// | Variable | Values | Description | +/// |----------|--------|-------------| +/// | COSMOS_PRESERVE_CONTAINERS | true / false | Controls whether to preserve test data after completion | +/// +/// Usage Examples: +/// - Run all tests in preserve mode: $env:COSMOS_PRESERVE_CONTAINERS="true"; dotnet test tests/Microsoft.Agents.AI.Abstractions.UnitTests/ +/// - Run specific test category in preserve mode: $env:COSMOS_PRESERVE_CONTAINERS="true"; dotnet test tests/Microsoft.Agents.AI.Abstractions.UnitTests/ --filter "Category=CosmosDB" +/// - Reset to cleanup mode: $env:COSMOS_PRESERVE_CONTAINERS=""; dotnet test tests/Microsoft.Agents.AI.Abstractions.UnitTests/ +/// +[Collection("CosmosDB")] +public sealed class CosmosChatMessageStoreTests : IAsyncLifetime, IDisposable +{ + // Cosmos DB Emulator connection settings + private const string EmulatorEndpoint = "https://localhost:8081"; + private const string EmulatorKey = "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw=="; + private const string TestDatabaseId = "AgentFrameworkTests"; + private const string TestContainerId = "ChatMessages"; + + private CosmosClient? _cosmosClient; + private Database? _database; + private Container? _container; + private bool _emulatorAvailable; + private bool _preserveContainer; + + public async Task InitializeAsync() + { + // Check environment variable to determine if we should preserve containers + // Set COSMOS_PRESERVE_CONTAINERS=true to keep containers and data for inspection + _preserveContainer = string.Equals(Environment.GetEnvironmentVariable("COSMOS_PRESERVE_CONTAINERS"), "true", StringComparison.OrdinalIgnoreCase); + + try + { + _cosmosClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); + + // Test connection by attempting to create database + _database = await _cosmosClient.CreateDatabaseIfNotExistsAsync(TestDatabaseId); + _container = await _database.CreateContainerIfNotExistsAsync( + TestContainerId, + "/conversationId", + throughput: 400); + + _emulatorAvailable = true; + } + catch (Exception) + { + // Emulator not available, tests will be skipped + _emulatorAvailable = false; + _cosmosClient?.Dispose(); + _cosmosClient = null; + } + } + + public async Task DisposeAsync() + { + if (_cosmosClient != null && _emulatorAvailable) + { + try + { + if (_preserveContainer) + { + // Preserve mode: Don't delete the database/container, keep data for inspection + // This allows viewing data in the Cosmos DB Emulator Data Explorer + // No cleanup needed - data persists for debugging + } + else + { + // Clean mode: Delete the test database and all data + await _database!.DeleteAsync(); + } + } + catch + { + // Ignore cleanup errors + } + finally + { + _cosmosClient.Dispose(); + } + } + } + + public void Dispose() + { + _cosmosClient?.Dispose(); + GC.SuppressFinalize(this); + } + + private void SkipIfEmulatorNotAvailable() + { + if (!_emulatorAvailable) + { + Assert.Fail("Cosmos DB Emulator is not available. Start the emulator to run these tests."); + } + } + + #region Constructor Tests + + [Fact] + [Trait("Category", "CosmosDB")] + public void Constructor_WithCosmosClient_ShouldCreateInstance() + { + // Arrange & Act + SkipIfEmulatorNotAvailable(); + + // Act + using var store = new CosmosChatMessageStore(_cosmosClient!, TestDatabaseId, TestContainerId, "test-conversation"); + + // Assert + Assert.NotNull(store); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public void Constructor_WithConnectionString_ShouldCreateInstance() + { + // Arrange + SkipIfEmulatorNotAvailable(); + var connectionString = $"AccountEndpoint={EmulatorEndpoint};AccountKey={EmulatorKey};"; + + // Act + using var store = new CosmosChatMessageStore(connectionString, TestDatabaseId, TestContainerId); + + // Assert + Assert.NotNull(store); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public void Constructor_WithNullCosmosClient_ShouldThrowArgumentNullException() + { + // Arrange & Act & Assert + Assert.Throws(() => + new CosmosChatMessageStore((CosmosClient)null!, TestDatabaseId, TestContainerId, "test-conversation")); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public void Constructor_WithEmptyConversationId_ShouldThrowArgumentException() + { + // Arrange & Act & Assert + SkipIfEmulatorNotAvailable(); + + Assert.Throws(() => + new CosmosChatMessageStore(_cosmosClient!, TestDatabaseId, TestContainerId, "")); + } + + #endregion + + #region AddMessagesAsync Tests + + [Fact] + [Trait("Category", "CosmosDB")] + public async Task AddMessagesAsync_WithSingleMessage_ShouldAddMessageAsync() + { + // Arrange + SkipIfEmulatorNotAvailable(); + var conversationId = Guid.NewGuid().ToString(); + using var store = new CosmosChatMessageStore(_cosmosClient!, TestDatabaseId, TestContainerId, conversationId); + var message = new ChatMessage(ChatRole.User, "Hello, world!"); + + // Act + await store.AddMessagesAsync([message]); + + // Wait a moment for eventual consistency + await Task.Delay(100); + + // Assert + var messages = await store.GetMessagesAsync(); + var messageList = messages.ToList(); + + // Simple assertion - if this fails, we know the deserialization is the issue + if (messageList.Count == 0) + { + // Let's check if we can find ANY items in the container for this conversation + var directQuery = new QueryDefinition("SELECT VALUE COUNT(1) FROM c WHERE c.conversationId = @conversationId") + .WithParameter("@conversationId", conversationId); + var countIterator = this._cosmosClient!.GetDatabase(TestDatabaseId).GetContainer(TestContainerId) + .GetItemQueryIterator(directQuery, requestOptions: new QueryRequestOptions + { + PartitionKey = new PartitionKey(conversationId) + }); + + var countResponse = await countIterator.ReadNextAsync(); + var count = countResponse.FirstOrDefault(); + + // Debug: Let's see what the raw query returns + var rawQuery = new QueryDefinition("SELECT * FROM c WHERE c.conversationId = @conversationId") + .WithParameter("@conversationId", conversationId); + var rawIterator = this._cosmosClient!.GetDatabase(TestDatabaseId).GetContainer(TestContainerId) + .GetItemQueryIterator(rawQuery, requestOptions: new QueryRequestOptions + { + PartitionKey = new PartitionKey(conversationId) + }); + + List rawResults = new List(); + while (rawIterator.HasMoreResults) + { + var rawResponse = await rawIterator.ReadNextAsync(); + rawResults.AddRange(rawResponse); + } + + string rawJson = rawResults.Count > 0 ? Newtonsoft.Json.JsonConvert.SerializeObject(rawResults[0], Newtonsoft.Json.Formatting.Indented) : "null"; + Assert.Fail($"GetMessagesAsync returned 0 messages, but direct count query found {count} items for conversation {conversationId}. Raw document: {rawJson}"); + } + + Assert.Single(messageList); + Assert.Equal("Hello, world!", messageList[0].Text); + Assert.Equal(ChatRole.User, messageList[0].Role); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public async Task AddMessagesAsync_WithMultipleMessages_ShouldAddAllMessagesAsync() + { + // Arrange + SkipIfEmulatorNotAvailable(); + var conversationId = Guid.NewGuid().ToString(); + using var store = new CosmosChatMessageStore(_cosmosClient!, TestDatabaseId, TestContainerId, conversationId); + var messages = new[] + { + new ChatMessage(ChatRole.User, "First message"), + new ChatMessage(ChatRole.Assistant, "Second message"), + new ChatMessage(ChatRole.User, "Third message") + }; + + // Act + await store.AddMessagesAsync(messages); + + // Assert + var retrievedMessages = await store.GetMessagesAsync(); + var messageList = retrievedMessages.ToList(); + Assert.Equal(3, messageList.Count); + Assert.Equal("First message", messageList[0].Text); + Assert.Equal("Second message", messageList[1].Text); + Assert.Equal("Third message", messageList[2].Text); + } + + #endregion + + #region GetMessagesAsync Tests + + [Fact] + [Trait("Category", "CosmosDB")] + public async Task GetMessagesAsync_WithNoMessages_ShouldReturnEmptyAsync() + { + // Arrange + SkipIfEmulatorNotAvailable(); + using var store = new CosmosChatMessageStore(_cosmosClient!, TestDatabaseId, TestContainerId, Guid.NewGuid().ToString()); + + // Act + var messages = await store.GetMessagesAsync(); + + // Assert + Assert.Empty(messages); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public async Task GetMessagesAsync_WithConversationIsolation_ShouldOnlyReturnMessagesForConversationAsync() + { + // Arrange + SkipIfEmulatorNotAvailable(); + var conversation1 = Guid.NewGuid().ToString(); + var conversation2 = Guid.NewGuid().ToString(); + + using var store1 = new CosmosChatMessageStore(_cosmosClient!, TestDatabaseId, TestContainerId, conversation1); + using var store2 = new CosmosChatMessageStore(_cosmosClient!, TestDatabaseId, TestContainerId, conversation2); + + await store1.AddMessagesAsync([new ChatMessage(ChatRole.User, "Message for conversation 1")]); + await store2.AddMessagesAsync([new ChatMessage(ChatRole.User, "Message for conversation 2")]); + + // Act + var messages1 = await store1.GetMessagesAsync(); + var messages2 = await store2.GetMessagesAsync(); + + // Assert + var messageList1 = messages1.ToList(); + var messageList2 = messages2.ToList(); + Assert.Single(messageList1); + Assert.Single(messageList2); + Assert.Equal("Message for conversation 1", messageList1[0].Text); + Assert.Equal("Message for conversation 2", messageList2[0].Text); + } + + #endregion + + #region Integration Tests + + [Fact] + [Trait("Category", "CosmosDB")] + public async Task FullWorkflow_AddAndGet_ShouldWorkCorrectlyAsync() + { + // Arrange + SkipIfEmulatorNotAvailable(); + using var originalStore = new CosmosChatMessageStore(_cosmosClient!, TestDatabaseId, TestContainerId, "test-conversation"); + + var messages = new[] + { + new ChatMessage(ChatRole.System, "You are a helpful assistant."), + new ChatMessage(ChatRole.User, "Hello!"), + new ChatMessage(ChatRole.Assistant, "Hi there! How can I help you today?"), + new ChatMessage(ChatRole.User, "What's the weather like?"), + new ChatMessage(ChatRole.Assistant, "I'm sorry, I don't have access to current weather data.") + }; + + // Act 1: Add messages + await originalStore.AddMessagesAsync(messages); + + // Act 2: Verify messages were added + var retrievedMessages = await originalStore.GetMessagesAsync(); + var retrievedList = retrievedMessages.ToList(); + Assert.Equal(5, retrievedList.Count); + + // Act 3: Create new store instance for same conversation (test persistence) + using var newStore = new CosmosChatMessageStore(_cosmosClient!, TestDatabaseId, TestContainerId, "test-conversation"); + var persistedMessages = await newStore.GetMessagesAsync(); + var persistedList = persistedMessages.ToList(); + + // Assert final state + Assert.Equal(5, persistedList.Count); + Assert.Equal("You are a helpful assistant.", persistedList[0].Text); + Assert.Equal("Hello!", persistedList[1].Text); + Assert.Equal("Hi there! How can I help you today?", persistedList[2].Text); + Assert.Equal("What's the weather like?", persistedList[3].Text); + Assert.Equal("I'm sorry, I don't have access to current weather data.", persistedList[4].Text); + } + + #endregion + + #region Disposal Tests + + [Fact] + [Trait("Category", "CosmosDB")] + public void Dispose_AfterUse_ShouldNotThrow() + { + // Arrange + SkipIfEmulatorNotAvailable(); + var store = new CosmosChatMessageStore(_cosmosClient!, TestDatabaseId, TestContainerId, Guid.NewGuid().ToString()); + + // Act & Assert + store.Dispose(); // Should not throw + } + + [Fact] + [Trait("Category", "CosmosDB")] + public void Dispose_MultipleCalls_ShouldNotThrow() + { + // Arrange + SkipIfEmulatorNotAvailable(); + var store = new CosmosChatMessageStore(_cosmosClient!, TestDatabaseId, TestContainerId, Guid.NewGuid().ToString()); + + // Act & Assert + store.Dispose(); // First call + store.Dispose(); // Second call - should not throw + } + + #endregion +} \ No newline at end of file diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/Microsoft.Agents.AI.Abstractions.UnitTests.csproj b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/Microsoft.Agents.AI.Abstractions.UnitTests.csproj index b7c5412a53..63b0bc4007 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/Microsoft.Agents.AI.Abstractions.UnitTests.csproj +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/Microsoft.Agents.AI.Abstractions.UnitTests.csproj @@ -16,6 +16,8 @@ + + diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Microsoft.Agents.AI.UnitTests.csproj b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Microsoft.Agents.AI.UnitTests.csproj index f871781d03..a567b8be58 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Microsoft.Agents.AI.UnitTests.csproj +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Microsoft.Agents.AI.UnitTests.csproj @@ -18,6 +18,9 @@ + + + diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/CosmosCheckpointStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/CosmosCheckpointStoreTests.cs new file mode 100644 index 0000000000..483988ba3b --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/CosmosCheckpointStoreTests.cs @@ -0,0 +1,436 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Linq; +using System.Text.Json; +using System.Threading.Tasks; +using Microsoft.Azure.Cosmos; +using Microsoft.Agents.AI.Workflows; +using Microsoft.Agents.AI.Workflows.Checkpointing; +using Xunit; + +namespace Microsoft.Agents.AI.Workflows.UnitTests; + +/// +/// Contains tests for . +/// +/// Test Modes: +/// - Default Mode: Cleans up all test data after each test run (deletes database) +/// - Preserve Mode: Keeps containers and data for inspection in Cosmos DB Emulator Data Explorer +/// +/// To enable Preserve Mode, set environment variable: COSMOS_PRESERVE_CONTAINERS=true +/// Example: $env:COSMOS_PRESERVE_CONTAINERS="true"; dotnet test +/// +/// In Preserve Mode, you can view the data in Cosmos DB Emulator Data Explorer at: +/// https://localhost:8081/_explorer/index.html +/// Database: AgentFrameworkTests +/// Container: Checkpoints +/// +public class CosmosCheckpointStoreTests : IAsyncLifetime, IDisposable +{ + // Cosmos DB Emulator connection settings + private const string EmulatorEndpoint = "https://localhost:8081"; + private const string EmulatorKey = "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw=="; + private const string TestDatabaseId = "AgentFrameworkTests"; + private const string TestContainerId = "Checkpoints"; + + private CosmosClient? _cosmosClient; + private Database? _database; + private Container? _container; + private bool _emulatorAvailable; + private bool _preserveContainer; + + public async Task InitializeAsync() + { + // Check environment variable to determine if we should preserve containers + // Set COSMOS_PRESERVE_CONTAINERS=true to keep containers and data for inspection + _preserveContainer = string.Equals(Environment.GetEnvironmentVariable("COSMOS_PRESERVE_CONTAINERS"), "true", StringComparison.OrdinalIgnoreCase); + + try + { + _cosmosClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); + + // Test connection by attempting to create database + _database = await _cosmosClient.CreateDatabaseIfNotExistsAsync(TestDatabaseId); + _container = await _database.CreateContainerIfNotExistsAsync( + TestContainerId, + "/runId", + throughput: 400); + + _emulatorAvailable = true; + } + catch (Exception) + { + // Emulator not available, tests will be skipped + _emulatorAvailable = false; + _cosmosClient?.Dispose(); + _cosmosClient = null; + } + } + + public async Task DisposeAsync() + { + if (_cosmosClient != null && _emulatorAvailable) + { + try + { + if (_preserveContainer) + { + // Preserve mode: Don't delete the database/container, keep data for inspection + // This allows viewing data in the Cosmos DB Emulator Data Explorer + // No cleanup needed - data persists for debugging + } + else + { + // Clean mode: Delete the test database and all data + await _database!.DeleteAsync(); + } + } + catch + { + // Ignore cleanup errors + } + finally + { + _cosmosClient.Dispose(); + } + } + } + + private void SkipIfEmulatorNotAvailable() + { + if (!_emulatorAvailable) + { + // For now let's just fail to see if tests work + Assert.Fail("Cosmos DB Emulator is not available. Start the emulator to run these tests."); + } + } + + #region Constructor Tests + + [Fact] + public void Constructor_WithCosmosClient_SetsProperties() + { + SkipIfEmulatorNotAvailable(); + + // Act + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + + // Assert + Assert.Equal(TestDatabaseId, store.DatabaseId); + Assert.Equal(TestContainerId, store.ContainerId); + } + + [Fact] + public void Constructor_WithConnectionString_SetsProperties() + { + SkipIfEmulatorNotAvailable(); + + // Arrange + var connectionString = $"AccountEndpoint={EmulatorEndpoint};AccountKey={EmulatorKey};"; + + // Act + using var store = new CosmosCheckpointStore(connectionString, TestDatabaseId, TestContainerId); + + // Assert + Assert.Equal(TestDatabaseId, store.DatabaseId); + Assert.Equal(TestContainerId, store.ContainerId); + } + + [Fact] + public void Constructor_WithNullCosmosClient_ThrowsArgumentNullException() + { + // Act & Assert + Assert.Throws(() => + new CosmosCheckpointStore((CosmosClient)null!, TestDatabaseId, TestContainerId)); + } + + [Fact] + public void Constructor_WithNullConnectionString_ThrowsArgumentException() + { + // Act & Assert + Assert.Throws(() => + new CosmosCheckpointStore((string)null!, TestDatabaseId, TestContainerId)); + } + + #endregion + + #region Checkpoint Operations Tests + + [Fact] + public async Task CreateCheckpointAsync_NewCheckpoint_CreatesSuccessfully() + { + SkipIfEmulatorNotAvailable(); + + // Arrange + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + var runId = Guid.NewGuid().ToString(); + var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test checkpoint" }); + + // Act + var checkpointInfo = await store.CreateCheckpointAsync(runId, checkpointValue); + + // Assert + Assert.NotNull(checkpointInfo); + Assert.Equal(runId, checkpointInfo.RunId); + Assert.NotNull(checkpointInfo.CheckpointId); + Assert.NotEmpty(checkpointInfo.CheckpointId); + } + + [Fact] + public async Task RetrieveCheckpointAsync_ExistingCheckpoint_ReturnsCorrectValue() + { + SkipIfEmulatorNotAvailable(); + + // Arrange + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + var runId = Guid.NewGuid().ToString(); + var originalData = new { message = "Hello, World!", timestamp = DateTimeOffset.UtcNow }; + var checkpointValue = JsonSerializer.SerializeToElement(originalData); + + // Act + var checkpointInfo = await store.CreateCheckpointAsync(runId, checkpointValue); + var retrievedValue = await store.RetrieveCheckpointAsync(runId, checkpointInfo); + + // Assert + Assert.Equal(JsonValueKind.Object, retrievedValue.ValueKind); + Assert.True(retrievedValue.TryGetProperty("message", out var messageProp)); + Assert.Equal("Hello, World!", messageProp.GetString()); + } + + [Fact] + public async Task RetrieveCheckpointAsync_NonExistentCheckpoint_ThrowsInvalidOperationException() + { + SkipIfEmulatorNotAvailable(); + + // Arrange + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + var runId = Guid.NewGuid().ToString(); + var fakeCheckpointInfo = new CheckpointInfo(runId, "nonexistent-checkpoint"); + + // Act & Assert + await Assert.ThrowsAsync(() => + store.RetrieveCheckpointAsync(runId, fakeCheckpointInfo).AsTask()); + } + + [Fact] + public async Task RetrieveIndexAsync_EmptyStore_ReturnsEmptyCollection() + { + SkipIfEmulatorNotAvailable(); + + // Arrange + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + var runId = Guid.NewGuid().ToString(); + + // Act + var index = await store.RetrieveIndexAsync(runId); + + // Assert + Assert.NotNull(index); + Assert.Empty(index); + } + + [Fact] + public async Task RetrieveIndexAsync_WithCheckpoints_ReturnsAllCheckpoints() + { + SkipIfEmulatorNotAvailable(); + + // Arrange + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + var runId = Guid.NewGuid().ToString(); + var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }); + + // Create multiple checkpoints + var checkpoint1 = await store.CreateCheckpointAsync(runId, checkpointValue); + var checkpoint2 = await store.CreateCheckpointAsync(runId, checkpointValue); + var checkpoint3 = await store.CreateCheckpointAsync(runId, checkpointValue); + + // Act + var index = (await store.RetrieveIndexAsync(runId)).ToList(); + + // Assert + Assert.Equal(3, index.Count); + Assert.Contains(index, c => c.CheckpointId == checkpoint1.CheckpointId); + Assert.Contains(index, c => c.CheckpointId == checkpoint2.CheckpointId); + Assert.Contains(index, c => c.CheckpointId == checkpoint3.CheckpointId); + } + + [Fact] + public async Task CreateCheckpointAsync_WithParent_CreatesHierarchy() + { + SkipIfEmulatorNotAvailable(); + + // Arrange + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + var runId = Guid.NewGuid().ToString(); + var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }); + + // Act + var parentCheckpoint = await store.CreateCheckpointAsync(runId, checkpointValue); + var childCheckpoint = await store.CreateCheckpointAsync(runId, checkpointValue, parentCheckpoint); + + // Assert + Assert.NotEqual(parentCheckpoint.CheckpointId, childCheckpoint.CheckpointId); + Assert.Equal(runId, parentCheckpoint.RunId); + Assert.Equal(runId, childCheckpoint.RunId); + } + + [Fact] + public async Task RetrieveIndexAsync_WithParentFilter_ReturnsFilteredResults() + { + SkipIfEmulatorNotAvailable(); + + // Arrange + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + var runId = Guid.NewGuid().ToString(); + var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }); + + // Create parent and child checkpoints + var parent = await store.CreateCheckpointAsync(runId, checkpointValue); + var child1 = await store.CreateCheckpointAsync(runId, checkpointValue, parent); + var child2 = await store.CreateCheckpointAsync(runId, checkpointValue, parent); + + // Create an orphan checkpoint + var orphan = await store.CreateCheckpointAsync(runId, checkpointValue); + + // Act + var allCheckpoints = (await store.RetrieveIndexAsync(runId)).ToList(); + var childrenOfParent = (await store.RetrieveIndexAsync(runId, parent)).ToList(); + + // Assert + Assert.Equal(4, allCheckpoints.Count); // parent + 2 children + orphan + Assert.Equal(2, childrenOfParent.Count); // only children + + Assert.Contains(childrenOfParent, c => c.CheckpointId == child1.CheckpointId); + Assert.Contains(childrenOfParent, c => c.CheckpointId == child2.CheckpointId); + Assert.DoesNotContain(childrenOfParent, c => c.CheckpointId == parent.CheckpointId); + Assert.DoesNotContain(childrenOfParent, c => c.CheckpointId == orphan.CheckpointId); + } + + #endregion + + #region Run Isolation Tests + + [Fact] + public async Task CheckpointOperations_DifferentRuns_IsolatesData() + { + SkipIfEmulatorNotAvailable(); + + // Arrange + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + var runId1 = Guid.NewGuid().ToString(); + var runId2 = Guid.NewGuid().ToString(); + var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }); + + // Act + var checkpoint1 = await store.CreateCheckpointAsync(runId1, checkpointValue); + var checkpoint2 = await store.CreateCheckpointAsync(runId2, checkpointValue); + + var index1 = (await store.RetrieveIndexAsync(runId1)).ToList(); + var index2 = (await store.RetrieveIndexAsync(runId2)).ToList(); + + // Assert + Assert.Single(index1); + Assert.Single(index2); + Assert.Equal(checkpoint1.CheckpointId, index1[0].CheckpointId); + Assert.Equal(checkpoint2.CheckpointId, index2[0].CheckpointId); + Assert.NotEqual(checkpoint1.CheckpointId, checkpoint2.CheckpointId); + } + + #endregion + + #region Error Handling Tests + + [Fact] + public async Task CreateCheckpointAsync_WithNullRunId_ThrowsArgumentException() + { + SkipIfEmulatorNotAvailable(); + + // Arrange + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }); + + // Act & Assert + await Assert.ThrowsAsync(() => + store.CreateCheckpointAsync(null!, checkpointValue).AsTask()); + } + + [Fact] + public async Task CreateCheckpointAsync_WithEmptyRunId_ThrowsArgumentException() + { + SkipIfEmulatorNotAvailable(); + + // Arrange + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }); + + // Act & Assert + await Assert.ThrowsAsync(() => + store.CreateCheckpointAsync("", checkpointValue).AsTask()); + } + + [Fact] + public async Task RetrieveCheckpointAsync_WithNullCheckpointInfo_ThrowsArgumentNullException() + { + SkipIfEmulatorNotAvailable(); + + // Arrange + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + var runId = Guid.NewGuid().ToString(); + + // Act & Assert + await Assert.ThrowsAsync(() => + store.RetrieveCheckpointAsync(runId, null!).AsTask()); + } + + #endregion + + #region Disposal Tests + + [Fact] + public async Task Dispose_AfterDisposal_ThrowsObjectDisposedException() + { + SkipIfEmulatorNotAvailable(); + + // Arrange + var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }); + + // Act + store.Dispose(); + + // Assert + await Assert.ThrowsAsync(() => + store.CreateCheckpointAsync("test-run", checkpointValue).AsTask()); + } + + [Fact] + public void Dispose_MultipleCalls_DoesNotThrow() + { + SkipIfEmulatorNotAvailable(); + + // Arrange + var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + + // Act & Assert (should not throw) + store.Dispose(); + store.Dispose(); + store.Dispose(); + } + + #endregion + + public void Dispose() + { + Dispose(true); + GC.SuppressFinalize(this); + } + + protected virtual void Dispose(bool disposing) + { + if (disposing) + { + this._cosmosClient?.Dispose(); + } + } +} \ No newline at end of file diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Microsoft.Agents.AI.Workflows.UnitTests.csproj b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Microsoft.Agents.AI.Workflows.UnitTests.csproj index bd9bc57915..cb4a81e5c5 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Microsoft.Agents.AI.Workflows.UnitTests.csproj +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Microsoft.Agents.AI.Workflows.UnitTests.csproj @@ -14,6 +14,8 @@ + + From 2208d3920ac7ccdaf1984e0334a7729c9a872e5e Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Thu, 30 Oct 2025 21:27:18 +0000 Subject: [PATCH 02/40] Added Cosmos agent thread and tests --- .../CosmosAgentThread.cs | 233 ++++++++ .../CosmosChatMessageStore.cs | 66 ++- .../CosmosAgentThreadTests.cs | 507 ++++++++++++++++++ .../CosmosChatMessageStoreTests.cs | 80 +-- 4 files changed, 822 insertions(+), 64 deletions(-) create mode 100644 dotnet/src/Microsoft.Agents.AI.Abstractions/CosmosAgentThread.cs create mode 100644 dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/CosmosAgentThreadTests.cs diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/CosmosAgentThread.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/CosmosAgentThread.cs new file mode 100644 index 0000000000..9c3be1ed30 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/CosmosAgentThread.cs @@ -0,0 +1,233 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Diagnostics; +using System.Text.Json; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Azure.Cosmos; +using Microsoft.Extensions.AI; + +namespace Microsoft.Agents.AI; + +/// +/// Provides an abstract base class for agent threads that maintain conversation state in Azure Cosmos DB. +/// +/// +/// +/// is designed for scenarios where conversation state should be persisted +/// in Azure Cosmos DB for durability, scalability, and cross-session availability. This approach provides +/// reliable persistence while maintaining efficient access to conversation data. +/// +/// +/// Cosmos threads persist conversation data across application restarts and can be shared across +/// multiple application instances. +/// +/// +[DebuggerDisplay("{DebuggerDisplay,nq}")] +public abstract class CosmosAgentThread : AgentThread +{ + /// + /// Initializes a new instance of the class. + /// + /// The Cosmos DB connection string. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// The unique identifier for this conversation thread. If null, a new GUID will be generated. + /// Thrown when any string parameter is null or whitespace. + /// + /// This constructor creates a new Cosmos DB message store with connection string authentication. + /// + protected CosmosAgentThread(string connectionString, string databaseId, string containerId, string? conversationId = null) + { + this.MessageStore = new CosmosChatMessageStore(connectionString, databaseId, containerId, conversationId ?? Guid.NewGuid().ToString("N")); + } + + /// + /// Initializes a new instance of the class using managed identity. + /// + /// The Cosmos DB account endpoint URI. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// The unique identifier for this conversation thread. If null, a new GUID will be generated. + /// Must be true to use this constructor. This parameter distinguishes this constructor from the connection string version. + /// Thrown when any string parameter is null or whitespace, or when useManagedIdentity is false. + /// + /// This constructor creates a new Cosmos DB message store with managed identity authentication. + /// + protected CosmosAgentThread(string accountEndpoint, string databaseId, string containerId, string? conversationId, bool useManagedIdentity) + { + if (!useManagedIdentity) + { + throw new ArgumentException("This constructor requires useManagedIdentity to be true.", nameof(useManagedIdentity)); + } + + this.MessageStore = new CosmosChatMessageStore(accountEndpoint, databaseId, containerId, conversationId ?? Guid.NewGuid().ToString("N"), useManagedIdentity: true); + } + + /// + /// Initializes a new instance of the class using an existing CosmosClient. + /// + /// The instance to use for Cosmos DB operations. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// The unique identifier for this conversation thread. If null, a new GUID will be generated. + /// Thrown when is null. + /// Thrown when any string parameter is null or whitespace. + /// + /// This constructor allows reuse of an existing CosmosClient instance across multiple threads. + /// + protected CosmosAgentThread(CosmosClient cosmosClient, string databaseId, string containerId, string? conversationId = null) + { + this.MessageStore = new CosmosChatMessageStore(cosmosClient, databaseId, containerId, conversationId ?? Guid.NewGuid().ToString("N")); + } + + /// + /// Initializes a new instance of the class with a pre-configured message store. + /// + /// + /// A instance to use for storing chat messages. + /// If , an exception will be thrown. + /// + /// is . + /// + /// This constructor allows sharing of message stores between threads or providing pre-configured + /// message stores with specific settings. + /// + protected CosmosAgentThread(CosmosChatMessageStore messageStore) + { + this.MessageStore = messageStore ?? throw new ArgumentNullException(nameof(messageStore)); + } + + /// + /// Initializes a new instance of the class from previously serialized state. + /// + /// A representing the serialized state of the thread. + /// + /// Factory function to create the from its serialized state. + /// This is required because Cosmos DB connection information cannot be reconstructed from serialized state. + /// + /// Optional settings for customizing the JSON deserialization process. + /// The is not a JSON object. + /// is . + /// The is invalid or cannot be deserialized to the expected type. + /// + /// This constructor enables restoration of Cosmos threads from previously saved state. Since Cosmos DB + /// connection information cannot be serialized for security reasons, a factory function must be provided + /// to reconstruct the message store with appropriate connection details. + /// + protected CosmosAgentThread( + JsonElement serializedThreadState, + Func messageStoreFactory, + JsonSerializerOptions? jsonSerializerOptions = null) + { +#if NET6_0_OR_GREATER + ArgumentNullException.ThrowIfNull(messageStoreFactory); +#else + if (messageStoreFactory is null) + { + throw new ArgumentNullException(nameof(messageStoreFactory)); + } +#endif + + if (serializedThreadState.ValueKind != JsonValueKind.Object) + { + throw new ArgumentException("The serialized thread state must be a JSON object.", nameof(serializedThreadState)); + } + + var state = serializedThreadState.Deserialize( + AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(CosmosAgentThreadState))) as CosmosAgentThreadState; + + this.MessageStore = messageStoreFactory.Invoke(state?.StoreState ?? default, jsonSerializerOptions); + } + + /// + /// Gets the used by this thread. + /// + public CosmosChatMessageStore MessageStore { get; } + + /// + /// Gets the conversation ID for this thread from the underlying message store. + /// + public string ConversationId => this.MessageStore.ConversationId; + + /// + /// Gets or sets the maximum number of messages to return in a single query batch. + /// This is delegated to the underlying message store. + /// + public int MaxItemCount + { + get => this.MessageStore.MaxItemCount; + set => this.MessageStore.MaxItemCount = value; + } + + /// + /// Gets or sets the Time-To-Live (TTL) in seconds for messages. + /// This is delegated to the underlying message store. + /// + public int? MessageTtlSeconds + { + get => this.MessageStore.MessageTtlSeconds; + set => this.MessageStore.MessageTtlSeconds = value; + } + + /// + /// Serializes the current object's state to a using the specified serialization options. + /// + /// The JSON serialization options to use. + /// A representation of the object's state. + /// + /// Note that connection strings and credentials are not included in the serialized state for security reasons. + /// When deserializing, you will need to provide connection information through the messageStoreFactory parameter. + /// + public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) + { + var storeState = this.MessageStore.Serialize(jsonSerializerOptions); + + var state = new CosmosAgentThreadState + { + StoreState = storeState, + }; + + return JsonSerializer.SerializeToElement(state, AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(CosmosAgentThreadState))); + } + + /// + public override object? GetService(Type serviceType, object? serviceKey = null) => + base.GetService(serviceType, serviceKey) ?? this.MessageStore?.GetService(serviceType, serviceKey); + + /// + protected internal override Task MessagesReceivedAsync(IEnumerable newMessages, CancellationToken cancellationToken = default) + => this.MessageStore.AddMessagesAsync(newMessages, cancellationToken); + + /// + /// Gets the total number of messages in this conversation. + /// This is a Cosmos-specific optimization that provides efficient message counting. + /// + /// The cancellation token. + /// The total number of messages in the conversation. + public async Task GetMessageCountAsync(CancellationToken cancellationToken = default) + { + return await this.MessageStore.GetMessageCountAsync(cancellationToken).ConfigureAwait(false); + } + + /// + /// Clears all messages in this conversation. + /// This is a Cosmos-specific utility method for conversation cleanup. + /// + /// The cancellation token. + /// The number of messages that were deleted. + public async Task ClearMessagesAsync(CancellationToken cancellationToken = default) + { + return await this.MessageStore.ClearMessagesAsync(cancellationToken).ConfigureAwait(false); + } + + [DebuggerBrowsable(DebuggerBrowsableState.Never)] + private string DebuggerDisplay => $"ConversationId = {this.ConversationId}"; + + internal sealed class CosmosAgentThreadState + { + public JsonElement? StoreState { get; set; } + } +} \ No newline at end of file diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/CosmosChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/CosmosChatMessageStore.cs index 0f00885c35..9fa5c343cd 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/CosmosChatMessageStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/CosmosChatMessageStore.cs @@ -22,6 +22,8 @@ public sealed class CosmosChatMessageStore : ChatMessageStore, IDisposable private readonly CosmosClient _cosmosClient; private readonly Container _container; private readonly string _conversationId; + private readonly string _databaseId; + private readonly string _containerId; private readonly bool _ownsClient; private bool _disposed; @@ -54,6 +56,21 @@ private static JsonSerializerOptions CreateDefaultJsonOptions() /// public int? MessageTtlSeconds { get; set; } = 86400; + /// + /// Gets the conversation ID associated with this message store. + /// + public string ConversationId => this._conversationId; + + /// + /// Gets the database ID associated with this message store. + /// + public string DatabaseId => this._databaseId; + + /// + /// Gets the container ID associated with this message store. + /// + public string ContainerId => this._containerId; + /// /// Initializes a new instance of the class using a connection string. /// @@ -101,6 +118,8 @@ public CosmosChatMessageStore(string connectionString, string databaseId, string this._cosmosClient = new CosmosClient(connectionString); this._container = this._cosmosClient.GetContainer(databaseId, containerId); this._conversationId = conversationId; + this._databaseId = databaseId; + this._containerId = containerId; this._ownsClient = true; } @@ -158,6 +177,8 @@ public CosmosChatMessageStore(string accountEndpoint, string databaseId, string this._cosmosClient = new CosmosClient(accountEndpoint, new DefaultAzureCredential()); this._container = this._cosmosClient.GetContainer(databaseId, containerId); this._conversationId = conversationId; + this._databaseId = databaseId; + this._containerId = containerId; this._ownsClient = true; } @@ -205,6 +226,8 @@ public CosmosChatMessageStore(CosmosClient cosmosClient, string databaseId, stri this._container = this._cosmosClient.GetContainer(databaseId, containerId); this._conversationId = conversationId; + this._databaseId = databaseId; + this._containerId = containerId; } /// @@ -225,9 +248,11 @@ public CosmosChatMessageStore(JsonElement serializedStoreState, CosmosClient cos if (serializedStoreState.ValueKind is JsonValueKind.Object) { var state = JsonSerializer.Deserialize(serializedStoreState, jsonSerializerOptions); - if (state?.ConversationId is { } conversationId && state.DatabaseId is { } databaseId && state.ContainerId is { } containerId) + if (state?.ConversationIdentifier is { } conversationId && state.DatabaseIdentifier is { } databaseId && state.ContainerIdentifier is { } containerId) { this._conversationId = conversationId; + this._databaseId = databaseId; + this._containerId = containerId; this._container = this._cosmosClient.GetContainer(databaseId, containerId); return; } @@ -236,21 +261,6 @@ public CosmosChatMessageStore(JsonElement serializedStoreState, CosmosClient cos throw new ArgumentException("Invalid serialized state", nameof(serializedStoreState)); } - /// - /// Gets the unique identifier for this conversation thread. - /// - public string ConversationId => this._conversationId; - - /// - /// Gets the identifier of the Cosmos DB database. - /// - public string DatabaseId => this._container.Database.Id; - - /// - /// Gets the identifier of the Cosmos DB container. - /// - public string ContainerId => this._container.Id; - /// [UnconditionalSuppressMessage("Trimming", "IL2026:Members annotated with 'RequiresUnreferencedCodeAttribute' require dynamic access", Justification = "ChatMessage deserialization is controlled")] [UnconditionalSuppressMessage("AOT", "IL3050:Calling members annotated with 'RequiresDynamicCodeAttribute' may break functionality when AOT compiling", Justification = "ChatMessage deserialization is controlled")] @@ -406,9 +416,9 @@ public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptio var state = new StoreState { - ConversationId = this._conversationId, - DatabaseId = this.DatabaseId, - ContainerId = this.ContainerId + ConversationIdentifier = this._conversationId, + DatabaseIdentifier = this.DatabaseId, + ContainerIdentifier = this.ContainerId }; var options = jsonSerializerOptions ?? s_defaultJsonOptions; @@ -431,7 +441,7 @@ public async Task GetMessageCountAsync(CancellationToken cancellationToken #pragma warning restore CA1513 // Efficient count query - var query = new QueryDefinition("SELECT VALUE COUNT(1) FROM c WHERE c.conversationId = @conversationId AND c.type = @type") + var query = new QueryDefinition("SELECT VALUE COUNT(1) FROM c WHERE c.conversationId = @conversationId AND c.Type = @type") .WithParameter("@conversationId", this._conversationId) .WithParameter("@type", "ChatMessage"); @@ -465,11 +475,11 @@ public async Task ClearMessagesAsync(CancellationToken cancellationToken = #pragma warning restore CA1513 // Batch delete for efficiency - var query = new QueryDefinition("SELECT c.id FROM c WHERE c.conversationId = @conversationId AND c.type = @type") + var query = new QueryDefinition("SELECT VALUE c.id FROM c WHERE c.conversationId = @conversationId AND c.Type = @type") .WithParameter("@conversationId", this._conversationId) .WithParameter("@type", "ChatMessage"); - var iterator = this._container.GetItemQueryIterator(query, requestOptions: new QueryRequestOptions + var iterator = this._container.GetItemQueryIterator(query, requestOptions: new QueryRequestOptions { PartitionKey = new PartitionKey(this._conversationId), MaxItemCount = this.MaxItemCount @@ -484,11 +494,11 @@ public async Task ClearMessagesAsync(CancellationToken cancellationToken = var batch = this._container.CreateTransactionalBatch(partitionKey); var batchItemCount = 0; - foreach (var item in response) + foreach (var itemId in response) { - if (item.TryGetProperty("id", out var idElement) && idElement.ValueKind == JsonValueKind.String) + if (!string.IsNullOrEmpty(itemId)) { - batch.DeleteItem(idElement.GetString()); + batch.DeleteItem(itemId); batchItemCount++; deletedCount++; } @@ -518,9 +528,9 @@ public void Dispose() private sealed class StoreState { - public string ConversationId { get; set; } = string.Empty; - public string DatabaseId { get; set; } = string.Empty; - public string ContainerId { get; set; } = string.Empty; + public string ConversationIdentifier { get; set; } = string.Empty; + public string DatabaseIdentifier { get; set; } = string.Empty; + public string ContainerIdentifier { get; set; } = string.Empty; } /// diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/CosmosAgentThreadTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/CosmosAgentThreadTests.cs new file mode 100644 index 0000000000..eecb3424b0 --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/CosmosAgentThreadTests.cs @@ -0,0 +1,507 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Linq; +using System.Text.Json; +using System.Threading.Tasks; +using Microsoft.Azure.Cosmos; +using Microsoft.Extensions.AI; +using Microsoft.Agents.AI; +using Xunit; + +namespace Microsoft.Agents.AI.Abstractions.UnitTests; + +/// +/// Collection definition to ensure Cosmos DB tests don't run in parallel +/// +[CollectionDefinition("CosmosDB", DisableParallelization = true)] +public class CosmosDBTestFixture : ICollectionFixture +{ +} + +/// +/// Contains tests for . +/// +/// These tests use the connection string constructor approach instead of manually creating CosmosClient instances, +/// making the tests simpler and more realistic. +/// +/// Test Modes: +/// - Default Mode: Cleans up all test data after each test run (deletes database) +/// - Preserve Mode: Keeps containers and data for inspection in Cosmos DB Emulator Data Explorer +/// +/// Usage Examples: +/// - Run all tests in preserve mode: $env:COSMOS_PRESERVE_CONTAINERS="true"; dotnet test tests/Microsoft.Agents.AI.Abstractions.UnitTests/ +/// - Run specific test category in preserve mode: $env:COSMOS_PRESERVE_CONTAINERS="true"; dotnet test tests/Microsoft.Agents.AI.Abstractions.UnitTests/ --filter "Category=CosmosDB" +/// - Reset to cleanup mode: $env:COSMOS_PRESERVE_CONTAINERS=""; dotnet test tests/Microsoft.Agents.AI.Abstractions.UnitTests/ +/// +[Collection("CosmosDB")] +[TestCaseOrderer("Xunit.Extensions.Ordering.TestCaseOrderer", "Xunit.Extensions.Ordering")] +public sealed class CosmosAgentThreadTests : IAsyncLifetime, IDisposable +{ + // Cosmos DB Emulator connection settings + private const string EmulatorEndpoint = "https://localhost:8081"; + private const string EmulatorKey = "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw=="; + + // Use static container names like CosmosChatMessageStoreTests for consistency + private readonly string _testDatabaseId = "AgentFrameworkTests"; + private readonly string _testContainerId = "ChatMessages"; // Use the same container as CosmosChatMessageStoreTests + + private string _connectionString = string.Empty; + private bool _emulatorAvailable; + private bool _preserveContainer; + private CosmosClient? _setupClient; // Only used for test setup/cleanup + + public async Task InitializeAsync() + { + // Check environment variable to determine if we should preserve containers + // Set COSMOS_PRESERVE_CONTAINERS=true to keep containers and data for inspection + this._preserveContainer = string.Equals(Environment.GetEnvironmentVariable("COSMOS_PRESERVE_CONTAINERS"), "true", StringComparison.OrdinalIgnoreCase); + + this._connectionString = $"AccountEndpoint={EmulatorEndpoint};AccountKey={EmulatorKey}"; + + try + { + // Only create CosmosClient for test setup - the actual tests will use connection string constructors + this._setupClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); + + // Test connection and ensure database/container exist + var databaseResponse = await this._setupClient.CreateDatabaseIfNotExistsAsync(this._testDatabaseId); + var containerResponse = await databaseResponse.Database.CreateContainerIfNotExistsAsync( + this._testContainerId, + "/conversationId", + throughput: 400); + + // Verify the container is actually accessible by doing a read operation + await containerResponse.Container.ReadContainerAsync(); + + // Wait a moment for container to be fully propagated across all clients + await Task.Delay(500); + + // Verify the container is accessible from a new client instance (simulating the test scenario) + using var verificationClient = new CosmosClient(this._connectionString); + var verificationContainer = verificationClient.GetContainer(this._testDatabaseId, this._testContainerId); + await verificationContainer.ReadContainerAsync(); + + this._emulatorAvailable = true; + } + catch (Exception) + { + // Emulator not available, tests will be skipped + this._emulatorAvailable = false; + this._setupClient?.Dispose(); + this._setupClient = null; + } + } + + public async Task DisposeAsync() + { + if (this._setupClient != null && this._emulatorAvailable) + { + try + { + // Clean up test database unless preserve mode is enabled + if (!this._preserveContainer) + { + // Delete the entire test database to clean up all containers + var database = this._setupClient.GetDatabase(this._testDatabaseId); + await database.DeleteAsync(); + } + } + catch (Exception ex) + { + // Ignore cleanup errors during test teardown + Console.WriteLine($"Warning: Cleanup failed: {ex.Message}"); + } + finally + { + this._setupClient.Dispose(); + } + } + } + + /// + /// Implements IDisposable to properly dispose of the setup client. + /// + public void Dispose() + { + this._setupClient?.Dispose(); + } + + private void SkipIfEmulatorNotAvailable() + { + if (!this._emulatorAvailable) + { + Assert.True(true, "Cosmos DB Emulator not available, test skipped"); + } + } + + private async Task EnsureContainerReadyAsync() + { + if (this._emulatorAvailable) + { + try + { + // Create a TestCosmosAgentThread just to let it set up the container infrastructure + // This ensures the container is created using the exact same pattern the test will use + var setupThread = new TestCosmosAgentThread(this._connectionString, this._testDatabaseId, this._testContainerId); + + // Trigger container creation by accessing the MessageStore and trying to get messages + // This will create the database and container if they don't exist + try + { + await setupThread.MessageStore.GetMessagesAsync(); + } + catch (Microsoft.Azure.Cosmos.CosmosException ex) when (ex.StatusCode == System.Net.HttpStatusCode.NotFound) + { + // This is expected if the container is empty - it means the container was created successfully + } + } + catch (Exception ex) + { + throw new InvalidOperationException($"Container {this._testContainerId} in database {this._testDatabaseId} is not ready: {ex.Message}", ex); + } + } + } + + #region Constructor Tests + + [Fact] + [Trait("Category", "CosmosDB")] + public async Task Constructor_WithConnectionString_CreatesValidThread() + { + // Arrange & Act - Skip if emulator not available + this.SkipIfEmulatorNotAvailable(); + + // Act + var thread = new TestCosmosAgentThread(this._connectionString, this._testDatabaseId, this._testContainerId); + + // Assert + Assert.NotNull(thread); + Assert.Equal(this._testDatabaseId, thread.MessageStore.DatabaseId); + Assert.Equal(this._testContainerId, thread.MessageStore.ContainerId); + Assert.NotNull(thread.ConversationId); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public async Task Constructor_WithConnectionStringAndConversationId_CreatesValidThread() + { + // Arrange & Act - Skip if emulator not available + this.SkipIfEmulatorNotAvailable(); + + const string ExpectedConversationId = "test-conversation-123"; + + // Act + var thread = new TestCosmosAgentThread(this._connectionString, this._testDatabaseId, this._testContainerId, ExpectedConversationId); + + // Assert + Assert.NotNull(thread); + Assert.Equal(ExpectedConversationId, thread.ConversationId); + Assert.Equal(this._testDatabaseId, thread.MessageStore.DatabaseId); + Assert.Equal(this._testContainerId, thread.MessageStore.ContainerId); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public async Task Constructor_WithManagedIdentity_CreatesValidThread() + { + // Arrange & Act - Skip if emulator not available + this.SkipIfEmulatorNotAvailable(); + + const string TestConversationId = "test-conversation-456"; + + // Act + var thread = new TestCosmosAgentThread(EmulatorEndpoint, this._testDatabaseId, this._testContainerId, TestConversationId, useManagedIdentity: true); + + // Assert + Assert.NotNull(thread); + Assert.Equal(TestConversationId, thread.ConversationId); + Assert.Equal(this._testDatabaseId, thread.MessageStore.DatabaseId); + Assert.Equal(this._testContainerId, thread.MessageStore.ContainerId); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public async Task Constructor_WithExistingMessageStore_CreatesValidThread() + { + // Arrange & Act - Skip if emulator not available + this.SkipIfEmulatorNotAvailable(); + + // Arrange + var messageStore = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, this._testContainerId); + + // Act + var thread = new TestCosmosAgentThread(messageStore); + + // Assert + Assert.NotNull(thread); + Assert.Equal(messageStore, thread.MessageStore); + Assert.NotNull(thread.ConversationId); + } + + [Fact(Skip = "Serialization requires additional JSON configuration for internal state types")] + [Trait("Category", "CosmosDB")] + public async Task Constructor_WithSerialization_RestoresCorrectly() + { + // Arrange & Act - Skip if emulator not available + this.SkipIfEmulatorNotAvailable(); + + // Arrange + var originalThread = new TestCosmosAgentThread(this._connectionString, this._testDatabaseId, this._testContainerId, "serialization-test"); + + // Add some messages to the original thread + await originalThread.MessageStore.AddMessagesAsync([ + new ChatMessage(ChatRole.User, "Test message for serialization"), + new ChatMessage(ChatRole.Assistant, "Response message") + ]); + + // Serialize the thread + var serialized = originalThread.Serialize(); + + // Factory function to recreate message store from serialized state + CosmosChatMessageStore MessageStoreFactory(JsonElement storeState, JsonSerializerOptions? options) + { + return new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, this._testContainerId); + } + + // Act - Restore from serialization + var restoredThread = new TestCosmosAgentThread(serialized, MessageStoreFactory); + + // Assert + Assert.NotNull(restoredThread); + Assert.Equal(originalThread.ConversationId, restoredThread.ConversationId); + + // Verify messages were preserved + var originalMessages = (await originalThread.MessageStore.GetMessagesAsync()).ToList(); + var restoredMessages = (await restoredThread.MessageStore.GetMessagesAsync()).ToList(); + + Assert.Equal(originalMessages.Count, restoredMessages.Count); + Assert.Equal(originalMessages.First().Contents, restoredMessages.First().Contents); + } + + #endregion + + #region Property and Delegation Tests + + [Fact] + [Trait("Category", "CosmosDB")] + public async Task MessageStore_DelegatesToUnderlyingStore() + { + // Arrange & Act - Skip if emulator not available + this.SkipIfEmulatorNotAvailable(); + + // Act + var thread = new TestCosmosAgentThread(this._connectionString, this._testDatabaseId, this._testContainerId); + + // Assert + Assert.NotNull(thread.MessageStore); + Assert.IsType(thread.MessageStore); + } + + #endregion + + #region Message Operations Tests + + [Fact] + [Trait("Category", "CosmosDB")] + public async Task AddAndGetMessagesAsync_WorksCorrectly() + { + // Arrange & Act - Skip if emulator not available + this.SkipIfEmulatorNotAvailable(); + + // Arrange - Use unique conversation ID to prevent test interference + var ConversationId = $"add-get-test-{Guid.NewGuid():N}"; + var thread = new TestCosmosAgentThread(this._connectionString, this._testDatabaseId, this._testContainerId, ConversationId); + + // Act + await thread.MessageStore.AddMessagesAsync([ + new ChatMessage(ChatRole.User, "Hello, this is a test message"), + new ChatMessage(ChatRole.Assistant, "Hello! I'm an AI assistant. How can I help you?") + ]); + + var messages = (await thread.MessageStore.GetMessagesAsync()).ToList(); + + // Assert + Assert.Equal(2, messages.Count); + Assert.Equal(ChatRole.User, messages[0].Role); + Assert.Equal("Hello, this is a test message", messages[0].Contents.FirstOrDefault()?.ToString()); + Assert.Equal(ChatRole.Assistant, messages[1].Role); + Assert.Equal("Hello! I'm an AI assistant. How can I help you?", messages[1].Contents.FirstOrDefault()?.ToString()); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public async Task DirectCosmosChatMessageStore_WithWorkingContainer_WorksCorrectlyAsync() + { + // Arrange & Act - Skip if emulator not available + this.SkipIfEmulatorNotAvailable(); + + // Arrange - Use the working container name from CosmosChatMessageStoreTests + var ConversationId = $"direct-test-{Guid.NewGuid():N}"; + using var messageStore = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, "ChatMessages", ConversationId); + + // Act + await messageStore.AddMessagesAsync([ + new ChatMessage(ChatRole.User, "Hello, this is a direct test message"), + new ChatMessage(ChatRole.Assistant, "Hello! I'm responding directly.") + ]); + + var messages = (await messageStore.GetMessagesAsync()).ToList(); + + // Assert + Assert.Equal(2, messages.Count); + Assert.Equal(ChatRole.User, messages[0].Role); + Assert.Equal("Hello, this is a direct test message", messages[0].Contents.FirstOrDefault()?.ToString()); + Assert.Equal(ChatRole.Assistant, messages[1].Role); + Assert.Equal("Hello! I'm responding directly.", messages[1].Contents.FirstOrDefault()?.ToString()); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public async Task GetMessageCountAsync_ReturnsCorrectCount() + { + // Arrange & Act - Skip if emulator not available + this.SkipIfEmulatorNotAvailable(); + await this.EnsureContainerReadyAsync(); + + // Arrange - Use unique conversation ID to prevent test interference + var ConversationId = $"count-test-{Guid.NewGuid():N}"; + var thread = new TestCosmosAgentThread(this._connectionString, this._testDatabaseId, this._testContainerId, ConversationId); + + // Act & Assert - Start with 0 messages + var initialCount = await thread.GetMessageCountAsync(); + Assert.Equal(0, initialCount); + + // Add some messages + await thread.MessageStore.AddMessagesAsync([ + new ChatMessage(ChatRole.User, "Message 1"), + new ChatMessage(ChatRole.User, "Message 2"), + new ChatMessage(ChatRole.Assistant, "Response") + ]); + + var finalCount = await thread.GetMessageCountAsync(); + Assert.Equal(3, finalCount); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public async Task ClearMessagesAsync_RemovesAllMessages() + { + // Arrange & Act - Skip if emulator not available + this.SkipIfEmulatorNotAvailable(); + await this.EnsureContainerReadyAsync(); + + // Arrange - Use unique conversation ID to prevent test interference + var ConversationId = $"clear-test-{Guid.NewGuid():N}"; + var thread = new TestCosmosAgentThread(this._connectionString, this._testDatabaseId, this._testContainerId, ConversationId); + + // Add messages first + await thread.MessageStore.AddMessagesAsync([ + new ChatMessage(ChatRole.User, "Message to be cleared"), + new ChatMessage(ChatRole.Assistant, "Response to be cleared") + ]); + + // Verify messages exist + var countBeforeClear = await thread.GetMessageCountAsync(); + Assert.Equal(2, countBeforeClear); + + // Act + await thread.ClearMessagesAsync(); + + // Assert + var countAfterClear = await thread.GetMessageCountAsync(); + Assert.Equal(0, countAfterClear); + + var messagesAfterClear = (await thread.MessageStore.GetMessagesAsync()).ToList(); + Assert.Empty(messagesAfterClear); + } + + #endregion + + #region Conversation Isolation Tests + + [Fact] + [Trait("Category", "CosmosDB")] + public async Task DifferentConversationsAsync_AreIsolated() + { + // Arrange & Act - Skip if emulator not available + this.SkipIfEmulatorNotAvailable(); + await this.EnsureContainerReadyAsync(); + + // Arrange - Use unique conversation IDs to avoid interference from other tests + var timestamp = Guid.NewGuid().ToString("N").Substring(0, 8); + var thread1 = new TestCosmosAgentThread(this._connectionString, this._testDatabaseId, this._testContainerId, $"isolation-conv-1-{timestamp}"); + var thread2 = new TestCosmosAgentThread(this._connectionString, this._testDatabaseId, this._testContainerId, $"isolation-conv-2-{timestamp}"); + + // Act + await thread1.MessageStore.AddMessagesAsync([new ChatMessage(ChatRole.User, "Message in conversation 1")]); + await thread2.MessageStore.AddMessagesAsync([new ChatMessage(ChatRole.User, "Message in conversation 2")]); + + var messages1 = (await thread1.MessageStore.GetMessagesAsync()).ToList(); + var messages2 = (await thread2.MessageStore.GetMessagesAsync()).ToList(); + + // Assert + Assert.Single(messages1); + Assert.Single(messages2); + Assert.Equal("Message in conversation 1", messages1[0].Contents.FirstOrDefault()?.ToString()); + Assert.Equal("Message in conversation 2", messages2[0].Contents.FirstOrDefault()?.ToString()); + + // Verify that clearing one conversation doesn't affect the other + await thread1.ClearMessagesAsync(); + + var messages1AfterClear = (await thread1.MessageStore.GetMessagesAsync()).ToList(); + var messages2AfterClear = (await thread2.MessageStore.GetMessagesAsync()).ToList(); + + Assert.Empty(messages1AfterClear); + Assert.Single(messages2AfterClear); + } + + #endregion + + #region Serialization Tests + + [Fact(Skip = "Serialization requires additional JSON configuration for internal state types")] + [Trait("Category", "CosmosDB")] + public async Task SerializeAsync_PreservesStateInformation() + { + // Arrange & Act - Skip if emulator not available + this.SkipIfEmulatorNotAvailable(); + + // Arrange + var thread = new TestCosmosAgentThread(this._connectionString, this._testDatabaseId, this._testContainerId, "serialization-conversation"); + + // Add some test data + await thread.MessageStore.AddMessagesAsync([new ChatMessage(ChatRole.User, "Serialization test message")]); + + // Act + var serialized = thread.Serialize(); + + // Assert + Assert.Equal(JsonValueKind.Object, serialized.ValueKind); + Assert.True(serialized.TryGetProperty("storeState", out var storeStateProperty)); + + // The store state should contain conversation information for restoration + Assert.Equal(JsonValueKind.Object, storeStateProperty.ValueKind); + } + + #endregion +} + +// Test implementation class to access protected constructors +file sealed class TestCosmosAgentThread : CosmosAgentThread +{ + public TestCosmosAgentThread(string connectionString, string databaseId, string containerId) + : base(connectionString, databaseId, containerId) { } + + public TestCosmosAgentThread(string connectionString, string databaseId, string containerId, string conversationId) + : base(connectionString, databaseId, containerId, conversationId) { } + + public TestCosmosAgentThread(string accountEndpoint, string databaseId, string containerId, string? conversationId, bool useManagedIdentity) + : base(accountEndpoint, databaseId, containerId, conversationId, useManagedIdentity) { } + + public TestCosmosAgentThread(CosmosChatMessageStore messageStore) + : base(messageStore) { } + + public TestCosmosAgentThread(JsonElement serializedThreadState, Func messageStoreFactory, JsonSerializerOptions? jsonSerializerOptions = null) + : base(serializedThreadState, messageStoreFactory, jsonSerializerOptions) { } +} diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/CosmosChatMessageStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/CosmosChatMessageStoreTests.cs index 388c753806..bdb9954b56 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/CosmosChatMessageStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/CosmosChatMessageStoreTests.cs @@ -45,11 +45,10 @@ public sealed class CosmosChatMessageStoreTests : IAsyncLifetime, IDisposable private const string TestDatabaseId = "AgentFrameworkTests"; private const string TestContainerId = "ChatMessages"; - private CosmosClient? _cosmosClient; - private Database? _database; - private Container? _container; + private string _connectionString = string.Empty; private bool _emulatorAvailable; private bool _preserveContainer; + private CosmosClient? _setupClient; // Only used for test setup/cleanup public async Task InitializeAsync() { @@ -57,13 +56,16 @@ public async Task InitializeAsync() // Set COSMOS_PRESERVE_CONTAINERS=true to keep containers and data for inspection _preserveContainer = string.Equals(Environment.GetEnvironmentVariable("COSMOS_PRESERVE_CONTAINERS"), "true", StringComparison.OrdinalIgnoreCase); + _connectionString = $"AccountEndpoint={EmulatorEndpoint};AccountKey={EmulatorKey}"; + try { - _cosmosClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); + // Only create CosmosClient for test setup - the actual tests will use connection string constructors + _setupClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); // Test connection by attempting to create database - _database = await _cosmosClient.CreateDatabaseIfNotExistsAsync(TestDatabaseId); - _container = await _database.CreateContainerIfNotExistsAsync( + var databaseResponse = await this._setupClient.CreateDatabaseIfNotExistsAsync(TestDatabaseId); + await databaseResponse.Database.CreateContainerIfNotExistsAsync( TestContainerId, "/conversationId", throughput: 400); @@ -74,14 +76,14 @@ public async Task InitializeAsync() { // Emulator not available, tests will be skipped _emulatorAvailable = false; - _cosmosClient?.Dispose(); - _cosmosClient = null; + _setupClient?.Dispose(); + _setupClient = null; } } public async Task DisposeAsync() { - if (_cosmosClient != null && _emulatorAvailable) + if (_setupClient != null && _emulatorAvailable) { try { @@ -94,7 +96,8 @@ public async Task DisposeAsync() else { // Clean mode: Delete the test database and all data - await _database!.DeleteAsync(); + var database = _setupClient.GetDatabase(TestDatabaseId); + await database.DeleteAsync(); } } catch @@ -103,14 +106,14 @@ public async Task DisposeAsync() } finally { - _cosmosClient.Dispose(); + _setupClient.Dispose(); } } } public void Dispose() { - _cosmosClient?.Dispose(); + _setupClient?.Dispose(); GC.SuppressFinalize(this); } @@ -126,40 +129,45 @@ private void SkipIfEmulatorNotAvailable() [Fact] [Trait("Category", "CosmosDB")] - public void Constructor_WithCosmosClient_ShouldCreateInstance() + public void Constructor_WithConnectionString_ShouldCreateInstance() { // Arrange & Act - SkipIfEmulatorNotAvailable(); + this.SkipIfEmulatorNotAvailable(); // Act - using var store = new CosmosChatMessageStore(_cosmosClient!, TestDatabaseId, TestContainerId, "test-conversation"); + using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, "test-conversation"); // Assert Assert.NotNull(store); + Assert.Equal("test-conversation", store.ConversationId); + Assert.Equal(TestDatabaseId, store.DatabaseId); + Assert.Equal(TestContainerId, store.ContainerId); } [Fact] [Trait("Category", "CosmosDB")] - public void Constructor_WithConnectionString_ShouldCreateInstance() + public void Constructor_WithConnectionStringNoConversationId_ShouldCreateInstance() { // Arrange - SkipIfEmulatorNotAvailable(); - var connectionString = $"AccountEndpoint={EmulatorEndpoint};AccountKey={EmulatorKey};"; + this.SkipIfEmulatorNotAvailable(); // Act - using var store = new CosmosChatMessageStore(connectionString, TestDatabaseId, TestContainerId); + using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId); // Assert Assert.NotNull(store); + Assert.NotNull(store.ConversationId); + Assert.Equal(TestDatabaseId, store.DatabaseId); + Assert.Equal(TestContainerId, store.ContainerId); } [Fact] [Trait("Category", "CosmosDB")] - public void Constructor_WithNullCosmosClient_ShouldThrowArgumentNullException() + public void Constructor_WithNullConnectionString_ShouldThrowArgumentException() { // Arrange & Act & Assert - Assert.Throws(() => - new CosmosChatMessageStore((CosmosClient)null!, TestDatabaseId, TestContainerId, "test-conversation")); + Assert.Throws(() => + new CosmosChatMessageStore((string)null!, TestDatabaseId, TestContainerId, "test-conversation")); } [Fact] @@ -170,7 +178,7 @@ public void Constructor_WithEmptyConversationId_ShouldThrowArgumentException() SkipIfEmulatorNotAvailable(); Assert.Throws(() => - new CosmosChatMessageStore(_cosmosClient!, TestDatabaseId, TestContainerId, "")); + new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, "")); } #endregion @@ -182,9 +190,9 @@ public void Constructor_WithEmptyConversationId_ShouldThrowArgumentException() public async Task AddMessagesAsync_WithSingleMessage_ShouldAddMessageAsync() { // Arrange - SkipIfEmulatorNotAvailable(); + this.SkipIfEmulatorNotAvailable(); var conversationId = Guid.NewGuid().ToString(); - using var store = new CosmosChatMessageStore(_cosmosClient!, TestDatabaseId, TestContainerId, conversationId); + using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, conversationId); var message = new ChatMessage(ChatRole.User, "Hello, world!"); // Act @@ -203,7 +211,7 @@ public async Task AddMessagesAsync_WithSingleMessage_ShouldAddMessageAsync() // Let's check if we can find ANY items in the container for this conversation var directQuery = new QueryDefinition("SELECT VALUE COUNT(1) FROM c WHERE c.conversationId = @conversationId") .WithParameter("@conversationId", conversationId); - var countIterator = this._cosmosClient!.GetDatabase(TestDatabaseId).GetContainer(TestContainerId) + var countIterator = this._setupClient!.GetDatabase(TestDatabaseId).GetContainer(TestContainerId) .GetItemQueryIterator(directQuery, requestOptions: new QueryRequestOptions { PartitionKey = new PartitionKey(conversationId) @@ -215,7 +223,7 @@ public async Task AddMessagesAsync_WithSingleMessage_ShouldAddMessageAsync() // Debug: Let's see what the raw query returns var rawQuery = new QueryDefinition("SELECT * FROM c WHERE c.conversationId = @conversationId") .WithParameter("@conversationId", conversationId); - var rawIterator = this._cosmosClient!.GetDatabase(TestDatabaseId).GetContainer(TestContainerId) + var rawIterator = this._setupClient!.GetDatabase(TestDatabaseId).GetContainer(TestContainerId) .GetItemQueryIterator(rawQuery, requestOptions: new QueryRequestOptions { PartitionKey = new PartitionKey(conversationId) @@ -244,7 +252,7 @@ public async Task AddMessagesAsync_WithMultipleMessages_ShouldAddAllMessagesAsyn // Arrange SkipIfEmulatorNotAvailable(); var conversationId = Guid.NewGuid().ToString(); - using var store = new CosmosChatMessageStore(_cosmosClient!, TestDatabaseId, TestContainerId, conversationId); + using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, conversationId); var messages = new[] { new ChatMessage(ChatRole.User, "First message"), @@ -274,7 +282,7 @@ public async Task GetMessagesAsync_WithNoMessages_ShouldReturnEmptyAsync() { // Arrange SkipIfEmulatorNotAvailable(); - using var store = new CosmosChatMessageStore(_cosmosClient!, TestDatabaseId, TestContainerId, Guid.NewGuid().ToString()); + using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, Guid.NewGuid().ToString()); // Act var messages = await store.GetMessagesAsync(); @@ -292,8 +300,8 @@ public async Task GetMessagesAsync_WithConversationIsolation_ShouldOnlyReturnMes var conversation1 = Guid.NewGuid().ToString(); var conversation2 = Guid.NewGuid().ToString(); - using var store1 = new CosmosChatMessageStore(_cosmosClient!, TestDatabaseId, TestContainerId, conversation1); - using var store2 = new CosmosChatMessageStore(_cosmosClient!, TestDatabaseId, TestContainerId, conversation2); + using var store1 = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, conversation1); + using var store2 = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, conversation2); await store1.AddMessagesAsync([new ChatMessage(ChatRole.User, "Message for conversation 1")]); await store2.AddMessagesAsync([new ChatMessage(ChatRole.User, "Message for conversation 2")]); @@ -321,7 +329,7 @@ public async Task FullWorkflow_AddAndGet_ShouldWorkCorrectlyAsync() { // Arrange SkipIfEmulatorNotAvailable(); - using var originalStore = new CosmosChatMessageStore(_cosmosClient!, TestDatabaseId, TestContainerId, "test-conversation"); + using var originalStore = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, "test-conversation"); var messages = new[] { @@ -341,7 +349,7 @@ public async Task FullWorkflow_AddAndGet_ShouldWorkCorrectlyAsync() Assert.Equal(5, retrievedList.Count); // Act 3: Create new store instance for same conversation (test persistence) - using var newStore = new CosmosChatMessageStore(_cosmosClient!, TestDatabaseId, TestContainerId, "test-conversation"); + using var newStore = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, "test-conversation"); var persistedMessages = await newStore.GetMessagesAsync(); var persistedList = persistedMessages.ToList(); @@ -364,7 +372,7 @@ public void Dispose_AfterUse_ShouldNotThrow() { // Arrange SkipIfEmulatorNotAvailable(); - var store = new CosmosChatMessageStore(_cosmosClient!, TestDatabaseId, TestContainerId, Guid.NewGuid().ToString()); + var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, Guid.NewGuid().ToString()); // Act & Assert store.Dispose(); // Should not throw @@ -376,7 +384,7 @@ public void Dispose_MultipleCalls_ShouldNotThrow() { // Arrange SkipIfEmulatorNotAvailable(); - var store = new CosmosChatMessageStore(_cosmosClient!, TestDatabaseId, TestContainerId, Guid.NewGuid().ToString()); + var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, Guid.NewGuid().ToString()); // Act & Assert store.Dispose(); // First call @@ -384,4 +392,4 @@ public void Dispose_MultipleCalls_ShouldNotThrow() } #endregion -} \ No newline at end of file +} From 67c65aabb34727a6ab2831bd23b63c718615b4f3 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Fri, 31 Oct 2025 12:09:33 +0000 Subject: [PATCH 03/40] revert unnecessary changes and fix tests --- .../Extensions/AgentProviderExtensions.cs | 18 ++--- ...oft.Agents.AI.Workflows.Declarative.csproj | 1 - .../CosmosAgentThreadTests.cs | 2 +- .../CosmosChatMessageStoreTests.cs | 74 ++++++++++--------- .../CosmosCheckpointStoreTests.cs | 52 ++++++------- 5 files changed, 73 insertions(+), 74 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Extensions/AgentProviderExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Extensions/AgentProviderExtensions.cs index 6c09043bb4..5b6bbbc297 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Extensions/AgentProviderExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Extensions/AgentProviderExtensions.cs @@ -4,24 +4,20 @@ using System.Linq; using System.Threading; using System.Threading.Tasks; -#if NET9_0_OR_GREATER using Azure.AI.Agents.Persistent; -#endif using Microsoft.Extensions.AI; namespace Microsoft.Agents.AI.Workflows.Declarative.Extensions; internal static class AgentProviderExtensions { -#if NET9_0_OR_GREATER - private static readonly HashSet s_failureStatus = + private static readonly HashSet s_failureStatus = [ - global::Azure.AI.Agents.Persistent.RunStatus.Failed, - global::Azure.AI.Agents.Persistent.RunStatus.Cancelled, - global::Azure.AI.Agents.Persistent.RunStatus.Cancelling, - global::Azure.AI.Agents.Persistent.RunStatus.Expired, + Azure.AI.Agents.Persistent.RunStatus.Failed, + Azure.AI.Agents.Persistent.RunStatus.Cancelled, + Azure.AI.Agents.Persistent.RunStatus.Cancelling, + Azure.AI.Agents.Persistent.RunStatus.Expired, ]; -#endif public static async ValueTask InvokeAgentAsync( this WorkflowAgentProvider agentProvider, @@ -64,14 +60,12 @@ inputMessages is not null ? updates.Add(update); -#if NET9_0_OR_GREATER if (update.RawRepresentation is ChatResponseUpdate chatUpdate && - chatUpdate.RawRepresentation is global::Azure.AI.Agents.Persistent.RunUpdate runUpdate && + chatUpdate.RawRepresentation is RunUpdate runUpdate && s_failureStatus.Contains(runUpdate.Value.Status)) { throw new DeclarativeActionException($"Unexpected failure invoking agent, run {runUpdate.Value.Status}: {agent.Name ?? agent.Id} [{runUpdate.Value.Id}/{conversationId}]"); } -#endif if (autoSend) { diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Microsoft.Agents.AI.Workflows.Declarative.csproj b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Microsoft.Agents.AI.Workflows.Declarative.csproj index 7b268f6c42..b3b2a86ab4 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Microsoft.Agents.AI.Workflows.Declarative.csproj +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Microsoft.Agents.AI.Workflows.Declarative.csproj @@ -22,7 +22,6 @@ - diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/CosmosAgentThreadTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/CosmosAgentThreadTests.cs index eecb3424b0..fc8dd2560d 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/CosmosAgentThreadTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/CosmosAgentThreadTests.cs @@ -43,7 +43,7 @@ public sealed class CosmosAgentThreadTests : IAsyncLifetime, IDisposable private const string EmulatorKey = "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw=="; // Use static container names like CosmosChatMessageStoreTests for consistency - private readonly string _testDatabaseId = "AgentFrameworkTests"; + private readonly string _testDatabaseId = $"AgentFrameworkTests-Thread-{Guid.NewGuid():N}"; private readonly string _testContainerId = "ChatMessages"; // Use the same container as CosmosChatMessageStoreTests private string _connectionString = string.Empty; diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/CosmosChatMessageStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/CosmosChatMessageStoreTests.cs index bdb9954b56..66266470a3 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/CosmosChatMessageStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/CosmosChatMessageStoreTests.cs @@ -42,9 +42,11 @@ public sealed class CosmosChatMessageStoreTests : IAsyncLifetime, IDisposable // Cosmos DB Emulator connection settings private const string EmulatorEndpoint = "https://localhost:8081"; private const string EmulatorKey = "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw=="; - private const string TestDatabaseId = "AgentFrameworkTests"; private const string TestContainerId = "ChatMessages"; + // Use unique database ID per test run to avoid conflicts between parallel test executions + private readonly string _testDatabaseId = $"AgentFrameworkTests-ChatStore-{Guid.NewGuid():N}"; + private string _connectionString = string.Empty; private bool _emulatorAvailable; private bool _preserveContainer; @@ -54,40 +56,40 @@ public async Task InitializeAsync() { // Check environment variable to determine if we should preserve containers // Set COSMOS_PRESERVE_CONTAINERS=true to keep containers and data for inspection - _preserveContainer = string.Equals(Environment.GetEnvironmentVariable("COSMOS_PRESERVE_CONTAINERS"), "true", StringComparison.OrdinalIgnoreCase); + this._preserveContainer = string.Equals(Environment.GetEnvironmentVariable("COSMOS_PRESERVE_CONTAINERS"), "true", StringComparison.OrdinalIgnoreCase); - _connectionString = $"AccountEndpoint={EmulatorEndpoint};AccountKey={EmulatorKey}"; + this._connectionString = $"AccountEndpoint={EmulatorEndpoint};AccountKey={EmulatorKey}"; try { // Only create CosmosClient for test setup - the actual tests will use connection string constructors - _setupClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); + this._setupClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); // Test connection by attempting to create database - var databaseResponse = await this._setupClient.CreateDatabaseIfNotExistsAsync(TestDatabaseId); + var databaseResponse = await this._setupClient.CreateDatabaseIfNotExistsAsync(this._testDatabaseId); await databaseResponse.Database.CreateContainerIfNotExistsAsync( TestContainerId, "/conversationId", throughput: 400); - _emulatorAvailable = true; + this._emulatorAvailable = true; } catch (Exception) { // Emulator not available, tests will be skipped - _emulatorAvailable = false; - _setupClient?.Dispose(); - _setupClient = null; + this._emulatorAvailable = false; + this._setupClient?.Dispose(); + this._setupClient = null; } } public async Task DisposeAsync() { - if (_setupClient != null && _emulatorAvailable) + if (this._setupClient != null && this._emulatorAvailable) { try { - if (_preserveContainer) + if (this._preserveContainer) { // Preserve mode: Don't delete the database/container, keep data for inspection // This allows viewing data in the Cosmos DB Emulator Data Explorer @@ -96,30 +98,31 @@ public async Task DisposeAsync() else { // Clean mode: Delete the test database and all data - var database = _setupClient.GetDatabase(TestDatabaseId); + var database = this._setupClient.GetDatabase(this._testDatabaseId); await database.DeleteAsync(); } } - catch + catch (Exception ex) { - // Ignore cleanup errors + // Ignore cleanup errors during test teardown + Console.WriteLine($"Warning: Cleanup failed: {ex.Message}"); } finally { - _setupClient.Dispose(); + this._setupClient.Dispose(); } } } public void Dispose() { - _setupClient?.Dispose(); + this._setupClient?.Dispose(); GC.SuppressFinalize(this); } private void SkipIfEmulatorNotAvailable() { - if (!_emulatorAvailable) + if (!this._emulatorAvailable) { Assert.Fail("Cosmos DB Emulator is not available. Start the emulator to run these tests."); } @@ -135,12 +138,12 @@ public void Constructor_WithConnectionString_ShouldCreateInstance() this.SkipIfEmulatorNotAvailable(); // Act - using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, "test-conversation"); + using var store = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, TestContainerId, "test-conversation"); // Assert Assert.NotNull(store); Assert.Equal("test-conversation", store.ConversationId); - Assert.Equal(TestDatabaseId, store.DatabaseId); + Assert.Equal(this._testDatabaseId, store.DatabaseId); Assert.Equal(TestContainerId, store.ContainerId); } @@ -152,12 +155,12 @@ public void Constructor_WithConnectionStringNoConversationId_ShouldCreateInstanc this.SkipIfEmulatorNotAvailable(); // Act - using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId); + using var store = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, TestContainerId); // Assert Assert.NotNull(store); Assert.NotNull(store.ConversationId); - Assert.Equal(TestDatabaseId, store.DatabaseId); + Assert.Equal(this._testDatabaseId, store.DatabaseId); Assert.Equal(TestContainerId, store.ContainerId); } @@ -167,7 +170,7 @@ public void Constructor_WithNullConnectionString_ShouldThrowArgumentException() { // Arrange & Act & Assert Assert.Throws(() => - new CosmosChatMessageStore((string)null!, TestDatabaseId, TestContainerId, "test-conversation")); + new CosmosChatMessageStore((string)null!, this._testDatabaseId, TestContainerId, "test-conversation")); } [Fact] @@ -175,10 +178,10 @@ public void Constructor_WithNullConnectionString_ShouldThrowArgumentException() public void Constructor_WithEmptyConversationId_ShouldThrowArgumentException() { // Arrange & Act & Assert - SkipIfEmulatorNotAvailable(); + this.SkipIfEmulatorNotAvailable(); Assert.Throws(() => - new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, "")); + new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, TestContainerId, "")); } #endregion @@ -192,7 +195,7 @@ public async Task AddMessagesAsync_WithSingleMessage_ShouldAddMessageAsync() // Arrange this.SkipIfEmulatorNotAvailable(); var conversationId = Guid.NewGuid().ToString(); - using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, conversationId); + using var store = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, TestContainerId, conversationId); var message = new ChatMessage(ChatRole.User, "Hello, world!"); // Act @@ -211,7 +214,7 @@ public async Task AddMessagesAsync_WithSingleMessage_ShouldAddMessageAsync() // Let's check if we can find ANY items in the container for this conversation var directQuery = new QueryDefinition("SELECT VALUE COUNT(1) FROM c WHERE c.conversationId = @conversationId") .WithParameter("@conversationId", conversationId); - var countIterator = this._setupClient!.GetDatabase(TestDatabaseId).GetContainer(TestContainerId) + var countIterator = this._setupClient!.GetDatabase(this._testDatabaseId).GetContainer(TestContainerId) .GetItemQueryIterator(directQuery, requestOptions: new QueryRequestOptions { PartitionKey = new PartitionKey(conversationId) @@ -223,7 +226,7 @@ public async Task AddMessagesAsync_WithSingleMessage_ShouldAddMessageAsync() // Debug: Let's see what the raw query returns var rawQuery = new QueryDefinition("SELECT * FROM c WHERE c.conversationId = @conversationId") .WithParameter("@conversationId", conversationId); - var rawIterator = this._setupClient!.GetDatabase(TestDatabaseId).GetContainer(TestContainerId) + var rawIterator = this._setupClient!.GetDatabase(this._testDatabaseId).GetContainer(TestContainerId) .GetItemQueryIterator(rawQuery, requestOptions: new QueryRequestOptions { PartitionKey = new PartitionKey(conversationId) @@ -252,7 +255,7 @@ public async Task AddMessagesAsync_WithMultipleMessages_ShouldAddAllMessagesAsyn // Arrange SkipIfEmulatorNotAvailable(); var conversationId = Guid.NewGuid().ToString(); - using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, conversationId); + using var store = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, TestContainerId, conversationId); var messages = new[] { new ChatMessage(ChatRole.User, "First message"), @@ -282,7 +285,7 @@ public async Task GetMessagesAsync_WithNoMessages_ShouldReturnEmptyAsync() { // Arrange SkipIfEmulatorNotAvailable(); - using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, Guid.NewGuid().ToString()); + using var store = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, TestContainerId, Guid.NewGuid().ToString()); // Act var messages = await store.GetMessagesAsync(); @@ -300,8 +303,8 @@ public async Task GetMessagesAsync_WithConversationIsolation_ShouldOnlyReturnMes var conversation1 = Guid.NewGuid().ToString(); var conversation2 = Guid.NewGuid().ToString(); - using var store1 = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, conversation1); - using var store2 = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, conversation2); + using var store1 = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, TestContainerId, conversation1); + using var store2 = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, TestContainerId, conversation2); await store1.AddMessagesAsync([new ChatMessage(ChatRole.User, "Message for conversation 1")]); await store2.AddMessagesAsync([new ChatMessage(ChatRole.User, "Message for conversation 2")]); @@ -329,7 +332,8 @@ public async Task FullWorkflow_AddAndGet_ShouldWorkCorrectlyAsync() { // Arrange SkipIfEmulatorNotAvailable(); - using var originalStore = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, "test-conversation"); + var conversationId = $"test-conversation-{Guid.NewGuid():N}"; // Use unique conversation ID + using var originalStore = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, TestContainerId, conversationId); var messages = new[] { @@ -349,7 +353,7 @@ public async Task FullWorkflow_AddAndGet_ShouldWorkCorrectlyAsync() Assert.Equal(5, retrievedList.Count); // Act 3: Create new store instance for same conversation (test persistence) - using var newStore = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, "test-conversation"); + using var newStore = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, TestContainerId, conversationId); var persistedMessages = await newStore.GetMessagesAsync(); var persistedList = persistedMessages.ToList(); @@ -372,7 +376,7 @@ public void Dispose_AfterUse_ShouldNotThrow() { // Arrange SkipIfEmulatorNotAvailable(); - var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, Guid.NewGuid().ToString()); + var store = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, TestContainerId, Guid.NewGuid().ToString()); // Act & Assert store.Dispose(); // Should not throw @@ -384,7 +388,7 @@ public void Dispose_MultipleCalls_ShouldNotThrow() { // Arrange SkipIfEmulatorNotAvailable(); - var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, Guid.NewGuid().ToString()); + var store = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, TestContainerId, Guid.NewGuid().ToString()); // Act & Assert store.Dispose(); // First call diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/CosmosCheckpointStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/CosmosCheckpointStoreTests.cs index 483988ba3b..9f323ff60e 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/CosmosCheckpointStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/CosmosCheckpointStoreTests.cs @@ -31,9 +31,12 @@ public class CosmosCheckpointStoreTests : IAsyncLifetime, IDisposable // Cosmos DB Emulator connection settings private const string EmulatorEndpoint = "https://localhost:8081"; private const string EmulatorKey = "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw=="; - private const string TestDatabaseId = "AgentFrameworkTests"; private const string TestContainerId = "Checkpoints"; + // Use unique database ID per test run to avoid conflicts between parallel test executions + private readonly string _testDatabaseId = $"AgentFrameworkTests-CheckpointStore-{Guid.NewGuid():N}"; + + private string _connectionString = string.Empty; private CosmosClient? _cosmosClient; private Database? _database; private Container? _container; @@ -46,12 +49,14 @@ public async Task InitializeAsync() // Set COSMOS_PRESERVE_CONTAINERS=true to keep containers and data for inspection _preserveContainer = string.Equals(Environment.GetEnvironmentVariable("COSMOS_PRESERVE_CONTAINERS"), "true", StringComparison.OrdinalIgnoreCase); + _connectionString = $"AccountEndpoint={EmulatorEndpoint};AccountKey={EmulatorKey}"; + try { _cosmosClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); // Test connection by attempting to create database - _database = await _cosmosClient.CreateDatabaseIfNotExistsAsync(TestDatabaseId); + _database = await _cosmosClient.CreateDatabaseIfNotExistsAsync(_testDatabaseId); _container = await _database.CreateContainerIfNotExistsAsync( TestContainerId, "/runId", @@ -114,10 +119,10 @@ public void Constructor_WithCosmosClient_SetsProperties() SkipIfEmulatorNotAvailable(); // Act - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(_cosmosClient!, _testDatabaseId, TestContainerId); // Assert - Assert.Equal(TestDatabaseId, store.DatabaseId); + Assert.Equal(_testDatabaseId, store.DatabaseId); Assert.Equal(TestContainerId, store.ContainerId); } @@ -126,14 +131,11 @@ public void Constructor_WithConnectionString_SetsProperties() { SkipIfEmulatorNotAvailable(); - // Arrange - var connectionString = $"AccountEndpoint={EmulatorEndpoint};AccountKey={EmulatorKey};"; - // Act - using var store = new CosmosCheckpointStore(connectionString, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(_connectionString, _testDatabaseId, TestContainerId); // Assert - Assert.Equal(TestDatabaseId, store.DatabaseId); + Assert.Equal(_testDatabaseId, store.DatabaseId); Assert.Equal(TestContainerId, store.ContainerId); } @@ -142,7 +144,7 @@ public void Constructor_WithNullCosmosClient_ThrowsArgumentNullException() { // Act & Assert Assert.Throws(() => - new CosmosCheckpointStore((CosmosClient)null!, TestDatabaseId, TestContainerId)); + new CosmosCheckpointStore((CosmosClient)null!, _testDatabaseId, TestContainerId)); } [Fact] @@ -150,7 +152,7 @@ public void Constructor_WithNullConnectionString_ThrowsArgumentException() { // Act & Assert Assert.Throws(() => - new CosmosCheckpointStore((string)null!, TestDatabaseId, TestContainerId)); + new CosmosCheckpointStore((string)null!, _testDatabaseId, TestContainerId)); } #endregion @@ -163,7 +165,7 @@ public async Task CreateCheckpointAsync_NewCheckpoint_CreatesSuccessfully() SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(_cosmosClient!, _testDatabaseId, TestContainerId); var runId = Guid.NewGuid().ToString(); var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test checkpoint" }); @@ -183,7 +185,7 @@ public async Task RetrieveCheckpointAsync_ExistingCheckpoint_ReturnsCorrectValue SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(_cosmosClient!, _testDatabaseId, TestContainerId); var runId = Guid.NewGuid().ToString(); var originalData = new { message = "Hello, World!", timestamp = DateTimeOffset.UtcNow }; var checkpointValue = JsonSerializer.SerializeToElement(originalData); @@ -204,7 +206,7 @@ public async Task RetrieveCheckpointAsync_NonExistentCheckpoint_ThrowsInvalidOpe SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(_cosmosClient!, _testDatabaseId, TestContainerId); var runId = Guid.NewGuid().ToString(); var fakeCheckpointInfo = new CheckpointInfo(runId, "nonexistent-checkpoint"); @@ -219,7 +221,7 @@ public async Task RetrieveIndexAsync_EmptyStore_ReturnsEmptyCollection() SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(_cosmosClient!, _testDatabaseId, TestContainerId); var runId = Guid.NewGuid().ToString(); // Act @@ -236,7 +238,7 @@ public async Task RetrieveIndexAsync_WithCheckpoints_ReturnsAllCheckpoints() SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(_cosmosClient!, _testDatabaseId, TestContainerId); var runId = Guid.NewGuid().ToString(); var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }); @@ -261,7 +263,7 @@ public async Task CreateCheckpointAsync_WithParent_CreatesHierarchy() SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(_cosmosClient!, _testDatabaseId, TestContainerId); var runId = Guid.NewGuid().ToString(); var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }); @@ -281,7 +283,7 @@ public async Task RetrieveIndexAsync_WithParentFilter_ReturnsFilteredResults() SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(_cosmosClient!, _testDatabaseId, TestContainerId); var runId = Guid.NewGuid().ToString(); var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }); @@ -317,7 +319,7 @@ public async Task CheckpointOperations_DifferentRuns_IsolatesData() SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(_cosmosClient!, _testDatabaseId, TestContainerId); var runId1 = Guid.NewGuid().ToString(); var runId2 = Guid.NewGuid().ToString(); var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }); @@ -347,7 +349,7 @@ public async Task CreateCheckpointAsync_WithNullRunId_ThrowsArgumentException() SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(_cosmosClient!, _testDatabaseId, TestContainerId); var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }); // Act & Assert @@ -361,7 +363,7 @@ public async Task CreateCheckpointAsync_WithEmptyRunId_ThrowsArgumentException() SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(_cosmosClient!, _testDatabaseId, TestContainerId); var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }); // Act & Assert @@ -375,7 +377,7 @@ public async Task RetrieveCheckpointAsync_WithNullCheckpointInfo_ThrowsArgumentN SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(_cosmosClient!, _testDatabaseId, TestContainerId); var runId = Guid.NewGuid().ToString(); // Act & Assert @@ -393,7 +395,7 @@ public async Task Dispose_AfterDisposal_ThrowsObjectDisposedException() SkipIfEmulatorNotAvailable(); // Arrange - var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + var store = new CosmosCheckpointStore(_cosmosClient!, _testDatabaseId, TestContainerId); var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }); // Act @@ -410,7 +412,7 @@ public void Dispose_MultipleCalls_DoesNotThrow() SkipIfEmulatorNotAvailable(); // Arrange - var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + var store = new CosmosCheckpointStore(_cosmosClient!, _testDatabaseId, TestContainerId); // Act & Assert (should not throw) store.Dispose(); @@ -433,4 +435,4 @@ protected virtual void Dispose(bool disposing) this._cosmosClient?.Dispose(); } } -} \ No newline at end of file +} From d020e188d527ad3774f3efbd37552c47318b35c7 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Fri, 31 Oct 2025 14:41:30 +0000 Subject: [PATCH 04/40] add multi-tenant support with hierarchical partition keys (and tests). --- .../CosmosChatMessageStore.cs | 273 ++++++++++++++++- .../CosmosChatMessageStoreTests.cs | 289 ++++++++++++++++++ 2 files changed, 555 insertions(+), 7 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/CosmosChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/CosmosChatMessageStore.cs index 9fa5c343cd..3709164c64 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/CosmosChatMessageStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/CosmosChatMessageStore.cs @@ -27,6 +27,12 @@ public sealed class CosmosChatMessageStore : ChatMessageStore, IDisposable private readonly bool _ownsClient; private bool _disposed; + // Hierarchical partition key support + private readonly string? _tenantId; + private readonly string? _userId; + private readonly PartitionKey _partitionKey; + private readonly bool _useHierarchicalPartitioning; + /// /// Cached JSON serializer options for .NET 9.0 compatibility. /// @@ -121,6 +127,12 @@ public CosmosChatMessageStore(string connectionString, string databaseId, string this._databaseId = databaseId; this._containerId = containerId; this._ownsClient = true; + + // Initialize simple partitioning mode + this._tenantId = null; + this._userId = null; + this._useHierarchicalPartitioning = false; + this._partitionKey = new PartitionKey(conversationId); } /// @@ -180,6 +192,12 @@ public CosmosChatMessageStore(string accountEndpoint, string databaseId, string this._databaseId = databaseId; this._containerId = containerId; this._ownsClient = true; + + // Initialize simple partitioning mode + this._tenantId = null; + this._userId = null; + this._useHierarchicalPartitioning = false; + this._partitionKey = new PartitionKey(conversationId); } /// @@ -228,6 +246,200 @@ public CosmosChatMessageStore(CosmosClient cosmosClient, string databaseId, stri this._conversationId = conversationId; this._databaseId = databaseId; this._containerId = containerId; + + // Initialize simple partitioning mode + this._tenantId = null; + this._userId = null; + this._useHierarchicalPartitioning = false; + this._partitionKey = new PartitionKey(conversationId); + } + + /// + /// Initializes a new instance of the class using a connection string with hierarchical partition keys. + /// + /// The Cosmos DB connection string. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// The tenant identifier for hierarchical partitioning. + /// The user identifier for hierarchical partitioning. + /// The session identifier for hierarchical partitioning. + /// Thrown when any required parameter is null. + /// Thrown when any string parameter is null or whitespace. + public CosmosChatMessageStore(string connectionString, string databaseId, string containerId, string tenantId, string userId, string sessionId) + { + if (string.IsNullOrWhiteSpace(connectionString)) + { + throw new ArgumentException("Connection string cannot be null or whitespace.", nameof(connectionString)); + } + + if (string.IsNullOrWhiteSpace(databaseId)) + { + throw new ArgumentException("Database ID cannot be null or whitespace.", nameof(databaseId)); + } + + if (string.IsNullOrWhiteSpace(containerId)) + { + throw new ArgumentException("Container ID cannot be null or whitespace.", nameof(containerId)); + } + + if (string.IsNullOrWhiteSpace(tenantId)) + { + throw new ArgumentException("Tenant ID cannot be null or whitespace.", nameof(tenantId)); + } + + if (string.IsNullOrWhiteSpace(userId)) + { + throw new ArgumentException("User ID cannot be null or whitespace.", nameof(userId)); + } + + if (string.IsNullOrWhiteSpace(sessionId)) + { + throw new ArgumentException("Session ID cannot be null or whitespace.", nameof(sessionId)); + } + + this._cosmosClient = new CosmosClient(connectionString); + this._container = this._cosmosClient.GetContainer(databaseId, containerId); + this._conversationId = sessionId; // Use sessionId as conversationId for compatibility + this._databaseId = databaseId; + this._containerId = containerId; + this._containerId = containerId; + + // Initialize hierarchical partitioning mode + this._tenantId = tenantId; + this._userId = userId; + this._useHierarchicalPartitioning = true; + // Use native hierarchical partition key with PartitionKeyBuilder + this._partitionKey = new PartitionKeyBuilder() + .Add(tenantId) + .Add(userId) + .Add(sessionId) + .Build(); + } + + /// + /// Initializes a new instance of the class using DefaultAzureCredential with hierarchical partition keys. + /// + /// The Cosmos DB account endpoint URI. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// The tenant identifier for hierarchical partitioning. + /// The user identifier for hierarchical partitioning. + /// The session identifier for hierarchical partitioning. + /// This parameter is used to distinguish this constructor from the connection string constructor. Always pass true. + /// Thrown when any required parameter is null. + /// Thrown when any string parameter is null or whitespace. + public CosmosChatMessageStore(string accountEndpoint, string databaseId, string containerId, string tenantId, string userId, string sessionId, bool useManagedIdentity) + { + if (string.IsNullOrWhiteSpace(accountEndpoint)) + { + throw new ArgumentException("Cannot be null or whitespace", nameof(accountEndpoint)); + } + + if (string.IsNullOrWhiteSpace(databaseId)) + { + throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); + } + + if (string.IsNullOrWhiteSpace(containerId)) + { + throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); + } + + if (string.IsNullOrWhiteSpace(tenantId)) + { + throw new ArgumentException("Tenant ID cannot be null or whitespace.", nameof(tenantId)); + } + + if (string.IsNullOrWhiteSpace(userId)) + { + throw new ArgumentException("User ID cannot be null or whitespace.", nameof(userId)); + } + + if (string.IsNullOrWhiteSpace(sessionId)) + { + throw new ArgumentException("Session ID cannot be null or whitespace.", nameof(sessionId)); + } + + if (!useManagedIdentity) + { + throw new ArgumentException("This constructor requires useManagedIdentity to be true. Use the connection string constructor for key-based authentication.", nameof(useManagedIdentity)); + } + + this._cosmosClient = new CosmosClient(accountEndpoint, new DefaultAzureCredential()); + this._container = this._cosmosClient.GetContainer(databaseId, containerId); + this._conversationId = sessionId; // Use sessionId as conversationId for compatibility + this._databaseId = databaseId; + this._containerId = containerId; + this._containerId = containerId; + + // Initialize hierarchical partitioning mode + this._tenantId = tenantId; + this._userId = userId; + this._useHierarchicalPartitioning = true; + // Use native hierarchical partition key with PartitionKeyBuilder + this._partitionKey = new PartitionKeyBuilder() + .Add(tenantId) + .Add(userId) + .Add(sessionId) + .Build(); + } + + /// + /// Initializes a new instance of the class using an existing with hierarchical partition keys. + /// + /// The instance to use for Cosmos DB operations. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// The tenant identifier for hierarchical partitioning. + /// The user identifier for hierarchical partitioning. + /// The session identifier for hierarchical partitioning. + /// Thrown when is null. + /// Thrown when any string parameter is null or whitespace. + public CosmosChatMessageStore(CosmosClient cosmosClient, string databaseId, string containerId, string tenantId, string userId, string sessionId) + { + this._cosmosClient = cosmosClient ?? throw new ArgumentNullException(nameof(cosmosClient)); + this._ownsClient = false; + + if (string.IsNullOrWhiteSpace(databaseId)) + { + throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); + } + + if (string.IsNullOrWhiteSpace(containerId)) + { + throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); + } + + if (string.IsNullOrWhiteSpace(tenantId)) + { + throw new ArgumentException("Tenant ID cannot be null or whitespace.", nameof(tenantId)); + } + + if (string.IsNullOrWhiteSpace(userId)) + { + throw new ArgumentException("User ID cannot be null or whitespace.", nameof(userId)); + } + + if (string.IsNullOrWhiteSpace(sessionId)) + { + throw new ArgumentException("Session ID cannot be null or whitespace.", nameof(sessionId)); + } + + this._container = this._cosmosClient.GetContainer(databaseId, containerId); + this._conversationId = sessionId; // Use sessionId as conversationId for compatibility + this._databaseId = databaseId; + this._containerId = containerId; + + // Initialize hierarchical partitioning mode + this._tenantId = tenantId; + this._userId = userId; + this._useHierarchicalPartitioning = true; + // Use native hierarchical partition key with PartitionKeyBuilder + this._partitionKey = new PartitionKeyBuilder() + .Add(tenantId) + .Add(userId) + .Add(sessionId) + .Build(); } /// @@ -254,6 +466,26 @@ public CosmosChatMessageStore(JsonElement serializedStoreState, CosmosClient cos this._databaseId = databaseId; this._containerId = containerId; this._container = this._cosmosClient.GetContainer(databaseId, containerId); + + // Initialize hierarchical partitioning if available in state + this._tenantId = state.TenantId; + this._userId = state.UserId; + this._useHierarchicalPartitioning = state.UseHierarchicalPartitioning; + + if (this._useHierarchicalPartitioning && this._tenantId != null && this._userId != null) + { + // Use native hierarchical partition key with PartitionKeyBuilder + this._partitionKey = new PartitionKeyBuilder() + .Add(this._tenantId) + .Add(this._userId) + .Add(conversationId) + .Build(); + } + else + { + this._partitionKey = new PartitionKey(conversationId); + } + return; } } @@ -280,7 +512,7 @@ public override async Task> GetMessagesAsync(Cancellati var iterator = this._container.GetItemQueryIterator(query, requestOptions: new QueryRequestOptions { - PartitionKey = new PartitionKey(this._conversationId), + PartitionKey = this._partitionKey, MaxItemCount = this.MaxItemCount // Configurable query performance }); @@ -345,8 +577,7 @@ public override async Task AddMessagesAsync(IEnumerable messages, C /// private async Task AddMessagesInBatchAsync(IList messages, CancellationToken cancellationToken) { - var partitionKey = new PartitionKey(this._conversationId); - var batch = this._container.CreateTransactionalBatch(partitionKey); + var batch = this._container.CreateTransactionalBatch(this._partitionKey); var currentTimestamp = DateTimeOffset.UtcNow.ToUnixTimeSeconds(); foreach (var message in messages) @@ -379,7 +610,7 @@ private async Task AddMessagesInBatchAsync(IList messages, Cancella private async Task AddSingleMessageAsync(ChatMessage message, CancellationToken cancellationToken) { var document = this.CreateMessageDocument(message, DateTimeOffset.UtcNow.ToUnixTimeSeconds()); - await this._container.CreateItemAsync(document, new PartitionKey(this._conversationId), cancellationToken: cancellationToken).ConfigureAwait(false); + await this._container.CreateItemAsync(document, this._partitionKey, cancellationToken: cancellationToken).ConfigureAwait(false); } /// @@ -398,7 +629,11 @@ private CosmosMessageDocument CreateMessageDocument(ChatMessage message, long ti Role = message.Role.Value ?? "unknown", Message = JsonSerializer.Serialize(message, s_defaultJsonOptions), Type = "ChatMessage", // Type discriminator - Ttl = this.MessageTtlSeconds // Configurable TTL + Ttl = this.MessageTtlSeconds, // Configurable TTL + // Include hierarchical metadata when using hierarchical partitioning + TenantId = this._useHierarchicalPartitioning ? this._tenantId : null, + UserId = this._useHierarchicalPartitioning ? this._userId : null, + SessionId = this._useHierarchicalPartitioning ? this._conversationId : null }; } @@ -418,7 +653,10 @@ public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptio { ConversationIdentifier = this._conversationId, DatabaseIdentifier = this.DatabaseId, - ContainerIdentifier = this.ContainerId + ContainerIdentifier = this.ContainerId, + TenantId = this._tenantId, + UserId = this._userId, + UseHierarchicalPartitioning = this._useHierarchicalPartitioning }; var options = jsonSerializerOptions ?? s_defaultJsonOptions; @@ -447,7 +685,7 @@ public async Task GetMessageCountAsync(CancellationToken cancellationToken var iterator = this._container.GetItemQueryIterator(query, requestOptions: new QueryRequestOptions { - PartitionKey = new PartitionKey(this._conversationId) + PartitionKey = this._partitionKey }); if (iterator.HasMoreResults) @@ -531,6 +769,9 @@ private sealed class StoreState public string ConversationIdentifier { get; set; } = string.Empty; public string DatabaseIdentifier { get; set; } = string.Empty; public string ContainerIdentifier { get; set; } = string.Empty; + public string? TenantId { get; set; } + public string? UserId { get; set; } + public bool UseHierarchicalPartitioning { get; set; } } /// @@ -562,5 +803,23 @@ private sealed class CosmosMessageDocument [Newtonsoft.Json.JsonProperty(nameof(Ttl))] public int? Ttl { get; set; } + + /// + /// Tenant ID for hierarchical partitioning scenarios (optional). + /// + [Newtonsoft.Json.JsonProperty("tenantId")] + public string? TenantId { get; set; } + + /// + /// User ID for hierarchical partitioning scenarios (optional). + /// + [Newtonsoft.Json.JsonProperty("userId")] + public string? UserId { get; set; } + + /// + /// Session ID for hierarchical partitioning scenarios (same as ConversationId for compatibility). + /// + [Newtonsoft.Json.JsonProperty("sessionId")] + public string? SessionId { get; set; } } } \ No newline at end of file diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/CosmosChatMessageStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/CosmosChatMessageStoreTests.cs index 66266470a3..29f71ff9fe 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/CosmosChatMessageStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/CosmosChatMessageStoreTests.cs @@ -3,6 +3,8 @@ using System; using System.Collections.Generic; using System.Linq; +using System.Text.Json; +using System.Text.Json.Serialization.Metadata; using System.Threading.Tasks; using Microsoft.Azure.Cosmos; using Microsoft.Extensions.AI; @@ -43,6 +45,7 @@ public sealed class CosmosChatMessageStoreTests : IAsyncLifetime, IDisposable private const string EmulatorEndpoint = "https://localhost:8081"; private const string EmulatorKey = "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw=="; private const string TestContainerId = "ChatMessages"; + private const string HierarchicalTestContainerId = "HierarchicalChatMessages"; // Use unique database ID per test run to avoid conflicts between parallel test executions private readonly string _testDatabaseId = $"AgentFrameworkTests-ChatStore-{Guid.NewGuid():N}"; @@ -67,11 +70,19 @@ public async Task InitializeAsync() // Test connection by attempting to create database var databaseResponse = await this._setupClient.CreateDatabaseIfNotExistsAsync(this._testDatabaseId); + + // Create container for simple partitioning tests await databaseResponse.Database.CreateContainerIfNotExistsAsync( TestContainerId, "/conversationId", throughput: 400); + // Create container for hierarchical partitioning tests with hierarchical partition key + var hierarchicalContainerProperties = new ContainerProperties(HierarchicalTestContainerId, new List { "/tenantId", "/userId", "/sessionId" }); + await databaseResponse.Database.CreateContainerIfNotExistsAsync( + hierarchicalContainerProperties, + throughput: 400); + this._emulatorAvailable = true; } catch (Exception) @@ -396,4 +407,282 @@ public void Dispose_MultipleCalls_ShouldNotThrow() } #endregion + + #region Hierarchical Partitioning Tests + + [Fact] + [Trait("Category", "CosmosDB")] + public void Constructor_WithHierarchicalConnectionString_ShouldCreateInstance() + { + // Arrange & Act + this.SkipIfEmulatorNotAvailable(); + + // Act + using var store = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789"); + + // Assert + Assert.NotNull(store); + Assert.Equal("session-789", store.ConversationId); + Assert.Equal(this._testDatabaseId, store.DatabaseId); + Assert.Equal(HierarchicalTestContainerId, store.ContainerId); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public void Constructor_WithHierarchicalEndpoint_ShouldCreateInstance() + { + // Arrange & Act + this.SkipIfEmulatorNotAvailable(); + + // Act + using var store = new CosmosChatMessageStore(EmulatorEndpoint, this._testDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789", useManagedIdentity: true); + + // Assert + Assert.NotNull(store); + Assert.Equal("session-789", store.ConversationId); + Assert.Equal(this._testDatabaseId, store.DatabaseId); + Assert.Equal(HierarchicalTestContainerId, store.ContainerId); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public void Constructor_WithHierarchicalCosmosClient_ShouldCreateInstance() + { + // Arrange & Act + this.SkipIfEmulatorNotAvailable(); + + using var cosmosClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); + using var store = new CosmosChatMessageStore(cosmosClient, this._testDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789"); + + // Assert + Assert.NotNull(store); + Assert.Equal("session-789", store.ConversationId); + Assert.Equal(this._testDatabaseId, store.DatabaseId); + Assert.Equal(HierarchicalTestContainerId, store.ContainerId); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public void Constructor_WithHierarchicalNullTenantId_ShouldThrowArgumentException() + { + // Arrange & Act & Assert + this.SkipIfEmulatorNotAvailable(); + + Assert.Throws(() => + new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, TestContainerId, null!, "user-456", "session-789")); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public void Constructor_WithHierarchicalEmptyUserId_ShouldThrowArgumentException() + { + // Arrange & Act & Assert + this.SkipIfEmulatorNotAvailable(); + + Assert.Throws(() => + new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, HierarchicalTestContainerId, "tenant-123", "", "session-789")); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public void Constructor_WithHierarchicalWhitespaceSessionId_ShouldThrowArgumentException() + { + // Arrange & Act & Assert + this.SkipIfEmulatorNotAvailable(); + + Assert.Throws(() => + new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", " ")); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public async Task AddMessagesAsync_WithHierarchicalPartitioning_ShouldAddMessageWithMetadataAsync() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + const string TenantId = "tenant-123"; + const string UserId = "user-456"; + const string SessionId = "session-789"; + // Test hierarchical partitioning constructor with connection string + using var store = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, HierarchicalTestContainerId, TenantId, UserId, SessionId); + var message = new ChatMessage(ChatRole.User, "Hello from hierarchical partitioning!"); + + // Act + await store.AddMessagesAsync([message]); + + // Wait a moment for eventual consistency + await Task.Delay(100); + + // Assert + var messages = await store.GetMessagesAsync(); + var messageList = messages.ToList(); + + Assert.Single(messageList); + Assert.Equal("Hello from hierarchical partitioning!", messageList[0].Text); + Assert.Equal(ChatRole.User, messageList[0].Role); + + // Verify that the document is stored with hierarchical partitioning metadata + var directQuery = new QueryDefinition("SELECT * FROM c WHERE c.conversationId = @conversationId AND c.Type = @type") + .WithParameter("@conversationId", SessionId) + .WithParameter("@type", "ChatMessage"); + + var iterator = this._setupClient!.GetDatabase(this._testDatabaseId).GetContainer(HierarchicalTestContainerId) + .GetItemQueryIterator(directQuery, requestOptions: new QueryRequestOptions + { + PartitionKey = new PartitionKeyBuilder().Add(TenantId).Add(UserId).Add(SessionId).Build() + }); + + var response = await iterator.ReadNextAsync(); + var document = response.FirstOrDefault(); + + Assert.NotNull(document); + // The document should have hierarchical metadata + Assert.Equal(SessionId, (string)document!.conversationId); + Assert.Equal(TenantId, (string)document!.tenantId); + Assert.Equal(UserId, (string)document!.userId); + Assert.Equal(SessionId, (string)document!.sessionId); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public async Task AddMessagesAsync_WithHierarchicalMultipleMessages_ShouldAddAllMessagesAsync() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + const string TenantId = "tenant-batch"; + const string UserId = "user-batch"; + const string SessionId = "session-batch"; + // Test hierarchical partitioning constructor with connection string + using var store = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, HierarchicalTestContainerId, TenantId, UserId, SessionId); + var messages = new[] + { + new ChatMessage(ChatRole.User, "First hierarchical message"), + new ChatMessage(ChatRole.Assistant, "Second hierarchical message"), + new ChatMessage(ChatRole.User, "Third hierarchical message") + }; + + // Act + await store.AddMessagesAsync(messages); + + // Wait a moment for eventual consistency + await Task.Delay(100); + + // Assert + var retrievedMessages = await store.GetMessagesAsync(); + var messageList = retrievedMessages.ToList(); + + Assert.Equal(3, messageList.Count); + Assert.Equal("First hierarchical message", messageList[0].Text); + Assert.Equal("Second hierarchical message", messageList[1].Text); + Assert.Equal("Third hierarchical message", messageList[2].Text); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public async Task GetMessagesAsync_WithHierarchicalPartitionIsolation_ShouldIsolateMessagesByUserIdAsync() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + const string TenantId = "tenant-isolation"; + const string UserId1 = "user-1"; + const string UserId2 = "user-2"; + const string SessionId = "session-isolation"; + + // Different userIds create different hierarchical partitions, providing proper isolation + using var store1 = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, HierarchicalTestContainerId, TenantId, UserId1, SessionId); + using var store2 = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, HierarchicalTestContainerId, TenantId, UserId2, SessionId); + + // Add messages to both stores + await store1.AddMessagesAsync([new ChatMessage(ChatRole.User, "Message from user 1")]); + await store2.AddMessagesAsync([new ChatMessage(ChatRole.User, "Message from user 2")]); + + // Wait a moment for eventual consistency + await Task.Delay(100); + + // Act & Assert + var messages1 = await store1.GetMessagesAsync(); + var messageList1 = messages1.ToList(); + + var messages2 = await store2.GetMessagesAsync(); + var messageList2 = messages2.ToList(); + + // With true hierarchical partitioning, each user sees only their own messages + Assert.Single(messageList1); + Assert.Single(messageList2); + Assert.Equal("Message from user 1", messageList1[0].Text); + Assert.Equal("Message from user 2", messageList2[0].Text); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public async Task SerializeDeserialize_WithHierarchicalPartitioning_ShouldPreserveStateAsync() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + const string TenantId = "tenant-serialize"; + const string UserId = "user-serialize"; + const string SessionId = "session-serialize"; + + using var originalStore = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, HierarchicalTestContainerId, TenantId, UserId, SessionId); + await originalStore.AddMessagesAsync([new ChatMessage(ChatRole.User, "Test serialization message")]); + + // Act - Serialize the store state + var serializedState = originalStore.Serialize(); + + // Create a new store from the serialized state + using var cosmosClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); + var serializerOptions = new JsonSerializerOptions + { + TypeInfoResolver = new DefaultJsonTypeInfoResolver() + }; + using var deserializedStore = new CosmosChatMessageStore(serializedState, cosmosClient, serializerOptions); + + // Wait a moment for eventual consistency + await Task.Delay(100); + + // Assert - The deserialized store should have the same functionality + var messages = await deserializedStore.GetMessagesAsync(); + var messageList = messages.ToList(); + + Assert.Single(messageList); + Assert.Equal("Test serialization message", messageList[0].Text); + Assert.Equal(SessionId, deserializedStore.ConversationId); + Assert.Equal(this._testDatabaseId, deserializedStore.DatabaseId); + Assert.Equal(HierarchicalTestContainerId, deserializedStore.ContainerId); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public async Task HierarchicalAndSimplePartitioning_ShouldCoexistAsync() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + const string SessionId = "coexist-session"; + + // Create simple store using simple partitioning container and hierarchical store using hierarchical container + using var simpleStore = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, TestContainerId, SessionId); + using var hierarchicalStore = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, HierarchicalTestContainerId, "tenant-coexist", "user-coexist", SessionId); + + // Add messages to both + await simpleStore.AddMessagesAsync([new ChatMessage(ChatRole.User, "Simple partitioning message")]); + await hierarchicalStore.AddMessagesAsync([new ChatMessage(ChatRole.User, "Hierarchical partitioning message")]); + + // Wait a moment for eventual consistency + await Task.Delay(100); + + // Act & Assert + var simpleMessages = await simpleStore.GetMessagesAsync(); + var simpleMessageList = simpleMessages.ToList(); + + var hierarchicalMessages = await hierarchicalStore.GetMessagesAsync(); + var hierarchicalMessageList = hierarchicalMessages.ToList(); + + // Each should only see its own messages since they use different containers + Assert.Single(simpleMessageList); + Assert.Single(hierarchicalMessageList); + Assert.Equal("Simple partitioning message", simpleMessageList[0].Text); + Assert.Equal("Hierarchical partitioning message", hierarchicalMessageList[0].Text); + } + + #endregion } From 9de2fe8735280678a5eb5df12e066209fe131e96 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Fri, 31 Oct 2025 15:20:45 +0000 Subject: [PATCH 05/40] enhance transactional batch --- .../CosmosChatMessageStore.cs | 50 ++++++++++++++++--- 1 file changed, 43 insertions(+), 7 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/CosmosChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/CosmosChatMessageStore.cs index 3709164c64..ad1dfd0eeb 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/CosmosChatMessageStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/CosmosChatMessageStore.cs @@ -56,6 +56,12 @@ private static JsonSerializerOptions CreateDefaultJsonOptions() /// public int MaxItemCount { get; set; } = 100; + /// + /// Gets or sets the maximum number of items per transactional batch operation. + /// Default is 100, maximum allowed by Cosmos DB is 100. + /// + public int MaxBatchSize { get; set; } = 100; + /// /// Gets or sets the Time-To-Live (TTL) in seconds for messages. /// Default is 86400 seconds (24 hours). Set to null to disable TTL. @@ -575,14 +581,28 @@ public override async Task AddMessagesAsync(IEnumerable messages, C /// /// Adds multiple messages using transactional batch operations for atomicity. /// - private async Task AddMessagesInBatchAsync(IList messages, CancellationToken cancellationToken) + private async Task AddMessagesInBatchAsync(List messages, CancellationToken cancellationToken) { - var batch = this._container.CreateTransactionalBatch(this._partitionKey); var currentTimestamp = DateTimeOffset.UtcNow.ToUnixTimeSeconds(); + // Process messages in optimal batch sizes + for (int i = 0; i < messages.Count; i += this.MaxBatchSize) + { + var batchMessages = messages.Skip(i).Take(this.MaxBatchSize).ToList(); + await this.ExecuteBatchOperationAsync(batchMessages, currentTimestamp, cancellationToken).ConfigureAwait(false); + } + } + + /// + /// Executes a single batch operation with retry logic and enhanced error handling. + /// + private async Task ExecuteBatchOperationAsync(List messages, long timestamp, CancellationToken cancellationToken) + { + var batch = this._container.CreateTransactionalBatch(this._partitionKey); + foreach (var message in messages) { - var document = this.CreateMessageDocument(message, currentTimestamp); + var document = this.CreateMessageDocument(message, timestamp); batch.CreateItem(document); } @@ -591,16 +611,32 @@ private async Task AddMessagesInBatchAsync(IList messages, Cancella var response = await batch.ExecuteAsync(cancellationToken).ConfigureAwait(false); if (!response.IsSuccessStatusCode) { - throw new InvalidOperationException($"Batch operation failed with status: {response.StatusCode}"); + throw new InvalidOperationException($"Batch operation failed with status: {response.StatusCode}. Details: {response.ErrorMessage}"); } } catch (CosmosException ex) when (ex.StatusCode == System.Net.HttpStatusCode.RequestEntityTooLarge) { - // Fallback to individual operations if batch is too large - foreach (var message in messages) + // If batch is too large, split into smaller batches + if (messages.Count == 1) { - await this.AddSingleMessageAsync(message, cancellationToken).ConfigureAwait(false); + // Can't split further, use single operation + await this.AddSingleMessageAsync(messages[0], cancellationToken).ConfigureAwait(false); + return; } + + // Split the batch in half and retry + var midpoint = messages.Count / 2; + var firstHalf = messages.Take(midpoint).ToList(); + var secondHalf = messages.Skip(midpoint).ToList(); + + await this.ExecuteBatchOperationAsync(firstHalf, timestamp, cancellationToken).ConfigureAwait(false); + await this.ExecuteBatchOperationAsync(secondHalf, timestamp, cancellationToken).ConfigureAwait(false); + } + catch (CosmosException ex) when (ex.StatusCode == (System.Net.HttpStatusCode)429) // TooManyRequests + { + // Handle rate limiting with exponential backoff + await Task.Delay(TimeSpan.FromMilliseconds(ex.RetryAfter?.TotalMilliseconds ?? 1000), cancellationToken).ConfigureAwait(false); + await this.ExecuteBatchOperationAsync(messages, timestamp, cancellationToken).ConfigureAwait(false); } } From b80c59ee5eafac792b9bf7f7ac70b62736194a06 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Fri, 31 Oct 2025 18:40:50 +0000 Subject: [PATCH 06/40] address review comments --- .../CosmosChatMessageStore.cs | 16 ++++------------ .../Checkpointing/CosmosCheckpointStore.cs | 15 ++++----------- .../CosmosCheckpointStoreTests.cs | 9 ++++++--- 3 files changed, 14 insertions(+), 26 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/CosmosChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/CosmosChatMessageStore.cs index ad1dfd0eeb..2790a8fb32 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/CosmosChatMessageStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/CosmosChatMessageStore.cs @@ -308,7 +308,6 @@ public CosmosChatMessageStore(string connectionString, string databaseId, string this._conversationId = sessionId; // Use sessionId as conversationId for compatibility this._databaseId = databaseId; this._containerId = containerId; - this._containerId = containerId; // Initialize hierarchical partitioning mode this._tenantId = tenantId; @@ -376,7 +375,6 @@ public CosmosChatMessageStore(string accountEndpoint, string databaseId, string this._conversationId = sessionId; // Use sessionId as conversationId for compatibility this._databaseId = databaseId; this._containerId = containerId; - this._containerId = containerId; // Initialize hierarchical partitioning mode this._tenantId = tenantId; @@ -478,19 +476,13 @@ public CosmosChatMessageStore(JsonElement serializedStoreState, CosmosClient cos this._userId = state.UserId; this._useHierarchicalPartitioning = state.UseHierarchicalPartitioning; - if (this._useHierarchicalPartitioning && this._tenantId != null && this._userId != null) - { - // Use native hierarchical partition key with PartitionKeyBuilder - this._partitionKey = new PartitionKeyBuilder() + this._partitionKey = (this._useHierarchicalPartitioning && this._tenantId != null && this._userId != null) + ? new PartitionKeyBuilder() .Add(this._tenantId) .Add(this._userId) .Add(conversationId) - .Build(); - } - else - { - this._partitionKey = new PartitionKey(conversationId); - } + .Build() + : new PartitionKey(conversationId); return; } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/CosmosCheckpointStore.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/CosmosCheckpointStore.cs index c9814e0456..3c3950f16d 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/CosmosCheckpointStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/CosmosCheckpointStore.cs @@ -175,19 +175,12 @@ public override async ValueTask> RetrieveIndexAsync( throw new ObjectDisposedException(GetType().FullName); #pragma warning restore CA1513 - QueryDefinition query; - - if (withParent == null) - { - query = new QueryDefinition("SELECT c.runId, c.checkpointId FROM c WHERE c.runId = @runId ORDER BY c.timestamp ASC") - .WithParameter("@runId", runId); - } - else - { - query = new QueryDefinition("SELECT c.runId, c.checkpointId FROM c WHERE c.runId = @runId AND c.parentCheckpointId = @parentCheckpointId ORDER BY c.timestamp ASC") + QueryDefinition query = withParent == null + ? new QueryDefinition("SELECT c.runId, c.checkpointId FROM c WHERE c.runId = @runId ORDER BY c.timestamp ASC") + .WithParameter("@runId", runId) + : new QueryDefinition("SELECT c.runId, c.checkpointId FROM c WHERE c.runId = @runId AND c.parentCheckpointId = @parentCheckpointId ORDER BY c.timestamp ASC") .WithParameter("@runId", runId) .WithParameter("@parentCheckpointId", withParent.CheckpointId); - } var iterator = _container.GetItemQueryIterator(query); var checkpoints = new List(); diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/CosmosCheckpointStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/CosmosCheckpointStoreTests.cs index 9f323ff60e..697c8879d8 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/CosmosCheckpointStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/CosmosCheckpointStoreTests.cs @@ -64,7 +64,7 @@ public async Task InitializeAsync() _emulatorAvailable = true; } - catch (Exception) + catch (Exception ex) when (!(ex is OutOfMemoryException || ex is StackOverflowException || ex is AccessViolationException)) { // Emulator not available, tests will be skipped _emulatorAvailable = false; @@ -91,9 +91,10 @@ public async Task DisposeAsync() await _database!.DeleteAsync(); } } - catch + catch (Exception ex) { - // Ignore cleanup errors + // Ignore cleanup errors, but log for diagnostics + Console.WriteLine($"[DisposeAsync] Cleanup error: {ex.Message}\n{ex.StackTrace}"); } finally { @@ -116,6 +117,7 @@ private void SkipIfEmulatorNotAvailable() [Fact] public void Constructor_WithCosmosClient_SetsProperties() { + // Arrange SkipIfEmulatorNotAvailable(); // Act @@ -129,6 +131,7 @@ public void Constructor_WithCosmosClient_SetsProperties() [Fact] public void Constructor_WithConnectionString_SetsProperties() { + // Arrange SkipIfEmulatorNotAvailable(); // Act From 042db6258589632ea0a59c7f50dbb984621c2149 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Fri, 7 Nov 2025 15:47:17 +0000 Subject: [PATCH 07/40] Address PR review comments from @westey-m --- dotnet/agent-framework-dotnet.slnx | 2 + .../CosmosChatMessageStore.cs | 18 +- .../CosmosCheckpointStore.cs | 5 + .../CosmosDBChatExtensions.cs | 7 + .../CosmosDBWorkflowExtensions.cs | 13 + ....Agents.AI.Abstractions.CosmosNoSql.csproj | 42 ++ .../CosmosAgentThread.cs | 233 -------- .../CosmosChatMessageStoreTests.cs | 91 ++-- .../CosmosCheckpointStoreTests.cs | 79 +-- ....Abstractions.CosmosNoSql.UnitTests.csproj | 23 + .../CosmosAgentThreadTests.cs | 507 ------------------ ...ft.Agents.AI.Abstractions.UnitTests.csproj | 2 - ...osoft.Agents.AI.Workflows.UnitTests.csproj | 2 - 13 files changed, 193 insertions(+), 831 deletions(-) rename dotnet/src/{Microsoft.Agents.AI.Abstractions => Microsoft.Agents.AI.Abstractions.CosmosNoSql}/CosmosChatMessageStore.cs (97%) rename dotnet/src/{Microsoft.Agents.AI.Workflows/Checkpointing => Microsoft.Agents.AI.Abstractions.CosmosNoSql}/CosmosCheckpointStore.cs (96%) rename dotnet/src/{Microsoft.Agents.AI => Microsoft.Agents.AI.Abstractions.CosmosNoSql}/CosmosDBChatExtensions.cs (87%) rename dotnet/src/{Microsoft.Agents.AI.Workflows => Microsoft.Agents.AI.Abstractions.CosmosNoSql}/CosmosDBWorkflowExtensions.cs (85%) create mode 100644 dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/Microsoft.Agents.AI.Abstractions.CosmosNoSql.csproj delete mode 100644 dotnet/src/Microsoft.Agents.AI.Abstractions/CosmosAgentThread.cs rename dotnet/tests/{Microsoft.Agents.AI.Abstractions.UnitTests => Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests}/CosmosChatMessageStoreTests.cs (87%) rename dotnet/tests/{Microsoft.Agents.AI.Workflows.UnitTests => Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests}/CosmosCheckpointStoreTests.cs (87%) create mode 100644 dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests.csproj delete mode 100644 dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/CosmosAgentThreadTests.cs diff --git a/dotnet/agent-framework-dotnet.slnx b/dotnet/agent-framework-dotnet.slnx index cf1e367293..e8f6b521d2 100644 --- a/dotnet/agent-framework-dotnet.slnx +++ b/dotnet/agent-framework-dotnet.slnx @@ -262,6 +262,7 @@ + @@ -288,6 +289,7 @@ + diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/CosmosChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosChatMessageStore.cs similarity index 97% rename from dotnet/src/Microsoft.Agents.AI.Abstractions/CosmosChatMessageStore.cs rename to dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosChatMessageStore.cs index 2790a8fb32..35e75581f9 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/CosmosChatMessageStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosChatMessageStore.cs @@ -17,6 +17,8 @@ namespace Microsoft.Agents.AI; /// /// Provides a Cosmos DB implementation of the abstract class. /// +[RequiresUnreferencedCode("The CosmosChatMessageStore uses JSON serialization which is incompatible with trimming.")] +[RequiresDynamicCode("The CosmosChatMessageStore uses JSON serialization which is incompatible with NativeAOT.")] public sealed class CosmosChatMessageStore : ChatMessageStore, IDisposable { private readonly CosmosClient _cosmosClient; @@ -38,8 +40,6 @@ public sealed class CosmosChatMessageStore : ChatMessageStore, IDisposable /// private static readonly JsonSerializerOptions s_defaultJsonOptions = CreateDefaultJsonOptions(); - [UnconditionalSuppressMessage("Trimming", "IL2026:Members annotated with 'RequiresUnreferencedCodeAttribute' require dynamic access", Justification = "JSON serialization is controlled")] - [UnconditionalSuppressMessage("AOT", "IL3050:Calling members annotated with 'RequiresDynamicCodeAttribute' may break functionality when AOT compiling", Justification = "JSON serialization is controlled")] private static JsonSerializerOptions CreateDefaultJsonOptions() { var options = new JsonSerializerOptions(); @@ -504,7 +504,7 @@ public override async Task> GetMessagesAsync(Cancellati #pragma warning restore CA1513 // Use type discriminator for efficient queries - var query = new QueryDefinition("SELECT * FROM c WHERE c.conversationId = @conversationId AND c.Type = @type ORDER BY c.Timestamp ASC") + var query = new QueryDefinition("SELECT * FROM c WHERE c.conversationId = @conversationId AND c.type = @type ORDER BY c.timestamp ASC") .WithParameter("@conversationId", this._conversationId) .WithParameter("@type", "ChatMessage"); @@ -814,22 +814,22 @@ private sealed class CosmosMessageDocument [Newtonsoft.Json.JsonProperty("conversationId")] public string ConversationId { get; set; } = string.Empty; - [Newtonsoft.Json.JsonProperty(nameof(Timestamp))] + [Newtonsoft.Json.JsonProperty("timestamp")] public long Timestamp { get; set; } - [Newtonsoft.Json.JsonProperty(nameof(MessageId))] + [Newtonsoft.Json.JsonProperty("messageId")] public string MessageId { get; set; } = string.Empty; - [Newtonsoft.Json.JsonProperty(nameof(Role))] + [Newtonsoft.Json.JsonProperty("role")] public string Role { get; set; } = string.Empty; - [Newtonsoft.Json.JsonProperty(nameof(Message))] + [Newtonsoft.Json.JsonProperty("message")] public string Message { get; set; } = string.Empty; - [Newtonsoft.Json.JsonProperty(nameof(Type))] + [Newtonsoft.Json.JsonProperty("type")] public string Type { get; set; } = string.Empty; - [Newtonsoft.Json.JsonProperty(nameof(Ttl))] + [Newtonsoft.Json.JsonProperty("ttl")] public int? Ttl { get; set; } /// diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/CosmosCheckpointStore.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosCheckpointStore.cs similarity index 96% rename from dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/CosmosCheckpointStore.cs rename to dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosCheckpointStore.cs index 3c3950f16d..2b9541ed9a 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Checkpointing/CosmosCheckpointStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosCheckpointStore.cs @@ -2,6 +2,7 @@ using System; using System.Collections.Generic; +using System.Diagnostics.CodeAnalysis; using System.Linq; using System.Text.Json; using System.Threading.Tasks; @@ -16,6 +17,8 @@ namespace Microsoft.Agents.AI.Workflows.Checkpointing; /// Provides a Cosmos DB implementation of the abstract class. /// /// The type of objects to store as checkpoint values. +[RequiresUnreferencedCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with trimming.")] +[RequiresDynamicCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with NativeAOT.")] public class CosmosCheckpointStore : JsonCheckpointStore, IDisposable { private readonly CosmosClient _cosmosClient; @@ -255,6 +258,8 @@ private sealed class CheckpointQueryResult /// /// Provides a non-generic Cosmos DB implementation of the abstract class. /// +[RequiresUnreferencedCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with trimming.")] +[RequiresDynamicCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with NativeAOT.")] public sealed class CosmosCheckpointStore : CosmosCheckpointStore { /// diff --git a/dotnet/src/Microsoft.Agents.AI/CosmosDBChatExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosDBChatExtensions.cs similarity index 87% rename from dotnet/src/Microsoft.Agents.AI/CosmosDBChatExtensions.cs rename to dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosDBChatExtensions.cs index ba42656bd6..a6f942ae11 100644 --- a/dotnet/src/Microsoft.Agents.AI/CosmosDBChatExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosDBChatExtensions.cs @@ -1,6 +1,7 @@ // Copyright (c) Microsoft. All rights reserved. using System; +using System.Diagnostics.CodeAnalysis; using Microsoft.Azure.Cosmos; namespace Microsoft.Agents.AI; @@ -20,6 +21,8 @@ public static class CosmosDBChatExtensions /// The configured . /// Thrown when is null. /// Thrown when any string parameter is null or whitespace. + [RequiresUnreferencedCode("The CosmosChatMessageStore uses JSON serialization which is incompatible with trimming.")] + [RequiresDynamicCode("The CosmosChatMessageStore uses JSON serialization which is incompatible with NativeAOT.")] public static ChatClientAgentOptions WithCosmosDBMessageStore( this ChatClientAgentOptions options, string connectionString, @@ -49,6 +52,8 @@ public static ChatClientAgentOptions WithCosmosDBMessageStore( /// The configured . /// Thrown when is null. /// Thrown when any string parameter is null or whitespace. + [RequiresUnreferencedCode("The CosmosChatMessageStore uses JSON serialization which is incompatible with trimming.")] + [RequiresDynamicCode("The CosmosChatMessageStore uses JSON serialization which is incompatible with NativeAOT.")] public static ChatClientAgentOptions WithCosmosDBMessageStoreUsingManagedIdentity( this ChatClientAgentOptions options, string accountEndpoint, @@ -78,6 +83,8 @@ public static ChatClientAgentOptions WithCosmosDBMessageStoreUsingManagedIdentit /// The configured . /// Thrown when any required parameter is null. /// Thrown when any string parameter is null or whitespace. + [RequiresUnreferencedCode("The CosmosChatMessageStore uses JSON serialization which is incompatible with trimming.")] + [RequiresDynamicCode("The CosmosChatMessageStore uses JSON serialization which is incompatible with NativeAOT.")] public static ChatClientAgentOptions WithCosmosDBMessageStore( this ChatClientAgentOptions options, CosmosClient cosmosClient, diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/CosmosDBWorkflowExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosDBWorkflowExtensions.cs similarity index 85% rename from dotnet/src/Microsoft.Agents.AI.Workflows/CosmosDBWorkflowExtensions.cs rename to dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosDBWorkflowExtensions.cs index b36651f3f3..98bdc6efd7 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/CosmosDBWorkflowExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosDBWorkflowExtensions.cs @@ -1,6 +1,7 @@ // Copyright (c) Microsoft. All rights reserved. using System; +using System.Diagnostics.CodeAnalysis; using System.Text.Json; using Microsoft.Azure.Cosmos; using Microsoft.Agents.AI.Workflows.Checkpointing; @@ -20,6 +21,8 @@ public static class CosmosDBWorkflowExtensions /// The identifier of the Cosmos DB container. /// A new instance of . /// Thrown when any string parameter is null or whitespace. + [RequiresUnreferencedCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with trimming.")] + [RequiresDynamicCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with NativeAOT.")] public static CosmosCheckpointStore CreateCheckpointStore( string connectionString, string databaseId, @@ -43,6 +46,8 @@ public static CosmosCheckpointStore CreateCheckpointStore( /// The identifier of the Cosmos DB container. /// A new instance of . /// Thrown when any string parameter is null or whitespace. + [RequiresUnreferencedCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with trimming.")] + [RequiresDynamicCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with NativeAOT.")] public static CosmosCheckpointStore CreateCheckpointStoreUsingManagedIdentity( string accountEndpoint, string databaseId, @@ -67,6 +72,8 @@ public static CosmosCheckpointStore CreateCheckpointStoreUsingManagedIdentity( /// A new instance of . /// Thrown when any required parameter is null. /// Thrown when any string parameter is null or whitespace. + [RequiresUnreferencedCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with trimming.")] + [RequiresDynamicCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with NativeAOT.")] public static CosmosCheckpointStore CreateCheckpointStore( CosmosClient cosmosClient, string databaseId, @@ -91,6 +98,8 @@ public static CosmosCheckpointStore CreateCheckpointStore( /// The identifier of the Cosmos DB container. /// A new instance of . /// Thrown when any string parameter is null or whitespace. + [RequiresUnreferencedCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with trimming.")] + [RequiresDynamicCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with NativeAOT.")] public static CosmosCheckpointStore CreateCheckpointStore( string connectionString, string databaseId, @@ -115,6 +124,8 @@ public static CosmosCheckpointStore CreateCheckpointStore( /// The identifier of the Cosmos DB container. /// A new instance of . /// Thrown when any string parameter is null or whitespace. + [RequiresUnreferencedCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with trimming.")] + [RequiresDynamicCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with NativeAOT.")] public static CosmosCheckpointStore CreateCheckpointStoreUsingManagedIdentity( string accountEndpoint, string databaseId, @@ -140,6 +151,8 @@ public static CosmosCheckpointStore CreateCheckpointStoreUsingManagedIdentity /// A new instance of . /// Thrown when any required parameter is null. /// Thrown when any string parameter is null or whitespace. + [RequiresUnreferencedCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with trimming.")] + [RequiresDynamicCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with NativeAOT.")] public static CosmosCheckpointStore CreateCheckpointStore( CosmosClient cosmosClient, string databaseId, diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/Microsoft.Agents.AI.Abstractions.CosmosNoSql.csproj b/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/Microsoft.Agents.AI.Abstractions.CosmosNoSql.csproj new file mode 100644 index 0000000000..7bdf5e227a --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/Microsoft.Agents.AI.Abstractions.CosmosNoSql.csproj @@ -0,0 +1,42 @@ + + + + $(ProjectsTargetFrameworks) + $(ProjectsDebugTargetFrameworks) + Microsoft.Agents.AI + $(NoWarn);MEAI001 + preview + + + + true + true + true + true + true + true + + + + + + + Microsoft Agent Framework Cosmos DB NoSQL Integration + Provides Cosmos DB NoSQL implementations for Microsoft Agent Framework storage abstractions including ChatMessageStore and CheckpointStore. + + + + + + + + + + + + + + + + + diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/CosmosAgentThread.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/CosmosAgentThread.cs deleted file mode 100644 index 9c3be1ed30..0000000000 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/CosmosAgentThread.cs +++ /dev/null @@ -1,233 +0,0 @@ -// Copyright (c) Microsoft. All rights reserved. - -using System; -using System.Collections.Generic; -using System.Diagnostics; -using System.Text.Json; -using System.Threading; -using System.Threading.Tasks; -using Microsoft.Azure.Cosmos; -using Microsoft.Extensions.AI; - -namespace Microsoft.Agents.AI; - -/// -/// Provides an abstract base class for agent threads that maintain conversation state in Azure Cosmos DB. -/// -/// -/// -/// is designed for scenarios where conversation state should be persisted -/// in Azure Cosmos DB for durability, scalability, and cross-session availability. This approach provides -/// reliable persistence while maintaining efficient access to conversation data. -/// -/// -/// Cosmos threads persist conversation data across application restarts and can be shared across -/// multiple application instances. -/// -/// -[DebuggerDisplay("{DebuggerDisplay,nq}")] -public abstract class CosmosAgentThread : AgentThread -{ - /// - /// Initializes a new instance of the class. - /// - /// The Cosmos DB connection string. - /// The identifier of the Cosmos DB database. - /// The identifier of the Cosmos DB container. - /// The unique identifier for this conversation thread. If null, a new GUID will be generated. - /// Thrown when any string parameter is null or whitespace. - /// - /// This constructor creates a new Cosmos DB message store with connection string authentication. - /// - protected CosmosAgentThread(string connectionString, string databaseId, string containerId, string? conversationId = null) - { - this.MessageStore = new CosmosChatMessageStore(connectionString, databaseId, containerId, conversationId ?? Guid.NewGuid().ToString("N")); - } - - /// - /// Initializes a new instance of the class using managed identity. - /// - /// The Cosmos DB account endpoint URI. - /// The identifier of the Cosmos DB database. - /// The identifier of the Cosmos DB container. - /// The unique identifier for this conversation thread. If null, a new GUID will be generated. - /// Must be true to use this constructor. This parameter distinguishes this constructor from the connection string version. - /// Thrown when any string parameter is null or whitespace, or when useManagedIdentity is false. - /// - /// This constructor creates a new Cosmos DB message store with managed identity authentication. - /// - protected CosmosAgentThread(string accountEndpoint, string databaseId, string containerId, string? conversationId, bool useManagedIdentity) - { - if (!useManagedIdentity) - { - throw new ArgumentException("This constructor requires useManagedIdentity to be true.", nameof(useManagedIdentity)); - } - - this.MessageStore = new CosmosChatMessageStore(accountEndpoint, databaseId, containerId, conversationId ?? Guid.NewGuid().ToString("N"), useManagedIdentity: true); - } - - /// - /// Initializes a new instance of the class using an existing CosmosClient. - /// - /// The instance to use for Cosmos DB operations. - /// The identifier of the Cosmos DB database. - /// The identifier of the Cosmos DB container. - /// The unique identifier for this conversation thread. If null, a new GUID will be generated. - /// Thrown when is null. - /// Thrown when any string parameter is null or whitespace. - /// - /// This constructor allows reuse of an existing CosmosClient instance across multiple threads. - /// - protected CosmosAgentThread(CosmosClient cosmosClient, string databaseId, string containerId, string? conversationId = null) - { - this.MessageStore = new CosmosChatMessageStore(cosmosClient, databaseId, containerId, conversationId ?? Guid.NewGuid().ToString("N")); - } - - /// - /// Initializes a new instance of the class with a pre-configured message store. - /// - /// - /// A instance to use for storing chat messages. - /// If , an exception will be thrown. - /// - /// is . - /// - /// This constructor allows sharing of message stores between threads or providing pre-configured - /// message stores with specific settings. - /// - protected CosmosAgentThread(CosmosChatMessageStore messageStore) - { - this.MessageStore = messageStore ?? throw new ArgumentNullException(nameof(messageStore)); - } - - /// - /// Initializes a new instance of the class from previously serialized state. - /// - /// A representing the serialized state of the thread. - /// - /// Factory function to create the from its serialized state. - /// This is required because Cosmos DB connection information cannot be reconstructed from serialized state. - /// - /// Optional settings for customizing the JSON deserialization process. - /// The is not a JSON object. - /// is . - /// The is invalid or cannot be deserialized to the expected type. - /// - /// This constructor enables restoration of Cosmos threads from previously saved state. Since Cosmos DB - /// connection information cannot be serialized for security reasons, a factory function must be provided - /// to reconstruct the message store with appropriate connection details. - /// - protected CosmosAgentThread( - JsonElement serializedThreadState, - Func messageStoreFactory, - JsonSerializerOptions? jsonSerializerOptions = null) - { -#if NET6_0_OR_GREATER - ArgumentNullException.ThrowIfNull(messageStoreFactory); -#else - if (messageStoreFactory is null) - { - throw new ArgumentNullException(nameof(messageStoreFactory)); - } -#endif - - if (serializedThreadState.ValueKind != JsonValueKind.Object) - { - throw new ArgumentException("The serialized thread state must be a JSON object.", nameof(serializedThreadState)); - } - - var state = serializedThreadState.Deserialize( - AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(CosmosAgentThreadState))) as CosmosAgentThreadState; - - this.MessageStore = messageStoreFactory.Invoke(state?.StoreState ?? default, jsonSerializerOptions); - } - - /// - /// Gets the used by this thread. - /// - public CosmosChatMessageStore MessageStore { get; } - - /// - /// Gets the conversation ID for this thread from the underlying message store. - /// - public string ConversationId => this.MessageStore.ConversationId; - - /// - /// Gets or sets the maximum number of messages to return in a single query batch. - /// This is delegated to the underlying message store. - /// - public int MaxItemCount - { - get => this.MessageStore.MaxItemCount; - set => this.MessageStore.MaxItemCount = value; - } - - /// - /// Gets or sets the Time-To-Live (TTL) in seconds for messages. - /// This is delegated to the underlying message store. - /// - public int? MessageTtlSeconds - { - get => this.MessageStore.MessageTtlSeconds; - set => this.MessageStore.MessageTtlSeconds = value; - } - - /// - /// Serializes the current object's state to a using the specified serialization options. - /// - /// The JSON serialization options to use. - /// A representation of the object's state. - /// - /// Note that connection strings and credentials are not included in the serialized state for security reasons. - /// When deserializing, you will need to provide connection information through the messageStoreFactory parameter. - /// - public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) - { - var storeState = this.MessageStore.Serialize(jsonSerializerOptions); - - var state = new CosmosAgentThreadState - { - StoreState = storeState, - }; - - return JsonSerializer.SerializeToElement(state, AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(CosmosAgentThreadState))); - } - - /// - public override object? GetService(Type serviceType, object? serviceKey = null) => - base.GetService(serviceType, serviceKey) ?? this.MessageStore?.GetService(serviceType, serviceKey); - - /// - protected internal override Task MessagesReceivedAsync(IEnumerable newMessages, CancellationToken cancellationToken = default) - => this.MessageStore.AddMessagesAsync(newMessages, cancellationToken); - - /// - /// Gets the total number of messages in this conversation. - /// This is a Cosmos-specific optimization that provides efficient message counting. - /// - /// The cancellation token. - /// The total number of messages in the conversation. - public async Task GetMessageCountAsync(CancellationToken cancellationToken = default) - { - return await this.MessageStore.GetMessageCountAsync(cancellationToken).ConfigureAwait(false); - } - - /// - /// Clears all messages in this conversation. - /// This is a Cosmos-specific utility method for conversation cleanup. - /// - /// The cancellation token. - /// The number of messages that were deleted. - public async Task ClearMessagesAsync(CancellationToken cancellationToken = default) - { - return await this.MessageStore.ClearMessagesAsync(cancellationToken).ConfigureAwait(false); - } - - [DebuggerBrowsable(DebuggerBrowsableState.Never)] - private string DebuggerDisplay => $"ConversationId = {this.ConversationId}"; - - internal sealed class CosmosAgentThreadState - { - public JsonElement? StoreState { get; set; } - } -} \ No newline at end of file diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/CosmosChatMessageStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs similarity index 87% rename from dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/CosmosChatMessageStoreTests.cs rename to dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs index 29f71ff9fe..23432f0fec 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/CosmosChatMessageStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs @@ -11,7 +11,7 @@ using Microsoft.Agents.AI; using Xunit; -namespace Microsoft.Agents.AI.Abstractions.UnitTests; +namespace Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests; /// /// Contains tests for . @@ -34,9 +34,9 @@ namespace Microsoft.Agents.AI.Abstractions.UnitTests; /// | COSMOS_PRESERVE_CONTAINERS | true / false | Controls whether to preserve test data after completion | /// /// Usage Examples: -/// - Run all tests in preserve mode: $env:COSMOS_PRESERVE_CONTAINERS="true"; dotnet test tests/Microsoft.Agents.AI.Abstractions.UnitTests/ -/// - Run specific test category in preserve mode: $env:COSMOS_PRESERVE_CONTAINERS="true"; dotnet test tests/Microsoft.Agents.AI.Abstractions.UnitTests/ --filter "Category=CosmosDB" -/// - Reset to cleanup mode: $env:COSMOS_PRESERVE_CONTAINERS=""; dotnet test tests/Microsoft.Agents.AI.Abstractions.UnitTests/ +/// - Run all tests in preserve mode: $env:COSMOS_PRESERVE_CONTAINERS="true"; dotnet test tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/ +/// - Run specific test category in preserve mode: $env:COSMOS_PRESERVE_CONTAINERS="true"; dotnet test tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/ --filter "Category=CosmosDB" +/// - Reset to cleanup mode: $env:COSMOS_PRESERVE_CONTAINERS=""; dotnet test tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/ /// [Collection("CosmosDB")] public sealed class CosmosChatMessageStoreTests : IAsyncLifetime, IDisposable @@ -46,9 +46,10 @@ public sealed class CosmosChatMessageStoreTests : IAsyncLifetime, IDisposable private const string EmulatorKey = "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw=="; private const string TestContainerId = "ChatMessages"; private const string HierarchicalTestContainerId = "HierarchicalChatMessages"; - - // Use unique database ID per test run to avoid conflicts between parallel test executions - private readonly string _testDatabaseId = $"AgentFrameworkTests-ChatStore-{Guid.NewGuid():N}"; + // Use fixed database ID for preserve mode inspection +#pragma warning disable CA1802 // Use literals where appropriate + private static readonly string TestDatabaseId = "AgentFrameworkTests-ChatStore"; +#pragma warning restore CA1802 private string _connectionString = string.Empty; private bool _emulatorAvailable; @@ -69,7 +70,7 @@ public async Task InitializeAsync() this._setupClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); // Test connection by attempting to create database - var databaseResponse = await this._setupClient.CreateDatabaseIfNotExistsAsync(this._testDatabaseId); + var databaseResponse = await this._setupClient.CreateDatabaseIfNotExistsAsync(TestDatabaseId); // Create container for simple partitioning tests await databaseResponse.Database.CreateContainerIfNotExistsAsync( @@ -109,7 +110,7 @@ public async Task DisposeAsync() else { // Clean mode: Delete the test database and all data - var database = this._setupClient.GetDatabase(this._testDatabaseId); + var database = this._setupClient.GetDatabase(TestDatabaseId); await database.DeleteAsync(); } } @@ -149,12 +150,12 @@ public void Constructor_WithConnectionString_ShouldCreateInstance() this.SkipIfEmulatorNotAvailable(); // Act - using var store = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, TestContainerId, "test-conversation"); + using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, "test-conversation"); // Assert Assert.NotNull(store); Assert.Equal("test-conversation", store.ConversationId); - Assert.Equal(this._testDatabaseId, store.DatabaseId); + Assert.Equal(TestDatabaseId, store.DatabaseId); Assert.Equal(TestContainerId, store.ContainerId); } @@ -166,12 +167,12 @@ public void Constructor_WithConnectionStringNoConversationId_ShouldCreateInstanc this.SkipIfEmulatorNotAvailable(); // Act - using var store = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, TestContainerId); + using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId); // Assert Assert.NotNull(store); Assert.NotNull(store.ConversationId); - Assert.Equal(this._testDatabaseId, store.DatabaseId); + Assert.Equal(TestDatabaseId, store.DatabaseId); Assert.Equal(TestContainerId, store.ContainerId); } @@ -181,7 +182,7 @@ public void Constructor_WithNullConnectionString_ShouldThrowArgumentException() { // Arrange & Act & Assert Assert.Throws(() => - new CosmosChatMessageStore((string)null!, this._testDatabaseId, TestContainerId, "test-conversation")); + new CosmosChatMessageStore((string)null!, TestDatabaseId, TestContainerId, "test-conversation")); } [Fact] @@ -192,7 +193,7 @@ public void Constructor_WithEmptyConversationId_ShouldThrowArgumentException() this.SkipIfEmulatorNotAvailable(); Assert.Throws(() => - new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, TestContainerId, "")); + new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, "")); } #endregion @@ -206,7 +207,7 @@ public async Task AddMessagesAsync_WithSingleMessage_ShouldAddMessageAsync() // Arrange this.SkipIfEmulatorNotAvailable(); var conversationId = Guid.NewGuid().ToString(); - using var store = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, TestContainerId, conversationId); + using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, conversationId); var message = new ChatMessage(ChatRole.User, "Hello, world!"); // Act @@ -225,7 +226,7 @@ public async Task AddMessagesAsync_WithSingleMessage_ShouldAddMessageAsync() // Let's check if we can find ANY items in the container for this conversation var directQuery = new QueryDefinition("SELECT VALUE COUNT(1) FROM c WHERE c.conversationId = @conversationId") .WithParameter("@conversationId", conversationId); - var countIterator = this._setupClient!.GetDatabase(this._testDatabaseId).GetContainer(TestContainerId) + var countIterator = this._setupClient!.GetDatabase(TestDatabaseId).GetContainer(TestContainerId) .GetItemQueryIterator(directQuery, requestOptions: new QueryRequestOptions { PartitionKey = new PartitionKey(conversationId) @@ -237,7 +238,7 @@ public async Task AddMessagesAsync_WithSingleMessage_ShouldAddMessageAsync() // Debug: Let's see what the raw query returns var rawQuery = new QueryDefinition("SELECT * FROM c WHERE c.conversationId = @conversationId") .WithParameter("@conversationId", conversationId); - var rawIterator = this._setupClient!.GetDatabase(this._testDatabaseId).GetContainer(TestContainerId) + var rawIterator = this._setupClient!.GetDatabase(TestDatabaseId).GetContainer(TestContainerId) .GetItemQueryIterator(rawQuery, requestOptions: new QueryRequestOptions { PartitionKey = new PartitionKey(conversationId) @@ -266,7 +267,7 @@ public async Task AddMessagesAsync_WithMultipleMessages_ShouldAddAllMessagesAsyn // Arrange SkipIfEmulatorNotAvailable(); var conversationId = Guid.NewGuid().ToString(); - using var store = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, TestContainerId, conversationId); + using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, conversationId); var messages = new[] { new ChatMessage(ChatRole.User, "First message"), @@ -296,7 +297,7 @@ public async Task GetMessagesAsync_WithNoMessages_ShouldReturnEmptyAsync() { // Arrange SkipIfEmulatorNotAvailable(); - using var store = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, TestContainerId, Guid.NewGuid().ToString()); + using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, Guid.NewGuid().ToString()); // Act var messages = await store.GetMessagesAsync(); @@ -314,8 +315,8 @@ public async Task GetMessagesAsync_WithConversationIsolation_ShouldOnlyReturnMes var conversation1 = Guid.NewGuid().ToString(); var conversation2 = Guid.NewGuid().ToString(); - using var store1 = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, TestContainerId, conversation1); - using var store2 = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, TestContainerId, conversation2); + using var store1 = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, conversation1); + using var store2 = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, conversation2); await store1.AddMessagesAsync([new ChatMessage(ChatRole.User, "Message for conversation 1")]); await store2.AddMessagesAsync([new ChatMessage(ChatRole.User, "Message for conversation 2")]); @@ -344,7 +345,7 @@ public async Task FullWorkflow_AddAndGet_ShouldWorkCorrectlyAsync() // Arrange SkipIfEmulatorNotAvailable(); var conversationId = $"test-conversation-{Guid.NewGuid():N}"; // Use unique conversation ID - using var originalStore = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, TestContainerId, conversationId); + using var originalStore = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, conversationId); var messages = new[] { @@ -364,7 +365,7 @@ public async Task FullWorkflow_AddAndGet_ShouldWorkCorrectlyAsync() Assert.Equal(5, retrievedList.Count); // Act 3: Create new store instance for same conversation (test persistence) - using var newStore = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, TestContainerId, conversationId); + using var newStore = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, conversationId); var persistedMessages = await newStore.GetMessagesAsync(); var persistedList = persistedMessages.ToList(); @@ -387,7 +388,7 @@ public void Dispose_AfterUse_ShouldNotThrow() { // Arrange SkipIfEmulatorNotAvailable(); - var store = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, TestContainerId, Guid.NewGuid().ToString()); + var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, Guid.NewGuid().ToString()); // Act & Assert store.Dispose(); // Should not throw @@ -399,7 +400,7 @@ public void Dispose_MultipleCalls_ShouldNotThrow() { // Arrange SkipIfEmulatorNotAvailable(); - var store = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, TestContainerId, Guid.NewGuid().ToString()); + var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, Guid.NewGuid().ToString()); // Act & Assert store.Dispose(); // First call @@ -418,12 +419,12 @@ public void Constructor_WithHierarchicalConnectionString_ShouldCreateInstance() this.SkipIfEmulatorNotAvailable(); // Act - using var store = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789"); + using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789"); // Assert Assert.NotNull(store); Assert.Equal("session-789", store.ConversationId); - Assert.Equal(this._testDatabaseId, store.DatabaseId); + Assert.Equal(TestDatabaseId, store.DatabaseId); Assert.Equal(HierarchicalTestContainerId, store.ContainerId); } @@ -435,12 +436,12 @@ public void Constructor_WithHierarchicalEndpoint_ShouldCreateInstance() this.SkipIfEmulatorNotAvailable(); // Act - using var store = new CosmosChatMessageStore(EmulatorEndpoint, this._testDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789", useManagedIdentity: true); + using var store = new CosmosChatMessageStore(EmulatorEndpoint, TestDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789", useManagedIdentity: true); // Assert Assert.NotNull(store); Assert.Equal("session-789", store.ConversationId); - Assert.Equal(this._testDatabaseId, store.DatabaseId); + Assert.Equal(TestDatabaseId, store.DatabaseId); Assert.Equal(HierarchicalTestContainerId, store.ContainerId); } @@ -452,12 +453,12 @@ public void Constructor_WithHierarchicalCosmosClient_ShouldCreateInstance() this.SkipIfEmulatorNotAvailable(); using var cosmosClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); - using var store = new CosmosChatMessageStore(cosmosClient, this._testDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789"); + using var store = new CosmosChatMessageStore(cosmosClient, TestDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789"); // Assert Assert.NotNull(store); Assert.Equal("session-789", store.ConversationId); - Assert.Equal(this._testDatabaseId, store.DatabaseId); + Assert.Equal(TestDatabaseId, store.DatabaseId); Assert.Equal(HierarchicalTestContainerId, store.ContainerId); } @@ -469,7 +470,7 @@ public void Constructor_WithHierarchicalNullTenantId_ShouldThrowArgumentExceptio this.SkipIfEmulatorNotAvailable(); Assert.Throws(() => - new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, TestContainerId, null!, "user-456", "session-789")); + new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, null!, "user-456", "session-789")); } [Fact] @@ -480,7 +481,7 @@ public void Constructor_WithHierarchicalEmptyUserId_ShouldThrowArgumentException this.SkipIfEmulatorNotAvailable(); Assert.Throws(() => - new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, HierarchicalTestContainerId, "tenant-123", "", "session-789")); + new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, "tenant-123", "", "session-789")); } [Fact] @@ -491,7 +492,7 @@ public void Constructor_WithHierarchicalWhitespaceSessionId_ShouldThrowArgumentE this.SkipIfEmulatorNotAvailable(); Assert.Throws(() => - new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", " ")); + new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", " ")); } [Fact] @@ -504,7 +505,7 @@ public async Task AddMessagesAsync_WithHierarchicalPartitioning_ShouldAddMessage const string UserId = "user-456"; const string SessionId = "session-789"; // Test hierarchical partitioning constructor with connection string - using var store = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, HierarchicalTestContainerId, TenantId, UserId, SessionId); + using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, TenantId, UserId, SessionId); var message = new ChatMessage(ChatRole.User, "Hello from hierarchical partitioning!"); // Act @@ -522,11 +523,11 @@ public async Task AddMessagesAsync_WithHierarchicalPartitioning_ShouldAddMessage Assert.Equal(ChatRole.User, messageList[0].Role); // Verify that the document is stored with hierarchical partitioning metadata - var directQuery = new QueryDefinition("SELECT * FROM c WHERE c.conversationId = @conversationId AND c.Type = @type") + var directQuery = new QueryDefinition("SELECT * FROM c WHERE c.conversationId = @conversationId AND c.type = @type") .WithParameter("@conversationId", SessionId) .WithParameter("@type", "ChatMessage"); - var iterator = this._setupClient!.GetDatabase(this._testDatabaseId).GetContainer(HierarchicalTestContainerId) + var iterator = this._setupClient!.GetDatabase(TestDatabaseId).GetContainer(HierarchicalTestContainerId) .GetItemQueryIterator(directQuery, requestOptions: new QueryRequestOptions { PartitionKey = new PartitionKeyBuilder().Add(TenantId).Add(UserId).Add(SessionId).Build() @@ -553,7 +554,7 @@ public async Task AddMessagesAsync_WithHierarchicalMultipleMessages_ShouldAddAll const string UserId = "user-batch"; const string SessionId = "session-batch"; // Test hierarchical partitioning constructor with connection string - using var store = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, HierarchicalTestContainerId, TenantId, UserId, SessionId); + using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, TenantId, UserId, SessionId); var messages = new[] { new ChatMessage(ChatRole.User, "First hierarchical message"), @@ -589,8 +590,8 @@ public async Task GetMessagesAsync_WithHierarchicalPartitionIsolation_ShouldIsol const string SessionId = "session-isolation"; // Different userIds create different hierarchical partitions, providing proper isolation - using var store1 = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, HierarchicalTestContainerId, TenantId, UserId1, SessionId); - using var store2 = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, HierarchicalTestContainerId, TenantId, UserId2, SessionId); + using var store1 = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, TenantId, UserId1, SessionId); + using var store2 = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, TenantId, UserId2, SessionId); // Add messages to both stores await store1.AddMessagesAsync([new ChatMessage(ChatRole.User, "Message from user 1")]); @@ -623,7 +624,7 @@ public async Task SerializeDeserialize_WithHierarchicalPartitioning_ShouldPreser const string UserId = "user-serialize"; const string SessionId = "session-serialize"; - using var originalStore = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, HierarchicalTestContainerId, TenantId, UserId, SessionId); + using var originalStore = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, TenantId, UserId, SessionId); await originalStore.AddMessagesAsync([new ChatMessage(ChatRole.User, "Test serialization message")]); // Act - Serialize the store state @@ -647,7 +648,7 @@ public async Task SerializeDeserialize_WithHierarchicalPartitioning_ShouldPreser Assert.Single(messageList); Assert.Equal("Test serialization message", messageList[0].Text); Assert.Equal(SessionId, deserializedStore.ConversationId); - Assert.Equal(this._testDatabaseId, deserializedStore.DatabaseId); + Assert.Equal(TestDatabaseId, deserializedStore.DatabaseId); Assert.Equal(HierarchicalTestContainerId, deserializedStore.ContainerId); } @@ -660,8 +661,8 @@ public async Task HierarchicalAndSimplePartitioning_ShouldCoexistAsync() const string SessionId = "coexist-session"; // Create simple store using simple partitioning container and hierarchical store using hierarchical container - using var simpleStore = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, TestContainerId, SessionId); - using var hierarchicalStore = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, HierarchicalTestContainerId, "tenant-coexist", "user-coexist", SessionId); + using var simpleStore = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, SessionId); + using var hierarchicalStore = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, "tenant-coexist", "user-coexist", SessionId); // Add messages to both await simpleStore.AddMessagesAsync([new ChatMessage(ChatRole.User, "Simple partitioning message")]); diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/CosmosCheckpointStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs similarity index 87% rename from dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/CosmosCheckpointStoreTests.cs rename to dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs index 697c8879d8..33c75545f3 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/CosmosCheckpointStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs @@ -9,7 +9,7 @@ using Microsoft.Agents.AI.Workflows.Checkpointing; using Xunit; -namespace Microsoft.Agents.AI.Workflows.UnitTests; +namespace Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests; /// /// Contains tests for . @@ -32,9 +32,10 @@ public class CosmosCheckpointStoreTests : IAsyncLifetime, IDisposable private const string EmulatorEndpoint = "https://localhost:8081"; private const string EmulatorKey = "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw=="; private const string TestContainerId = "Checkpoints"; - - // Use unique database ID per test run to avoid conflicts between parallel test executions - private readonly string _testDatabaseId = $"AgentFrameworkTests-CheckpointStore-{Guid.NewGuid():N}"; + // Use fixed database ID for preserve mode inspection +#pragma warning disable CA1802 // Use literals where appropriate + private static readonly string TestDatabaseId = "AgentFrameworkTests-CheckpointStore"; +#pragma warning restore CA1802 private string _connectionString = string.Empty; private CosmosClient? _cosmosClient; @@ -43,6 +44,18 @@ public class CosmosCheckpointStoreTests : IAsyncLifetime, IDisposable private bool _emulatorAvailable; private bool _preserveContainer; + // JsonSerializerOptions configured for .NET 9+ compatibility + private static readonly JsonSerializerOptions s_jsonOptions = CreateJsonOptions(); + + private static JsonSerializerOptions CreateJsonOptions() + { + var options = new JsonSerializerOptions(); +#if NET9_0_OR_GREATER + options.TypeInfoResolver = new System.Text.Json.Serialization.Metadata.DefaultJsonTypeInfoResolver(); +#endif + return options; + } + public async Task InitializeAsync() { // Check environment variable to determine if we should preserve containers @@ -56,7 +69,7 @@ public async Task InitializeAsync() _cosmosClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); // Test connection by attempting to create database - _database = await _cosmosClient.CreateDatabaseIfNotExistsAsync(_testDatabaseId); + _database = await _cosmosClient.CreateDatabaseIfNotExistsAsync(TestDatabaseId); _container = await _database.CreateContainerIfNotExistsAsync( TestContainerId, "/runId", @@ -121,10 +134,10 @@ public void Constructor_WithCosmosClient_SetsProperties() SkipIfEmulatorNotAvailable(); // Act - using var store = new CosmosCheckpointStore(_cosmosClient!, _testDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); // Assert - Assert.Equal(_testDatabaseId, store.DatabaseId); + Assert.Equal(TestDatabaseId, store.DatabaseId); Assert.Equal(TestContainerId, store.ContainerId); } @@ -135,10 +148,10 @@ public void Constructor_WithConnectionString_SetsProperties() SkipIfEmulatorNotAvailable(); // Act - using var store = new CosmosCheckpointStore(_connectionString, _testDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(_connectionString, TestDatabaseId, TestContainerId); // Assert - Assert.Equal(_testDatabaseId, store.DatabaseId); + Assert.Equal(TestDatabaseId, store.DatabaseId); Assert.Equal(TestContainerId, store.ContainerId); } @@ -147,7 +160,7 @@ public void Constructor_WithNullCosmosClient_ThrowsArgumentNullException() { // Act & Assert Assert.Throws(() => - new CosmosCheckpointStore((CosmosClient)null!, _testDatabaseId, TestContainerId)); + new CosmosCheckpointStore((CosmosClient)null!, TestDatabaseId, TestContainerId)); } [Fact] @@ -155,7 +168,7 @@ public void Constructor_WithNullConnectionString_ThrowsArgumentException() { // Act & Assert Assert.Throws(() => - new CosmosCheckpointStore((string)null!, _testDatabaseId, TestContainerId)); + new CosmosCheckpointStore((string)null!, TestDatabaseId, TestContainerId)); } #endregion @@ -168,9 +181,9 @@ public async Task CreateCheckpointAsync_NewCheckpoint_CreatesSuccessfully() SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, _testDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); var runId = Guid.NewGuid().ToString(); - var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test checkpoint" }); + var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test checkpoint" }, s_jsonOptions); // Act var checkpointInfo = await store.CreateCheckpointAsync(runId, checkpointValue); @@ -188,10 +201,10 @@ public async Task RetrieveCheckpointAsync_ExistingCheckpoint_ReturnsCorrectValue SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, _testDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); var runId = Guid.NewGuid().ToString(); var originalData = new { message = "Hello, World!", timestamp = DateTimeOffset.UtcNow }; - var checkpointValue = JsonSerializer.SerializeToElement(originalData); + var checkpointValue = JsonSerializer.SerializeToElement(originalData, s_jsonOptions); // Act var checkpointInfo = await store.CreateCheckpointAsync(runId, checkpointValue); @@ -209,7 +222,7 @@ public async Task RetrieveCheckpointAsync_NonExistentCheckpoint_ThrowsInvalidOpe SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, _testDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); var runId = Guid.NewGuid().ToString(); var fakeCheckpointInfo = new CheckpointInfo(runId, "nonexistent-checkpoint"); @@ -224,7 +237,7 @@ public async Task RetrieveIndexAsync_EmptyStore_ReturnsEmptyCollection() SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, _testDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); var runId = Guid.NewGuid().ToString(); // Act @@ -241,9 +254,9 @@ public async Task RetrieveIndexAsync_WithCheckpoints_ReturnsAllCheckpoints() SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, _testDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); var runId = Guid.NewGuid().ToString(); - var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }); + var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); // Create multiple checkpoints var checkpoint1 = await store.CreateCheckpointAsync(runId, checkpointValue); @@ -266,9 +279,9 @@ public async Task CreateCheckpointAsync_WithParent_CreatesHierarchy() SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, _testDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); var runId = Guid.NewGuid().ToString(); - var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }); + var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); // Act var parentCheckpoint = await store.CreateCheckpointAsync(runId, checkpointValue); @@ -286,9 +299,9 @@ public async Task RetrieveIndexAsync_WithParentFilter_ReturnsFilteredResults() SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, _testDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); var runId = Guid.NewGuid().ToString(); - var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }); + var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); // Create parent and child checkpoints var parent = await store.CreateCheckpointAsync(runId, checkpointValue); @@ -322,10 +335,10 @@ public async Task CheckpointOperations_DifferentRuns_IsolatesData() SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, _testDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); var runId1 = Guid.NewGuid().ToString(); var runId2 = Guid.NewGuid().ToString(); - var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }); + var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); // Act var checkpoint1 = await store.CreateCheckpointAsync(runId1, checkpointValue); @@ -352,8 +365,8 @@ public async Task CreateCheckpointAsync_WithNullRunId_ThrowsArgumentException() SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, _testDatabaseId, TestContainerId); - var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }); + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); // Act & Assert await Assert.ThrowsAsync(() => @@ -366,8 +379,8 @@ public async Task CreateCheckpointAsync_WithEmptyRunId_ThrowsArgumentException() SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, _testDatabaseId, TestContainerId); - var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }); + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); // Act & Assert await Assert.ThrowsAsync(() => @@ -380,7 +393,7 @@ public async Task RetrieveCheckpointAsync_WithNullCheckpointInfo_ThrowsArgumentN SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, _testDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); var runId = Guid.NewGuid().ToString(); // Act & Assert @@ -398,8 +411,8 @@ public async Task Dispose_AfterDisposal_ThrowsObjectDisposedException() SkipIfEmulatorNotAvailable(); // Arrange - var store = new CosmosCheckpointStore(_cosmosClient!, _testDatabaseId, TestContainerId); - var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }); + var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); // Act store.Dispose(); @@ -415,7 +428,7 @@ public void Dispose_MultipleCalls_DoesNotThrow() SkipIfEmulatorNotAvailable(); // Arrange - var store = new CosmosCheckpointStore(_cosmosClient!, _testDatabaseId, TestContainerId); + var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); // Act & Assert (should not throw) store.Dispose(); diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests.csproj b/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests.csproj new file mode 100644 index 0000000000..13096c44da --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests.csproj @@ -0,0 +1,23 @@ + + + + $(ProjectsTargetFrameworks) + $(NoWarn);MEAI001 + + + + false + + + + + + + + + + + + + + diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/CosmosAgentThreadTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/CosmosAgentThreadTests.cs deleted file mode 100644 index fc8dd2560d..0000000000 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/CosmosAgentThreadTests.cs +++ /dev/null @@ -1,507 +0,0 @@ -// Copyright (c) Microsoft. All rights reserved. - -using System; -using System.Linq; -using System.Text.Json; -using System.Threading.Tasks; -using Microsoft.Azure.Cosmos; -using Microsoft.Extensions.AI; -using Microsoft.Agents.AI; -using Xunit; - -namespace Microsoft.Agents.AI.Abstractions.UnitTests; - -/// -/// Collection definition to ensure Cosmos DB tests don't run in parallel -/// -[CollectionDefinition("CosmosDB", DisableParallelization = true)] -public class CosmosDBTestFixture : ICollectionFixture -{ -} - -/// -/// Contains tests for . -/// -/// These tests use the connection string constructor approach instead of manually creating CosmosClient instances, -/// making the tests simpler and more realistic. -/// -/// Test Modes: -/// - Default Mode: Cleans up all test data after each test run (deletes database) -/// - Preserve Mode: Keeps containers and data for inspection in Cosmos DB Emulator Data Explorer -/// -/// Usage Examples: -/// - Run all tests in preserve mode: $env:COSMOS_PRESERVE_CONTAINERS="true"; dotnet test tests/Microsoft.Agents.AI.Abstractions.UnitTests/ -/// - Run specific test category in preserve mode: $env:COSMOS_PRESERVE_CONTAINERS="true"; dotnet test tests/Microsoft.Agents.AI.Abstractions.UnitTests/ --filter "Category=CosmosDB" -/// - Reset to cleanup mode: $env:COSMOS_PRESERVE_CONTAINERS=""; dotnet test tests/Microsoft.Agents.AI.Abstractions.UnitTests/ -/// -[Collection("CosmosDB")] -[TestCaseOrderer("Xunit.Extensions.Ordering.TestCaseOrderer", "Xunit.Extensions.Ordering")] -public sealed class CosmosAgentThreadTests : IAsyncLifetime, IDisposable -{ - // Cosmos DB Emulator connection settings - private const string EmulatorEndpoint = "https://localhost:8081"; - private const string EmulatorKey = "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw=="; - - // Use static container names like CosmosChatMessageStoreTests for consistency - private readonly string _testDatabaseId = $"AgentFrameworkTests-Thread-{Guid.NewGuid():N}"; - private readonly string _testContainerId = "ChatMessages"; // Use the same container as CosmosChatMessageStoreTests - - private string _connectionString = string.Empty; - private bool _emulatorAvailable; - private bool _preserveContainer; - private CosmosClient? _setupClient; // Only used for test setup/cleanup - - public async Task InitializeAsync() - { - // Check environment variable to determine if we should preserve containers - // Set COSMOS_PRESERVE_CONTAINERS=true to keep containers and data for inspection - this._preserveContainer = string.Equals(Environment.GetEnvironmentVariable("COSMOS_PRESERVE_CONTAINERS"), "true", StringComparison.OrdinalIgnoreCase); - - this._connectionString = $"AccountEndpoint={EmulatorEndpoint};AccountKey={EmulatorKey}"; - - try - { - // Only create CosmosClient for test setup - the actual tests will use connection string constructors - this._setupClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); - - // Test connection and ensure database/container exist - var databaseResponse = await this._setupClient.CreateDatabaseIfNotExistsAsync(this._testDatabaseId); - var containerResponse = await databaseResponse.Database.CreateContainerIfNotExistsAsync( - this._testContainerId, - "/conversationId", - throughput: 400); - - // Verify the container is actually accessible by doing a read operation - await containerResponse.Container.ReadContainerAsync(); - - // Wait a moment for container to be fully propagated across all clients - await Task.Delay(500); - - // Verify the container is accessible from a new client instance (simulating the test scenario) - using var verificationClient = new CosmosClient(this._connectionString); - var verificationContainer = verificationClient.GetContainer(this._testDatabaseId, this._testContainerId); - await verificationContainer.ReadContainerAsync(); - - this._emulatorAvailable = true; - } - catch (Exception) - { - // Emulator not available, tests will be skipped - this._emulatorAvailable = false; - this._setupClient?.Dispose(); - this._setupClient = null; - } - } - - public async Task DisposeAsync() - { - if (this._setupClient != null && this._emulatorAvailable) - { - try - { - // Clean up test database unless preserve mode is enabled - if (!this._preserveContainer) - { - // Delete the entire test database to clean up all containers - var database = this._setupClient.GetDatabase(this._testDatabaseId); - await database.DeleteAsync(); - } - } - catch (Exception ex) - { - // Ignore cleanup errors during test teardown - Console.WriteLine($"Warning: Cleanup failed: {ex.Message}"); - } - finally - { - this._setupClient.Dispose(); - } - } - } - - /// - /// Implements IDisposable to properly dispose of the setup client. - /// - public void Dispose() - { - this._setupClient?.Dispose(); - } - - private void SkipIfEmulatorNotAvailable() - { - if (!this._emulatorAvailable) - { - Assert.True(true, "Cosmos DB Emulator not available, test skipped"); - } - } - - private async Task EnsureContainerReadyAsync() - { - if (this._emulatorAvailable) - { - try - { - // Create a TestCosmosAgentThread just to let it set up the container infrastructure - // This ensures the container is created using the exact same pattern the test will use - var setupThread = new TestCosmosAgentThread(this._connectionString, this._testDatabaseId, this._testContainerId); - - // Trigger container creation by accessing the MessageStore and trying to get messages - // This will create the database and container if they don't exist - try - { - await setupThread.MessageStore.GetMessagesAsync(); - } - catch (Microsoft.Azure.Cosmos.CosmosException ex) when (ex.StatusCode == System.Net.HttpStatusCode.NotFound) - { - // This is expected if the container is empty - it means the container was created successfully - } - } - catch (Exception ex) - { - throw new InvalidOperationException($"Container {this._testContainerId} in database {this._testDatabaseId} is not ready: {ex.Message}", ex); - } - } - } - - #region Constructor Tests - - [Fact] - [Trait("Category", "CosmosDB")] - public async Task Constructor_WithConnectionString_CreatesValidThread() - { - // Arrange & Act - Skip if emulator not available - this.SkipIfEmulatorNotAvailable(); - - // Act - var thread = new TestCosmosAgentThread(this._connectionString, this._testDatabaseId, this._testContainerId); - - // Assert - Assert.NotNull(thread); - Assert.Equal(this._testDatabaseId, thread.MessageStore.DatabaseId); - Assert.Equal(this._testContainerId, thread.MessageStore.ContainerId); - Assert.NotNull(thread.ConversationId); - } - - [Fact] - [Trait("Category", "CosmosDB")] - public async Task Constructor_WithConnectionStringAndConversationId_CreatesValidThread() - { - // Arrange & Act - Skip if emulator not available - this.SkipIfEmulatorNotAvailable(); - - const string ExpectedConversationId = "test-conversation-123"; - - // Act - var thread = new TestCosmosAgentThread(this._connectionString, this._testDatabaseId, this._testContainerId, ExpectedConversationId); - - // Assert - Assert.NotNull(thread); - Assert.Equal(ExpectedConversationId, thread.ConversationId); - Assert.Equal(this._testDatabaseId, thread.MessageStore.DatabaseId); - Assert.Equal(this._testContainerId, thread.MessageStore.ContainerId); - } - - [Fact] - [Trait("Category", "CosmosDB")] - public async Task Constructor_WithManagedIdentity_CreatesValidThread() - { - // Arrange & Act - Skip if emulator not available - this.SkipIfEmulatorNotAvailable(); - - const string TestConversationId = "test-conversation-456"; - - // Act - var thread = new TestCosmosAgentThread(EmulatorEndpoint, this._testDatabaseId, this._testContainerId, TestConversationId, useManagedIdentity: true); - - // Assert - Assert.NotNull(thread); - Assert.Equal(TestConversationId, thread.ConversationId); - Assert.Equal(this._testDatabaseId, thread.MessageStore.DatabaseId); - Assert.Equal(this._testContainerId, thread.MessageStore.ContainerId); - } - - [Fact] - [Trait("Category", "CosmosDB")] - public async Task Constructor_WithExistingMessageStore_CreatesValidThread() - { - // Arrange & Act - Skip if emulator not available - this.SkipIfEmulatorNotAvailable(); - - // Arrange - var messageStore = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, this._testContainerId); - - // Act - var thread = new TestCosmosAgentThread(messageStore); - - // Assert - Assert.NotNull(thread); - Assert.Equal(messageStore, thread.MessageStore); - Assert.NotNull(thread.ConversationId); - } - - [Fact(Skip = "Serialization requires additional JSON configuration for internal state types")] - [Trait("Category", "CosmosDB")] - public async Task Constructor_WithSerialization_RestoresCorrectly() - { - // Arrange & Act - Skip if emulator not available - this.SkipIfEmulatorNotAvailable(); - - // Arrange - var originalThread = new TestCosmosAgentThread(this._connectionString, this._testDatabaseId, this._testContainerId, "serialization-test"); - - // Add some messages to the original thread - await originalThread.MessageStore.AddMessagesAsync([ - new ChatMessage(ChatRole.User, "Test message for serialization"), - new ChatMessage(ChatRole.Assistant, "Response message") - ]); - - // Serialize the thread - var serialized = originalThread.Serialize(); - - // Factory function to recreate message store from serialized state - CosmosChatMessageStore MessageStoreFactory(JsonElement storeState, JsonSerializerOptions? options) - { - return new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, this._testContainerId); - } - - // Act - Restore from serialization - var restoredThread = new TestCosmosAgentThread(serialized, MessageStoreFactory); - - // Assert - Assert.NotNull(restoredThread); - Assert.Equal(originalThread.ConversationId, restoredThread.ConversationId); - - // Verify messages were preserved - var originalMessages = (await originalThread.MessageStore.GetMessagesAsync()).ToList(); - var restoredMessages = (await restoredThread.MessageStore.GetMessagesAsync()).ToList(); - - Assert.Equal(originalMessages.Count, restoredMessages.Count); - Assert.Equal(originalMessages.First().Contents, restoredMessages.First().Contents); - } - - #endregion - - #region Property and Delegation Tests - - [Fact] - [Trait("Category", "CosmosDB")] - public async Task MessageStore_DelegatesToUnderlyingStore() - { - // Arrange & Act - Skip if emulator not available - this.SkipIfEmulatorNotAvailable(); - - // Act - var thread = new TestCosmosAgentThread(this._connectionString, this._testDatabaseId, this._testContainerId); - - // Assert - Assert.NotNull(thread.MessageStore); - Assert.IsType(thread.MessageStore); - } - - #endregion - - #region Message Operations Tests - - [Fact] - [Trait("Category", "CosmosDB")] - public async Task AddAndGetMessagesAsync_WorksCorrectly() - { - // Arrange & Act - Skip if emulator not available - this.SkipIfEmulatorNotAvailable(); - - // Arrange - Use unique conversation ID to prevent test interference - var ConversationId = $"add-get-test-{Guid.NewGuid():N}"; - var thread = new TestCosmosAgentThread(this._connectionString, this._testDatabaseId, this._testContainerId, ConversationId); - - // Act - await thread.MessageStore.AddMessagesAsync([ - new ChatMessage(ChatRole.User, "Hello, this is a test message"), - new ChatMessage(ChatRole.Assistant, "Hello! I'm an AI assistant. How can I help you?") - ]); - - var messages = (await thread.MessageStore.GetMessagesAsync()).ToList(); - - // Assert - Assert.Equal(2, messages.Count); - Assert.Equal(ChatRole.User, messages[0].Role); - Assert.Equal("Hello, this is a test message", messages[0].Contents.FirstOrDefault()?.ToString()); - Assert.Equal(ChatRole.Assistant, messages[1].Role); - Assert.Equal("Hello! I'm an AI assistant. How can I help you?", messages[1].Contents.FirstOrDefault()?.ToString()); - } - - [Fact] - [Trait("Category", "CosmosDB")] - public async Task DirectCosmosChatMessageStore_WithWorkingContainer_WorksCorrectlyAsync() - { - // Arrange & Act - Skip if emulator not available - this.SkipIfEmulatorNotAvailable(); - - // Arrange - Use the working container name from CosmosChatMessageStoreTests - var ConversationId = $"direct-test-{Guid.NewGuid():N}"; - using var messageStore = new CosmosChatMessageStore(this._connectionString, this._testDatabaseId, "ChatMessages", ConversationId); - - // Act - await messageStore.AddMessagesAsync([ - new ChatMessage(ChatRole.User, "Hello, this is a direct test message"), - new ChatMessage(ChatRole.Assistant, "Hello! I'm responding directly.") - ]); - - var messages = (await messageStore.GetMessagesAsync()).ToList(); - - // Assert - Assert.Equal(2, messages.Count); - Assert.Equal(ChatRole.User, messages[0].Role); - Assert.Equal("Hello, this is a direct test message", messages[0].Contents.FirstOrDefault()?.ToString()); - Assert.Equal(ChatRole.Assistant, messages[1].Role); - Assert.Equal("Hello! I'm responding directly.", messages[1].Contents.FirstOrDefault()?.ToString()); - } - - [Fact] - [Trait("Category", "CosmosDB")] - public async Task GetMessageCountAsync_ReturnsCorrectCount() - { - // Arrange & Act - Skip if emulator not available - this.SkipIfEmulatorNotAvailable(); - await this.EnsureContainerReadyAsync(); - - // Arrange - Use unique conversation ID to prevent test interference - var ConversationId = $"count-test-{Guid.NewGuid():N}"; - var thread = new TestCosmosAgentThread(this._connectionString, this._testDatabaseId, this._testContainerId, ConversationId); - - // Act & Assert - Start with 0 messages - var initialCount = await thread.GetMessageCountAsync(); - Assert.Equal(0, initialCount); - - // Add some messages - await thread.MessageStore.AddMessagesAsync([ - new ChatMessage(ChatRole.User, "Message 1"), - new ChatMessage(ChatRole.User, "Message 2"), - new ChatMessage(ChatRole.Assistant, "Response") - ]); - - var finalCount = await thread.GetMessageCountAsync(); - Assert.Equal(3, finalCount); - } - - [Fact] - [Trait("Category", "CosmosDB")] - public async Task ClearMessagesAsync_RemovesAllMessages() - { - // Arrange & Act - Skip if emulator not available - this.SkipIfEmulatorNotAvailable(); - await this.EnsureContainerReadyAsync(); - - // Arrange - Use unique conversation ID to prevent test interference - var ConversationId = $"clear-test-{Guid.NewGuid():N}"; - var thread = new TestCosmosAgentThread(this._connectionString, this._testDatabaseId, this._testContainerId, ConversationId); - - // Add messages first - await thread.MessageStore.AddMessagesAsync([ - new ChatMessage(ChatRole.User, "Message to be cleared"), - new ChatMessage(ChatRole.Assistant, "Response to be cleared") - ]); - - // Verify messages exist - var countBeforeClear = await thread.GetMessageCountAsync(); - Assert.Equal(2, countBeforeClear); - - // Act - await thread.ClearMessagesAsync(); - - // Assert - var countAfterClear = await thread.GetMessageCountAsync(); - Assert.Equal(0, countAfterClear); - - var messagesAfterClear = (await thread.MessageStore.GetMessagesAsync()).ToList(); - Assert.Empty(messagesAfterClear); - } - - #endregion - - #region Conversation Isolation Tests - - [Fact] - [Trait("Category", "CosmosDB")] - public async Task DifferentConversationsAsync_AreIsolated() - { - // Arrange & Act - Skip if emulator not available - this.SkipIfEmulatorNotAvailable(); - await this.EnsureContainerReadyAsync(); - - // Arrange - Use unique conversation IDs to avoid interference from other tests - var timestamp = Guid.NewGuid().ToString("N").Substring(0, 8); - var thread1 = new TestCosmosAgentThread(this._connectionString, this._testDatabaseId, this._testContainerId, $"isolation-conv-1-{timestamp}"); - var thread2 = new TestCosmosAgentThread(this._connectionString, this._testDatabaseId, this._testContainerId, $"isolation-conv-2-{timestamp}"); - - // Act - await thread1.MessageStore.AddMessagesAsync([new ChatMessage(ChatRole.User, "Message in conversation 1")]); - await thread2.MessageStore.AddMessagesAsync([new ChatMessage(ChatRole.User, "Message in conversation 2")]); - - var messages1 = (await thread1.MessageStore.GetMessagesAsync()).ToList(); - var messages2 = (await thread2.MessageStore.GetMessagesAsync()).ToList(); - - // Assert - Assert.Single(messages1); - Assert.Single(messages2); - Assert.Equal("Message in conversation 1", messages1[0].Contents.FirstOrDefault()?.ToString()); - Assert.Equal("Message in conversation 2", messages2[0].Contents.FirstOrDefault()?.ToString()); - - // Verify that clearing one conversation doesn't affect the other - await thread1.ClearMessagesAsync(); - - var messages1AfterClear = (await thread1.MessageStore.GetMessagesAsync()).ToList(); - var messages2AfterClear = (await thread2.MessageStore.GetMessagesAsync()).ToList(); - - Assert.Empty(messages1AfterClear); - Assert.Single(messages2AfterClear); - } - - #endregion - - #region Serialization Tests - - [Fact(Skip = "Serialization requires additional JSON configuration for internal state types")] - [Trait("Category", "CosmosDB")] - public async Task SerializeAsync_PreservesStateInformation() - { - // Arrange & Act - Skip if emulator not available - this.SkipIfEmulatorNotAvailable(); - - // Arrange - var thread = new TestCosmosAgentThread(this._connectionString, this._testDatabaseId, this._testContainerId, "serialization-conversation"); - - // Add some test data - await thread.MessageStore.AddMessagesAsync([new ChatMessage(ChatRole.User, "Serialization test message")]); - - // Act - var serialized = thread.Serialize(); - - // Assert - Assert.Equal(JsonValueKind.Object, serialized.ValueKind); - Assert.True(serialized.TryGetProperty("storeState", out var storeStateProperty)); - - // The store state should contain conversation information for restoration - Assert.Equal(JsonValueKind.Object, storeStateProperty.ValueKind); - } - - #endregion -} - -// Test implementation class to access protected constructors -file sealed class TestCosmosAgentThread : CosmosAgentThread -{ - public TestCosmosAgentThread(string connectionString, string databaseId, string containerId) - : base(connectionString, databaseId, containerId) { } - - public TestCosmosAgentThread(string connectionString, string databaseId, string containerId, string conversationId) - : base(connectionString, databaseId, containerId, conversationId) { } - - public TestCosmosAgentThread(string accountEndpoint, string databaseId, string containerId, string? conversationId, bool useManagedIdentity) - : base(accountEndpoint, databaseId, containerId, conversationId, useManagedIdentity) { } - - public TestCosmosAgentThread(CosmosChatMessageStore messageStore) - : base(messageStore) { } - - public TestCosmosAgentThread(JsonElement serializedThreadState, Func messageStoreFactory, JsonSerializerOptions? jsonSerializerOptions = null) - : base(serializedThreadState, messageStoreFactory, jsonSerializerOptions) { } -} diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/Microsoft.Agents.AI.Abstractions.UnitTests.csproj b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/Microsoft.Agents.AI.Abstractions.UnitTests.csproj index 63b0bc4007..b7c5412a53 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/Microsoft.Agents.AI.Abstractions.UnitTests.csproj +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/Microsoft.Agents.AI.Abstractions.UnitTests.csproj @@ -16,8 +16,6 @@ - - diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Microsoft.Agents.AI.Workflows.UnitTests.csproj b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Microsoft.Agents.AI.Workflows.UnitTests.csproj index cb4a81e5c5..bd9bc57915 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Microsoft.Agents.AI.Workflows.UnitTests.csproj +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Microsoft.Agents.AI.Workflows.UnitTests.csproj @@ -14,8 +14,6 @@ - - From 405d0ca60200d8dfec0410e64984af1b48a38c45 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Fri, 7 Nov 2025 17:10:15 +0000 Subject: [PATCH 08/40] Merge upstream/main - resolve slnx conflicts --- dotnet/agent-framework-dotnet.slnx | 10 +++++++--- .../Microsoft.Agents.AI.Abstractions.csproj | 3 --- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/dotnet/agent-framework-dotnet.slnx b/dotnet/agent-framework-dotnet.slnx index e8f6b521d2..441b4ba92e 100644 --- a/dotnet/agent-framework-dotnet.slnx +++ b/dotnet/agent-framework-dotnet.slnx @@ -263,7 +263,8 @@ - + + @@ -280,6 +281,7 @@ + @@ -290,8 +292,10 @@ - - + + + + diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/Microsoft.Agents.AI.Abstractions.csproj b/dotnet/src/Microsoft.Agents.AI.Abstractions/Microsoft.Agents.AI.Abstractions.csproj index d58a173b3e..4add7f427c 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/Microsoft.Agents.AI.Abstractions.csproj +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/Microsoft.Agents.AI.Abstractions.csproj @@ -28,9 +28,6 @@ - - - From 7672a76b350ca2927b00788a10e9f55cceb90d58 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Fri, 7 Nov 2025 19:05:02 +0000 Subject: [PATCH 09/40] use param validation helpers --- .../CosmosChatMessageStore.cs | 189 +++--------------- .../CosmosCheckpointStore.cs | 38 +--- .../CosmosChatMessageStoreTests.cs | 4 +- .../CosmosCheckpointStoreTests.cs | 2 +- 4 files changed, 39 insertions(+), 194 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosChatMessageStore.cs index 35e75581f9..65092ff743 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosChatMessageStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosChatMessageStore.cs @@ -11,6 +11,7 @@ using Azure.Identity; using Microsoft.Azure.Cosmos; using Microsoft.Extensions.AI; +using Microsoft.Shared.Diagnostics; namespace Microsoft.Agents.AI; @@ -107,29 +108,9 @@ public CosmosChatMessageStore(string connectionString, string databaseId, string /// Thrown when any string parameter is null or whitespace. public CosmosChatMessageStore(string connectionString, string databaseId, string containerId, string conversationId) { - if (string.IsNullOrWhiteSpace(connectionString)) - { - throw new ArgumentException("Connection string cannot be null or whitespace.", nameof(connectionString)); - } - - if (string.IsNullOrWhiteSpace(databaseId)) - { - throw new ArgumentException("Database ID cannot be null or whitespace.", nameof(databaseId)); - } - - if (string.IsNullOrWhiteSpace(containerId)) - { - throw new ArgumentException("Container ID cannot be null or whitespace.", nameof(containerId)); - } - - if (string.IsNullOrWhiteSpace(conversationId)) - { - throw new ArgumentException("Conversation ID cannot be null or whitespace.", nameof(conversationId)); - } - - this._cosmosClient = new CosmosClient(connectionString); - this._container = this._cosmosClient.GetContainer(databaseId, containerId); - this._conversationId = conversationId; + this._cosmosClient = new CosmosClient(Throw.IfNullOrWhitespace(connectionString)); + this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); + this._conversationId = Throw.IfNullOrWhitespace(conversationId); this._databaseId = databaseId; this._containerId = containerId; this._ownsClient = true; @@ -167,34 +148,14 @@ public CosmosChatMessageStore(string accountEndpoint, string databaseId, string /// Thrown when any string parameter is null or whitespace. public CosmosChatMessageStore(string accountEndpoint, string databaseId, string containerId, string conversationId, bool useManagedIdentity) { - if (string.IsNullOrWhiteSpace(accountEndpoint)) - { - throw new ArgumentException("Cannot be null or whitespace", nameof(accountEndpoint)); - } - - if (string.IsNullOrWhiteSpace(databaseId)) - { - throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); - } - - if (string.IsNullOrWhiteSpace(containerId)) - { - throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); - } - - if (string.IsNullOrWhiteSpace(conversationId)) - { - throw new ArgumentException("Cannot be null or whitespace", nameof(conversationId)); - } - if (!useManagedIdentity) { throw new ArgumentException("This constructor requires useManagedIdentity to be true. Use the connection string constructor for key-based authentication.", nameof(useManagedIdentity)); } - this._cosmosClient = new CosmosClient(accountEndpoint, new DefaultAzureCredential()); - this._container = this._cosmosClient.GetContainer(databaseId, containerId); - this._conversationId = conversationId; + this._cosmosClient = new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), new DefaultAzureCredential()); + this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); + this._conversationId = Throw.IfNullOrWhitespace(conversationId); this._databaseId = databaseId; this._containerId = containerId; this._ownsClient = true; @@ -230,26 +191,11 @@ public CosmosChatMessageStore(CosmosClient cosmosClient, string databaseId, stri /// Thrown when any string parameter is null or whitespace. public CosmosChatMessageStore(CosmosClient cosmosClient, string databaseId, string containerId, string conversationId) { - this._cosmosClient = cosmosClient ?? throw new ArgumentNullException(nameof(cosmosClient)); + this._cosmosClient = Throw.IfNull(cosmosClient); this._ownsClient = false; - if (string.IsNullOrWhiteSpace(databaseId)) - { - throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); - } - - if (string.IsNullOrWhiteSpace(containerId)) - { - throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); - } - - if (string.IsNullOrWhiteSpace(conversationId)) - { - throw new ArgumentException("Cannot be null or whitespace", nameof(conversationId)); - } - - this._container = this._cosmosClient.GetContainer(databaseId, containerId); - this._conversationId = conversationId; + this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); + this._conversationId = Throw.IfNullOrWhitespace(conversationId); this._databaseId = databaseId; this._containerId = containerId; @@ -273,45 +219,15 @@ public CosmosChatMessageStore(CosmosClient cosmosClient, string databaseId, stri /// Thrown when any string parameter is null or whitespace. public CosmosChatMessageStore(string connectionString, string databaseId, string containerId, string tenantId, string userId, string sessionId) { - if (string.IsNullOrWhiteSpace(connectionString)) - { - throw new ArgumentException("Connection string cannot be null or whitespace.", nameof(connectionString)); - } - - if (string.IsNullOrWhiteSpace(databaseId)) - { - throw new ArgumentException("Database ID cannot be null or whitespace.", nameof(databaseId)); - } - - if (string.IsNullOrWhiteSpace(containerId)) - { - throw new ArgumentException("Container ID cannot be null or whitespace.", nameof(containerId)); - } - - if (string.IsNullOrWhiteSpace(tenantId)) - { - throw new ArgumentException("Tenant ID cannot be null or whitespace.", nameof(tenantId)); - } - - if (string.IsNullOrWhiteSpace(userId)) - { - throw new ArgumentException("User ID cannot be null or whitespace.", nameof(userId)); - } - - if (string.IsNullOrWhiteSpace(sessionId)) - { - throw new ArgumentException("Session ID cannot be null or whitespace.", nameof(sessionId)); - } - - this._cosmosClient = new CosmosClient(connectionString); - this._container = this._cosmosClient.GetContainer(databaseId, containerId); - this._conversationId = sessionId; // Use sessionId as conversationId for compatibility + this._cosmosClient = new CosmosClient(Throw.IfNullOrWhitespace(connectionString)); + this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); + this._conversationId = Throw.IfNullOrWhitespace(sessionId); // Use sessionId as conversationId for compatibility this._databaseId = databaseId; this._containerId = containerId; // Initialize hierarchical partitioning mode - this._tenantId = tenantId; - this._userId = userId; + this._tenantId = Throw.IfNullOrWhitespace(tenantId); + this._userId = Throw.IfNullOrWhitespace(userId); this._useHierarchicalPartitioning = true; // Use native hierarchical partition key with PartitionKeyBuilder this._partitionKey = new PartitionKeyBuilder() @@ -335,50 +251,20 @@ public CosmosChatMessageStore(string connectionString, string databaseId, string /// Thrown when any string parameter is null or whitespace. public CosmosChatMessageStore(string accountEndpoint, string databaseId, string containerId, string tenantId, string userId, string sessionId, bool useManagedIdentity) { - if (string.IsNullOrWhiteSpace(accountEndpoint)) - { - throw new ArgumentException("Cannot be null or whitespace", nameof(accountEndpoint)); - } - - if (string.IsNullOrWhiteSpace(databaseId)) - { - throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); - } - - if (string.IsNullOrWhiteSpace(containerId)) - { - throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); - } - - if (string.IsNullOrWhiteSpace(tenantId)) - { - throw new ArgumentException("Tenant ID cannot be null or whitespace.", nameof(tenantId)); - } - - if (string.IsNullOrWhiteSpace(userId)) - { - throw new ArgumentException("User ID cannot be null or whitespace.", nameof(userId)); - } - - if (string.IsNullOrWhiteSpace(sessionId)) - { - throw new ArgumentException("Session ID cannot be null or whitespace.", nameof(sessionId)); - } - if (!useManagedIdentity) { throw new ArgumentException("This constructor requires useManagedIdentity to be true. Use the connection string constructor for key-based authentication.", nameof(useManagedIdentity)); } - this._cosmosClient = new CosmosClient(accountEndpoint, new DefaultAzureCredential()); - this._container = this._cosmosClient.GetContainer(databaseId, containerId); - this._conversationId = sessionId; // Use sessionId as conversationId for compatibility + this._cosmosClient = new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), new DefaultAzureCredential()); + this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); + this._conversationId = Throw.IfNullOrWhitespace(sessionId); // Use sessionId as conversationId for compatibility this._databaseId = databaseId; this._containerId = containerId; // Initialize hierarchical partitioning mode - this._tenantId = tenantId; - this._userId = userId; + this._tenantId = Throw.IfNullOrWhitespace(tenantId); + this._userId = Throw.IfNullOrWhitespace(userId); this._useHierarchicalPartitioning = true; // Use native hierarchical partition key with PartitionKeyBuilder this._partitionKey = new PartitionKeyBuilder() @@ -401,42 +287,17 @@ public CosmosChatMessageStore(string accountEndpoint, string databaseId, string /// Thrown when any string parameter is null or whitespace. public CosmosChatMessageStore(CosmosClient cosmosClient, string databaseId, string containerId, string tenantId, string userId, string sessionId) { - this._cosmosClient = cosmosClient ?? throw new ArgumentNullException(nameof(cosmosClient)); + this._cosmosClient = Throw.IfNull(cosmosClient); this._ownsClient = false; - if (string.IsNullOrWhiteSpace(databaseId)) - { - throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); - } - - if (string.IsNullOrWhiteSpace(containerId)) - { - throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); - } - - if (string.IsNullOrWhiteSpace(tenantId)) - { - throw new ArgumentException("Tenant ID cannot be null or whitespace.", nameof(tenantId)); - } - - if (string.IsNullOrWhiteSpace(userId)) - { - throw new ArgumentException("User ID cannot be null or whitespace.", nameof(userId)); - } - - if (string.IsNullOrWhiteSpace(sessionId)) - { - throw new ArgumentException("Session ID cannot be null or whitespace.", nameof(sessionId)); - } - - this._container = this._cosmosClient.GetContainer(databaseId, containerId); - this._conversationId = sessionId; // Use sessionId as conversationId for compatibility + this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); + this._conversationId = Throw.IfNullOrWhitespace(sessionId); // Use sessionId as conversationId for compatibility this._databaseId = databaseId; this._containerId = containerId; // Initialize hierarchical partitioning mode - this._tenantId = tenantId; - this._userId = userId; + this._tenantId = Throw.IfNullOrWhitespace(tenantId); + this._userId = Throw.IfNullOrWhitespace(userId); this._useHierarchicalPartitioning = true; // Use native hierarchical partition key with PartitionKeyBuilder this._partitionKey = new PartitionKeyBuilder() diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosCheckpointStore.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosCheckpointStore.cs index 2b9541ed9a..021d551678 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosCheckpointStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosCheckpointStore.cs @@ -8,6 +8,7 @@ using System.Threading.Tasks; using Azure.Identity; using Microsoft.Azure.Cosmos; +using Microsoft.Shared.Diagnostics; using Newtonsoft.Json; using Newtonsoft.Json.Linq; @@ -36,17 +37,10 @@ public class CosmosCheckpointStore : JsonCheckpointStore, IDisposable /// Thrown when any string parameter is null or whitespace. public CosmosCheckpointStore(string connectionString, string databaseId, string containerId) { - if (string.IsNullOrWhiteSpace(connectionString)) - throw new ArgumentException("Cannot be null or whitespace", nameof(connectionString)); - if (string.IsNullOrWhiteSpace(databaseId)) - throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); - if (string.IsNullOrWhiteSpace(containerId)) - throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); - var cosmosClientOptions = new CosmosClientOptions(); - _cosmosClient = new CosmosClient(connectionString, cosmosClientOptions); - _container = _cosmosClient.GetContainer(databaseId, containerId); + _cosmosClient = new CosmosClient(Throw.IfNullOrWhitespace(connectionString), cosmosClientOptions); + _container = _cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); _ownsClient = true; } @@ -61,15 +55,10 @@ public CosmosCheckpointStore(string connectionString, string databaseId, string /// Thrown when any string parameter is null or whitespace. public CosmosCheckpointStore(string accountEndpoint, string databaseId, string containerId, bool useManagedIdentity) { - if (string.IsNullOrWhiteSpace(accountEndpoint)) - throw new ArgumentException("Cannot be null or whitespace", nameof(accountEndpoint)); - if (string.IsNullOrWhiteSpace(databaseId)) - throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); - if (string.IsNullOrWhiteSpace(containerId)) - throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); - if (!useManagedIdentity) + { throw new ArgumentException("This constructor requires useManagedIdentity to be true. Use the connection string constructor for key-based authentication.", nameof(useManagedIdentity)); + } var cosmosClientOptions = new CosmosClientOptions { @@ -79,9 +68,9 @@ public CosmosCheckpointStore(string accountEndpoint, string databaseId, string c } }; - _cosmosClient = new CosmosClient(accountEndpoint, new DefaultAzureCredential(), cosmosClientOptions); - _container = _cosmosClient.GetContainer(databaseId, containerId); - _ownsClient = true; + this._cosmosClient = new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), new DefaultAzureCredential(), cosmosClientOptions); + this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); + this._ownsClient = true; } /// @@ -94,15 +83,10 @@ public CosmosCheckpointStore(string accountEndpoint, string databaseId, string c /// Thrown when any string parameter is null or whitespace. public CosmosCheckpointStore(CosmosClient cosmosClient, string databaseId, string containerId) { - _cosmosClient = cosmosClient ?? throw new ArgumentNullException(nameof(cosmosClient)); - - if (string.IsNullOrWhiteSpace(databaseId)) - throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); - if (string.IsNullOrWhiteSpace(containerId)) - throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); + this._cosmosClient = Throw.IfNull(cosmosClient); - _container = _cosmosClient.GetContainer(databaseId, containerId); - _ownsClient = false; + this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); + this._ownsClient = false; } /// diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs index 23432f0fec..a716b9b66f 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs @@ -181,7 +181,7 @@ public void Constructor_WithConnectionStringNoConversationId_ShouldCreateInstanc public void Constructor_WithNullConnectionString_ShouldThrowArgumentException() { // Arrange & Act & Assert - Assert.Throws(() => + Assert.Throws(() => new CosmosChatMessageStore((string)null!, TestDatabaseId, TestContainerId, "test-conversation")); } @@ -469,7 +469,7 @@ public void Constructor_WithHierarchicalNullTenantId_ShouldThrowArgumentExceptio // Arrange & Act & Assert this.SkipIfEmulatorNotAvailable(); - Assert.Throws(() => + Assert.Throws(() => new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, null!, "user-456", "session-789")); } diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs index 33c75545f3..06f99ed8e3 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs @@ -167,7 +167,7 @@ public void Constructor_WithNullCosmosClient_ThrowsArgumentNullException() public void Constructor_WithNullConnectionString_ThrowsArgumentException() { // Act & Assert - Assert.Throws(() => + Assert.Throws(() => new CosmosCheckpointStore((string)null!, TestDatabaseId, TestContainerId)); } From 394b749ffe574a5d14f6c08b425b4226b71d674a Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Fri, 7 Nov 2025 19:30:22 +0000 Subject: [PATCH 10/40] Replace useManagedIdentity boolean with TokenCredential parameter --- .../CosmosChatMessageStore.cs | 35 +++++++------------ .../CosmosCheckpointStore.cs | 18 ++++------ .../CosmosDBChatExtensions.cs | 3 +- .../CosmosDBWorkflowExtensions.cs | 5 +-- .../CosmosChatMessageStoreTests.cs | 5 ++- 5 files changed, 29 insertions(+), 37 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosChatMessageStore.cs index 65092ff743..7cad7194d6 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosChatMessageStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosChatMessageStore.cs @@ -8,6 +8,7 @@ using System.Text.Json.Serialization; using System.Threading; using System.Threading.Tasks; +using Azure.Core; using Azure.Identity; using Microsoft.Azure.Cosmos; using Microsoft.Extensions.AI; @@ -123,37 +124,32 @@ public CosmosChatMessageStore(string connectionString, string databaseId, string } /// - /// Initializes a new instance of the class using DefaultAzureCredential. + /// Initializes a new instance of the class using TokenCredential for authentication. /// /// The Cosmos DB account endpoint URI. + /// The TokenCredential to use for authentication (e.g., DefaultAzureCredential, ManagedIdentityCredential). /// The identifier of the Cosmos DB database. /// The identifier of the Cosmos DB container. - /// This parameter is used to distinguish this constructor from the connection string constructor. Always pass true. /// Thrown when any required parameter is null. /// Thrown when any string parameter is null or whitespace. - public CosmosChatMessageStore(string accountEndpoint, string databaseId, string containerId, bool useManagedIdentity) - : this(accountEndpoint, databaseId, containerId, Guid.NewGuid().ToString("N"), useManagedIdentity) + public CosmosChatMessageStore(string accountEndpoint, TokenCredential tokenCredential, string databaseId, string containerId) + : this(accountEndpoint, tokenCredential, databaseId, containerId, Guid.NewGuid().ToString("N")) { } /// - /// Initializes a new instance of the class using DefaultAzureCredential. + /// Initializes a new instance of the class using a TokenCredential for authentication. /// /// The Cosmos DB account endpoint URI. + /// The TokenCredential to use for authentication (e.g., DefaultAzureCredential, ManagedIdentityCredential). /// The identifier of the Cosmos DB database. /// The identifier of the Cosmos DB container. /// The unique identifier for this conversation thread. - /// This parameter is used to distinguish this constructor from the connection string constructor. Always pass true. /// Thrown when any required parameter is null. /// Thrown when any string parameter is null or whitespace. - public CosmosChatMessageStore(string accountEndpoint, string databaseId, string containerId, string conversationId, bool useManagedIdentity) + public CosmosChatMessageStore(string accountEndpoint, TokenCredential tokenCredential, string databaseId, string containerId, string conversationId) { - if (!useManagedIdentity) - { - throw new ArgumentException("This constructor requires useManagedIdentity to be true. Use the connection string constructor for key-based authentication.", nameof(useManagedIdentity)); - } - - this._cosmosClient = new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), new DefaultAzureCredential()); + this._cosmosClient = new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), Throw.IfNull(tokenCredential)); this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); this._conversationId = Throw.IfNullOrWhitespace(conversationId); this._databaseId = databaseId; @@ -238,25 +234,20 @@ public CosmosChatMessageStore(string connectionString, string databaseId, string } /// - /// Initializes a new instance of the class using DefaultAzureCredential with hierarchical partition keys. + /// Initializes a new instance of the class using a TokenCredential for authentication with hierarchical partition keys. /// /// The Cosmos DB account endpoint URI. + /// The TokenCredential to use for authentication (e.g., DefaultAzureCredential, ManagedIdentityCredential). /// The identifier of the Cosmos DB database. /// The identifier of the Cosmos DB container. /// The tenant identifier for hierarchical partitioning. /// The user identifier for hierarchical partitioning. /// The session identifier for hierarchical partitioning. - /// This parameter is used to distinguish this constructor from the connection string constructor. Always pass true. /// Thrown when any required parameter is null. /// Thrown when any string parameter is null or whitespace. - public CosmosChatMessageStore(string accountEndpoint, string databaseId, string containerId, string tenantId, string userId, string sessionId, bool useManagedIdentity) + public CosmosChatMessageStore(string accountEndpoint, TokenCredential tokenCredential, string databaseId, string containerId, string tenantId, string userId, string sessionId) { - if (!useManagedIdentity) - { - throw new ArgumentException("This constructor requires useManagedIdentity to be true. Use the connection string constructor for key-based authentication.", nameof(useManagedIdentity)); - } - - this._cosmosClient = new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), new DefaultAzureCredential()); + this._cosmosClient = new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), Throw.IfNull(tokenCredential)); this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); this._conversationId = Throw.IfNullOrWhitespace(sessionId); // Use sessionId as conversationId for compatibility this._databaseId = databaseId; diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosCheckpointStore.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosCheckpointStore.cs index 021d551678..daae028d05 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosCheckpointStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosCheckpointStore.cs @@ -6,6 +6,7 @@ using System.Linq; using System.Text.Json; using System.Threading.Tasks; +using Azure.Core; using Azure.Identity; using Microsoft.Azure.Cosmos; using Microsoft.Shared.Diagnostics; @@ -45,21 +46,16 @@ public CosmosCheckpointStore(string connectionString, string databaseId, string } /// - /// Initializes a new instance of the class using DefaultAzureCredential. + /// Initializes a new instance of the class using a TokenCredential for authentication. /// /// The Cosmos DB account endpoint URI. + /// The TokenCredential to use for authentication (e.g., DefaultAzureCredential, ManagedIdentityCredential). /// The identifier of the Cosmos DB database. /// The identifier of the Cosmos DB container. - /// This parameter is used to distinguish this constructor from the connection string constructor. Always pass true. /// Thrown when any required parameter is null. /// Thrown when any string parameter is null or whitespace. - public CosmosCheckpointStore(string accountEndpoint, string databaseId, string containerId, bool useManagedIdentity) + public CosmosCheckpointStore(string accountEndpoint, TokenCredential tokenCredential, string databaseId, string containerId) { - if (!useManagedIdentity) - { - throw new ArgumentException("This constructor requires useManagedIdentity to be true. Use the connection string constructor for key-based authentication.", nameof(useManagedIdentity)); - } - var cosmosClientOptions = new CosmosClientOptions { SerializerOptions = new CosmosSerializationOptions @@ -68,7 +64,7 @@ public CosmosCheckpointStore(string accountEndpoint, string databaseId, string c } }; - this._cosmosClient = new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), new DefaultAzureCredential(), cosmosClientOptions); + this._cosmosClient = new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), Throw.IfNull(tokenCredential), cosmosClientOptions); this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); this._ownsClient = true; } @@ -253,8 +249,8 @@ public CosmosCheckpointStore(string connectionString, string databaseId, string } /// - public CosmosCheckpointStore(string accountEndpoint, string databaseId, string containerId, bool useManagedIdentity) - : base(accountEndpoint, databaseId, containerId, useManagedIdentity) + public CosmosCheckpointStore(string accountEndpoint, TokenCredential tokenCredential, string databaseId, string containerId) + : base(accountEndpoint, tokenCredential, databaseId, containerId) { } diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosDBChatExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosDBChatExtensions.cs index a6f942ae11..9b9c76bd34 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosDBChatExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosDBChatExtensions.cs @@ -2,6 +2,7 @@ using System; using System.Diagnostics.CodeAnalysis; +using Azure.Identity; using Microsoft.Azure.Cosmos; namespace Microsoft.Agents.AI; @@ -69,7 +70,7 @@ public static ChatClientAgentOptions WithCosmosDBMessageStoreUsingManagedIdentit if (string.IsNullOrWhiteSpace(containerId)) throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); - options.ChatMessageStoreFactory = context => new CosmosChatMessageStore(accountEndpoint, databaseId, containerId, useManagedIdentity: true); + options.ChatMessageStoreFactory = context => new CosmosChatMessageStore(accountEndpoint, new DefaultAzureCredential(), databaseId, containerId); return options; } diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosDBWorkflowExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosDBWorkflowExtensions.cs index 98bdc6efd7..cc55a7081b 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosDBWorkflowExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosDBWorkflowExtensions.cs @@ -3,6 +3,7 @@ using System; using System.Diagnostics.CodeAnalysis; using System.Text.Json; +using Azure.Identity; using Microsoft.Azure.Cosmos; using Microsoft.Agents.AI.Workflows.Checkpointing; @@ -60,7 +61,7 @@ public static CosmosCheckpointStore CreateCheckpointStoreUsingManagedIdentity( if (string.IsNullOrWhiteSpace(containerId)) throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); - return new CosmosCheckpointStore(accountEndpoint, databaseId, containerId, useManagedIdentity: true); + return new CosmosCheckpointStore(accountEndpoint, new DefaultAzureCredential(), databaseId, containerId); } /// @@ -138,7 +139,7 @@ public static CosmosCheckpointStore CreateCheckpointStoreUsingManagedIdentity if (string.IsNullOrWhiteSpace(containerId)) throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); - return new CosmosCheckpointStore(accountEndpoint, databaseId, containerId, useManagedIdentity: true); + return new CosmosCheckpointStore(accountEndpoint, new DefaultAzureCredential(), databaseId, containerId); } /// diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs index a716b9b66f..00bea31c99 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs @@ -6,6 +6,8 @@ using System.Text.Json; using System.Text.Json.Serialization.Metadata; using System.Threading.Tasks; +using Azure.Core; +using Azure.Identity; using Microsoft.Azure.Cosmos; using Microsoft.Extensions.AI; using Microsoft.Agents.AI; @@ -436,7 +438,8 @@ public void Constructor_WithHierarchicalEndpoint_ShouldCreateInstance() this.SkipIfEmulatorNotAvailable(); // Act - using var store = new CosmosChatMessageStore(EmulatorEndpoint, TestDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789", useManagedIdentity: true); + TokenCredential credential = new DefaultAzureCredential(); + using var store = new CosmosChatMessageStore(EmulatorEndpoint, credential, TestDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789"); // Assert Assert.NotNull(store); From 76f300d107ca60b4af512bc848b8da8453fd18a5 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Fri, 7 Nov 2025 20:06:10 +0000 Subject: [PATCH 11/40] Remove redundant suppressions and fix tests --- .../CosmosChatMessageStore.cs | 10 ---------- .../CosmosChatMessageStoreTests.cs | 4 ++-- .../CosmosCheckpointStoreTests.cs | 5 +++-- .../CosmosDBCollectionFixture.cs | 18 ++++++++++++++++++ 4 files changed, 23 insertions(+), 14 deletions(-) create mode 100644 dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosDBCollectionFixture.cs diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosChatMessageStore.cs index 7cad7194d6..4ca653a84c 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosChatMessageStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosChatMessageStore.cs @@ -306,8 +306,6 @@ public CosmosChatMessageStore(CosmosClient cosmosClient, string databaseId, stri /// Optional settings for customizing the JSON deserialization process. /// Thrown when is null. /// Thrown when the serialized state cannot be deserialized. - [UnconditionalSuppressMessage("Trimming", "IL2026:Members annotated with 'RequiresUnreferencedCodeAttribute' require dynamic access", Justification = "StoreState type is controlled and used for serialization")] - [UnconditionalSuppressMessage("AOT", "IL3050:Calling members annotated with 'RequiresDynamicCodeAttribute' may break functionality when AOT compiling", Justification = "StoreState type is controlled and used for serialization")] public CosmosChatMessageStore(JsonElement serializedStoreState, CosmosClient cosmosClient, JsonSerializerOptions? jsonSerializerOptions = null) { this._cosmosClient = cosmosClient ?? throw new ArgumentNullException(nameof(cosmosClient)); @@ -344,8 +342,6 @@ public CosmosChatMessageStore(JsonElement serializedStoreState, CosmosClient cos } /// - [UnconditionalSuppressMessage("Trimming", "IL2026:Members annotated with 'RequiresUnreferencedCodeAttribute' require dynamic access", Justification = "ChatMessage deserialization is controlled")] - [UnconditionalSuppressMessage("AOT", "IL3050:Calling members annotated with 'RequiresDynamicCodeAttribute' may break functionality when AOT compiling", Justification = "ChatMessage deserialization is controlled")] public override async Task> GetMessagesAsync(CancellationToken cancellationToken = default) { #pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks @@ -389,8 +385,6 @@ public override async Task> GetMessagesAsync(Cancellati } /// - [UnconditionalSuppressMessage("Trimming", "IL2026:Members annotated with 'RequiresUnreferencedCodeAttribute' require dynamic access", Justification = "ChatMessage serialization is controlled")] - [UnconditionalSuppressMessage("AOT", "IL3050:Calling members annotated with 'RequiresDynamicCodeAttribute' may break functionality when AOT compiling", Justification = "ChatMessage serialization is controlled")] public override async Task AddMessagesAsync(IEnumerable messages, CancellationToken cancellationToken = default) { if (messages is null) @@ -496,8 +490,6 @@ private async Task AddSingleMessageAsync(ChatMessage message, CancellationToken /// /// Creates a message document with enhanced metadata. /// - [UnconditionalSuppressMessage("Trimming", "IL2026:Members annotated with 'RequiresUnreferencedCodeAttribute' require dynamic access", Justification = "ChatMessage serialization is controlled")] - [UnconditionalSuppressMessage("AOT", "IL3050:Calling members annotated with 'RequiresDynamicCodeAttribute' may break functionality when AOT compiling", Justification = "ChatMessage serialization is controlled")] private CosmosMessageDocument CreateMessageDocument(ChatMessage message, long timestamp) { return new CosmosMessageDocument @@ -518,8 +510,6 @@ private CosmosMessageDocument CreateMessageDocument(ChatMessage message, long ti } /// - [UnconditionalSuppressMessage("Trimming", "IL2026:Members annotated with 'RequiresUnreferencedCodeAttribute' require dynamic access", Justification = "StoreState serialization is controlled")] - [UnconditionalSuppressMessage("AOT", "IL3050:Calling members annotated with 'RequiresDynamicCodeAttribute' may break functionality when AOT compiling", Justification = "StoreState serialization is controlled")] public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) { #pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs index 00bea31c99..9f8daf4bdb 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs @@ -48,9 +48,9 @@ public sealed class CosmosChatMessageStoreTests : IAsyncLifetime, IDisposable private const string EmulatorKey = "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw=="; private const string TestContainerId = "ChatMessages"; private const string HierarchicalTestContainerId = "HierarchicalChatMessages"; - // Use fixed database ID for preserve mode inspection + // Use unique database ID per test class instance to avoid conflicts #pragma warning disable CA1802 // Use literals where appropriate - private static readonly string TestDatabaseId = "AgentFrameworkTests-ChatStore"; + private static readonly string TestDatabaseId = $"AgentFrameworkTests-ChatStore-{Guid.NewGuid():N}"; #pragma warning restore CA1802 private string _connectionString = string.Empty; diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs index 06f99ed8e3..0df3cecb18 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs @@ -26,15 +26,16 @@ namespace Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests; /// Database: AgentFrameworkTests /// Container: Checkpoints /// +[Collection("CosmosDB")] public class CosmosCheckpointStoreTests : IAsyncLifetime, IDisposable { // Cosmos DB Emulator connection settings private const string EmulatorEndpoint = "https://localhost:8081"; private const string EmulatorKey = "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw=="; private const string TestContainerId = "Checkpoints"; - // Use fixed database ID for preserve mode inspection + // Use unique database ID per test class instance to avoid conflicts #pragma warning disable CA1802 // Use literals where appropriate - private static readonly string TestDatabaseId = "AgentFrameworkTests-CheckpointStore"; + private static readonly string TestDatabaseId = $"AgentFrameworkTests-CheckpointStore-{Guid.NewGuid():N}"; #pragma warning restore CA1802 private string _connectionString = string.Empty; diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosDBCollectionFixture.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosDBCollectionFixture.cs new file mode 100644 index 0000000000..ac64b85d2a --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosDBCollectionFixture.cs @@ -0,0 +1,18 @@ +// Copyright (c) Microsoft. All rights reserved. + +using Xunit; + +namespace Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests; + +/// +/// Defines a collection fixture for Cosmos DB tests to ensure they run sequentially. +/// This prevents race conditions and resource conflicts when tests create and delete +/// databases in the Cosmos DB Emulator. +/// +[CollectionDefinition("CosmosDB", DisableParallelization = true)] +public sealed class CosmosDBCollectionFixture +{ + // This class has no code, and is never created. Its purpose is simply + // to be the place to apply [CollectionDefinition] and all the + // ICollectionFixture<> interfaces. +} From bbbc651f6347ab42c6b302cb9049f5cb363bf931 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Fri, 7 Nov 2025 20:19:32 +0000 Subject: [PATCH 12/40] Rename project from Microsoft.Agents.AI.Abstractions.CosmosNoSql to Microsoft.Agents.AI.CosmosNoSql --- dotnet/agent-framework-dotnet.slnx | 4 ++-- .../CosmosChatMessageStore.cs | 0 .../CosmosCheckpointStore.cs | 0 .../CosmosDBChatExtensions.cs | 0 .../CosmosDBWorkflowExtensions.cs | 0 .../Microsoft.Agents.AI.CosmosNoSql.csproj} | 2 +- .../CosmosChatMessageStoreTests.cs | 8 ++++---- .../CosmosCheckpointStoreTests.cs | 2 +- .../CosmosDBCollectionFixture.cs | 2 +- .../Microsoft.Agents.AI.CosmosNoSql.UnitTests.csproj} | 2 +- 10 files changed, 10 insertions(+), 10 deletions(-) rename dotnet/src/{Microsoft.Agents.AI.Abstractions.CosmosNoSql => Microsoft.Agents.AI.CosmosNoSql}/CosmosChatMessageStore.cs (100%) rename dotnet/src/{Microsoft.Agents.AI.Abstractions.CosmosNoSql => Microsoft.Agents.AI.CosmosNoSql}/CosmosCheckpointStore.cs (100%) rename dotnet/src/{Microsoft.Agents.AI.Abstractions.CosmosNoSql => Microsoft.Agents.AI.CosmosNoSql}/CosmosDBChatExtensions.cs (100%) rename dotnet/src/{Microsoft.Agents.AI.Abstractions.CosmosNoSql => Microsoft.Agents.AI.CosmosNoSql}/CosmosDBWorkflowExtensions.cs (100%) rename dotnet/src/{Microsoft.Agents.AI.Abstractions.CosmosNoSql/Microsoft.Agents.AI.Abstractions.CosmosNoSql.csproj => Microsoft.Agents.AI.CosmosNoSql/Microsoft.Agents.AI.CosmosNoSql.csproj} (94%) rename dotnet/tests/{Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests => Microsoft.Agents.AI.CosmosNoSql.UnitTests}/CosmosChatMessageStoreTests.cs (99%) rename dotnet/tests/{Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests => Microsoft.Agents.AI.CosmosNoSql.UnitTests}/CosmosCheckpointStoreTests.cs (99%) rename dotnet/tests/{Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests => Microsoft.Agents.AI.CosmosNoSql.UnitTests}/CosmosDBCollectionFixture.cs (90%) rename dotnet/tests/{Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests.csproj => Microsoft.Agents.AI.CosmosNoSql.UnitTests/Microsoft.Agents.AI.CosmosNoSql.UnitTests.csproj} (82%) diff --git a/dotnet/agent-framework-dotnet.slnx b/dotnet/agent-framework-dotnet.slnx index ef9c64512c..777104a70f 100644 --- a/dotnet/agent-framework-dotnet.slnx +++ b/dotnet/agent-framework-dotnet.slnx @@ -282,7 +282,7 @@ - + @@ -313,7 +313,7 @@ - + diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs similarity index 100% rename from dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosChatMessageStore.cs rename to dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosCheckpointStore.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosCheckpointStore.cs similarity index 100% rename from dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosCheckpointStore.cs rename to dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosCheckpointStore.cs diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosDBChatExtensions.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBChatExtensions.cs similarity index 100% rename from dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosDBChatExtensions.cs rename to dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBChatExtensions.cs diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosDBWorkflowExtensions.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBWorkflowExtensions.cs similarity index 100% rename from dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosDBWorkflowExtensions.cs rename to dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBWorkflowExtensions.cs diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/Microsoft.Agents.AI.Abstractions.CosmosNoSql.csproj b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/Microsoft.Agents.AI.CosmosNoSql.csproj similarity index 94% rename from dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/Microsoft.Agents.AI.Abstractions.CosmosNoSql.csproj rename to dotnet/src/Microsoft.Agents.AI.CosmosNoSql/Microsoft.Agents.AI.CosmosNoSql.csproj index 7bdf5e227a..b6f38b81be 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/Microsoft.Agents.AI.Abstractions.CosmosNoSql.csproj +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/Microsoft.Agents.AI.CosmosNoSql.csproj @@ -37,6 +37,6 @@ - + diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs similarity index 99% rename from dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs rename to dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs index 9f8daf4bdb..43891f67e9 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs @@ -13,7 +13,7 @@ using Microsoft.Agents.AI; using Xunit; -namespace Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests; +namespace Microsoft.Agents.AI.CosmosNoSql.UnitTests; /// /// Contains tests for . @@ -36,9 +36,9 @@ namespace Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests; /// | COSMOS_PRESERVE_CONTAINERS | true / false | Controls whether to preserve test data after completion | /// /// Usage Examples: -/// - Run all tests in preserve mode: $env:COSMOS_PRESERVE_CONTAINERS="true"; dotnet test tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/ -/// - Run specific test category in preserve mode: $env:COSMOS_PRESERVE_CONTAINERS="true"; dotnet test tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/ --filter "Category=CosmosDB" -/// - Reset to cleanup mode: $env:COSMOS_PRESERVE_CONTAINERS=""; dotnet test tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/ +/// - Run all tests in preserve mode: $env:COSMOS_PRESERVE_CONTAINERS="true"; dotnet test tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/ +/// - Run specific test category in preserve mode: $env:COSMOS_PRESERVE_CONTAINERS="true"; dotnet test tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/ --filter "Category=CosmosDB" +/// - Reset to cleanup mode: $env:COSMOS_PRESERVE_CONTAINERS=""; dotnet test tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/ /// [Collection("CosmosDB")] public sealed class CosmosChatMessageStoreTests : IAsyncLifetime, IDisposable diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs similarity index 99% rename from dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs rename to dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs index 0df3cecb18..e3bebce494 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs @@ -9,7 +9,7 @@ using Microsoft.Agents.AI.Workflows.Checkpointing; using Xunit; -namespace Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests; +namespace Microsoft.Agents.AI.CosmosNoSql.UnitTests; /// /// Contains tests for . diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosDBCollectionFixture.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosDBCollectionFixture.cs similarity index 90% rename from dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosDBCollectionFixture.cs rename to dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosDBCollectionFixture.cs index ac64b85d2a..c2567ed2ac 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosDBCollectionFixture.cs +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosDBCollectionFixture.cs @@ -2,7 +2,7 @@ using Xunit; -namespace Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests; +namespace Microsoft.Agents.AI.CosmosNoSql.UnitTests; /// /// Defines a collection fixture for Cosmos DB tests to ensure they run sequentially. diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests.csproj b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/Microsoft.Agents.AI.CosmosNoSql.UnitTests.csproj similarity index 82% rename from dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests.csproj rename to dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/Microsoft.Agents.AI.CosmosNoSql.UnitTests.csproj index 13096c44da..d75ef79f99 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests.csproj +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/Microsoft.Agents.AI.CosmosNoSql.UnitTests.csproj @@ -10,7 +10,7 @@ - + From 2015a765fe08b76255a9bbc1ab348b3c17a82fb8 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Fri, 7 Nov 2025 20:40:56 +0000 Subject: [PATCH 13/40] Refactor constructors to use chaining pattern --- .../CosmosChatMessageStore.cs | 126 ++++++------------ 1 file changed, 39 insertions(+), 87 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs index 4ca653a84c..6078c5bc2c 100644 --- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs @@ -85,6 +85,39 @@ private static JsonSerializerOptions CreateDefaultJsonOptions() /// public string ContainerId => this._containerId; + /// + /// Internal primary constructor used by all public constructors. + /// + /// The instance to use for Cosmos DB operations. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// The unique identifier for this conversation thread. + /// Whether this instance owns the CosmosClient and should dispose it. + /// Optional tenant identifier for hierarchical partitioning. + /// Optional user identifier for hierarchical partitioning. + internal CosmosChatMessageStore(CosmosClient cosmosClient, string databaseId, string containerId, string conversationId, bool ownsClient, string? tenantId = null, string? userId = null) + { + this._cosmosClient = Throw.IfNull(cosmosClient); + this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); + this._conversationId = Throw.IfNullOrWhitespace(conversationId); + this._databaseId = databaseId; + this._containerId = containerId; + this._ownsClient = ownsClient; + + // Initialize partitioning mode + this._tenantId = tenantId; + this._userId = userId; + this._useHierarchicalPartitioning = tenantId != null && userId != null; + + this._partitionKey = this._useHierarchicalPartitioning + ? new PartitionKeyBuilder() + .Add(tenantId!) + .Add(userId!) + .Add(conversationId) + .Build() + : new PartitionKey(conversationId); + } + /// /// Initializes a new instance of the class using a connection string. /// @@ -108,19 +141,8 @@ public CosmosChatMessageStore(string connectionString, string databaseId, string /// Thrown when any required parameter is null. /// Thrown when any string parameter is null or whitespace. public CosmosChatMessageStore(string connectionString, string databaseId, string containerId, string conversationId) + : this(new CosmosClient(Throw.IfNullOrWhitespace(connectionString)), databaseId, containerId, conversationId, ownsClient: true) { - this._cosmosClient = new CosmosClient(Throw.IfNullOrWhitespace(connectionString)); - this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); - this._conversationId = Throw.IfNullOrWhitespace(conversationId); - this._databaseId = databaseId; - this._containerId = containerId; - this._ownsClient = true; - - // Initialize simple partitioning mode - this._tenantId = null; - this._userId = null; - this._useHierarchicalPartitioning = false; - this._partitionKey = new PartitionKey(conversationId); } /// @@ -148,19 +170,8 @@ public CosmosChatMessageStore(string accountEndpoint, TokenCredential tokenCrede /// Thrown when any required parameter is null. /// Thrown when any string parameter is null or whitespace. public CosmosChatMessageStore(string accountEndpoint, TokenCredential tokenCredential, string databaseId, string containerId, string conversationId) + : this(new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), Throw.IfNull(tokenCredential)), databaseId, containerId, conversationId, ownsClient: true) { - this._cosmosClient = new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), Throw.IfNull(tokenCredential)); - this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); - this._conversationId = Throw.IfNullOrWhitespace(conversationId); - this._databaseId = databaseId; - this._containerId = containerId; - this._ownsClient = true; - - // Initialize simple partitioning mode - this._tenantId = null; - this._userId = null; - this._useHierarchicalPartitioning = false; - this._partitionKey = new PartitionKey(conversationId); } /// @@ -186,20 +197,8 @@ public CosmosChatMessageStore(CosmosClient cosmosClient, string databaseId, stri /// Thrown when is null. /// Thrown when any string parameter is null or whitespace. public CosmosChatMessageStore(CosmosClient cosmosClient, string databaseId, string containerId, string conversationId) + : this(cosmosClient, databaseId, containerId, conversationId, ownsClient: false) { - this._cosmosClient = Throw.IfNull(cosmosClient); - this._ownsClient = false; - - this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); - this._conversationId = Throw.IfNullOrWhitespace(conversationId); - this._databaseId = databaseId; - this._containerId = containerId; - - // Initialize simple partitioning mode - this._tenantId = null; - this._userId = null; - this._useHierarchicalPartitioning = false; - this._partitionKey = new PartitionKey(conversationId); } /// @@ -214,23 +213,8 @@ public CosmosChatMessageStore(CosmosClient cosmosClient, string databaseId, stri /// Thrown when any required parameter is null. /// Thrown when any string parameter is null or whitespace. public CosmosChatMessageStore(string connectionString, string databaseId, string containerId, string tenantId, string userId, string sessionId) + : this(new CosmosClient(Throw.IfNullOrWhitespace(connectionString)), databaseId, containerId, Throw.IfNullOrWhitespace(sessionId), ownsClient: true, Throw.IfNullOrWhitespace(tenantId), Throw.IfNullOrWhitespace(userId)) { - this._cosmosClient = new CosmosClient(Throw.IfNullOrWhitespace(connectionString)); - this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); - this._conversationId = Throw.IfNullOrWhitespace(sessionId); // Use sessionId as conversationId for compatibility - this._databaseId = databaseId; - this._containerId = containerId; - - // Initialize hierarchical partitioning mode - this._tenantId = Throw.IfNullOrWhitespace(tenantId); - this._userId = Throw.IfNullOrWhitespace(userId); - this._useHierarchicalPartitioning = true; - // Use native hierarchical partition key with PartitionKeyBuilder - this._partitionKey = new PartitionKeyBuilder() - .Add(tenantId) - .Add(userId) - .Add(sessionId) - .Build(); } /// @@ -246,23 +230,8 @@ public CosmosChatMessageStore(string connectionString, string databaseId, string /// Thrown when any required parameter is null. /// Thrown when any string parameter is null or whitespace. public CosmosChatMessageStore(string accountEndpoint, TokenCredential tokenCredential, string databaseId, string containerId, string tenantId, string userId, string sessionId) + : this(new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), Throw.IfNull(tokenCredential)), databaseId, containerId, Throw.IfNullOrWhitespace(sessionId), ownsClient: true, Throw.IfNullOrWhitespace(tenantId), Throw.IfNullOrWhitespace(userId)) { - this._cosmosClient = new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), Throw.IfNull(tokenCredential)); - this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); - this._conversationId = Throw.IfNullOrWhitespace(sessionId); // Use sessionId as conversationId for compatibility - this._databaseId = databaseId; - this._containerId = containerId; - - // Initialize hierarchical partitioning mode - this._tenantId = Throw.IfNullOrWhitespace(tenantId); - this._userId = Throw.IfNullOrWhitespace(userId); - this._useHierarchicalPartitioning = true; - // Use native hierarchical partition key with PartitionKeyBuilder - this._partitionKey = new PartitionKeyBuilder() - .Add(tenantId) - .Add(userId) - .Add(sessionId) - .Build(); } /// @@ -277,25 +246,8 @@ public CosmosChatMessageStore(string accountEndpoint, TokenCredential tokenCrede /// Thrown when is null. /// Thrown when any string parameter is null or whitespace. public CosmosChatMessageStore(CosmosClient cosmosClient, string databaseId, string containerId, string tenantId, string userId, string sessionId) + : this(cosmosClient, databaseId, containerId, Throw.IfNullOrWhitespace(sessionId), ownsClient: false, Throw.IfNullOrWhitespace(tenantId), Throw.IfNullOrWhitespace(userId)) { - this._cosmosClient = Throw.IfNull(cosmosClient); - this._ownsClient = false; - - this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); - this._conversationId = Throw.IfNullOrWhitespace(sessionId); // Use sessionId as conversationId for compatibility - this._databaseId = databaseId; - this._containerId = containerId; - - // Initialize hierarchical partitioning mode - this._tenantId = Throw.IfNullOrWhitespace(tenantId); - this._userId = Throw.IfNullOrWhitespace(userId); - this._useHierarchicalPartitioning = true; - // Use native hierarchical partition key with PartitionKeyBuilder - this._partitionKey = new PartitionKeyBuilder() - .Add(tenantId) - .Add(userId) - .Add(sessionId) - .Build(); } /// From bea1b474e74595e5054bb2a77c22da2fe9a2ee47 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Fri, 7 Nov 2025 21:10:37 +0000 Subject: [PATCH 14/40] Reorder deserialization constructor parameters for consistency --- .../Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs | 4 ++-- .../CosmosChatMessageStoreTests.cs | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs index 6078c5bc2c..a4fbd8e0a7 100644 --- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs @@ -253,12 +253,12 @@ public CosmosChatMessageStore(CosmosClient cosmosClient, string databaseId, stri /// /// Initializes a new instance of the class from previously serialized state. /// - /// A representing the serialized state of the message store. /// The instance to use for Cosmos DB operations. + /// A representing the serialized state of the message store. /// Optional settings for customizing the JSON deserialization process. /// Thrown when is null. /// Thrown when the serialized state cannot be deserialized. - public CosmosChatMessageStore(JsonElement serializedStoreState, CosmosClient cosmosClient, JsonSerializerOptions? jsonSerializerOptions = null) + public CosmosChatMessageStore(CosmosClient cosmosClient, JsonElement serializedStoreState, JsonSerializerOptions? jsonSerializerOptions = null) { this._cosmosClient = cosmosClient ?? throw new ArgumentNullException(nameof(cosmosClient)); this._ownsClient = false; diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs index 43891f67e9..f821d088e7 100644 --- a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs @@ -639,7 +639,7 @@ public async Task SerializeDeserialize_WithHierarchicalPartitioning_ShouldPreser { TypeInfoResolver = new DefaultJsonTypeInfoResolver() }; - using var deserializedStore = new CosmosChatMessageStore(serializedState, cosmosClient, serializerOptions); + using var deserializedStore = new CosmosChatMessageStore(cosmosClient, serializedState, serializerOptions); // Wait a moment for eventual consistency await Task.Delay(100); From 721a0b093617bc8e0ae075fc9bba308e040268e3 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Fri, 7 Nov 2025 21:23:15 +0000 Subject: [PATCH 15/40] Remove database/container IDs from serialized state --- .../CosmosChatMessageStore.cs | 16 +++++++--------- .../CosmosChatMessageStoreTests.cs | 2 +- 2 files changed, 8 insertions(+), 10 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs index a4fbd8e0a7..796a213ab1 100644 --- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs @@ -255,23 +255,25 @@ public CosmosChatMessageStore(CosmosClient cosmosClient, string databaseId, stri /// /// The instance to use for Cosmos DB operations. /// A representing the serialized state of the message store. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. /// Optional settings for customizing the JSON deserialization process. /// Thrown when is null. /// Thrown when the serialized state cannot be deserialized. - public CosmosChatMessageStore(CosmosClient cosmosClient, JsonElement serializedStoreState, JsonSerializerOptions? jsonSerializerOptions = null) + public CosmosChatMessageStore(CosmosClient cosmosClient, JsonElement serializedStoreState, string databaseId, string containerId, JsonSerializerOptions? jsonSerializerOptions = null) { this._cosmosClient = cosmosClient ?? throw new ArgumentNullException(nameof(cosmosClient)); + this._databaseId = Throw.IfNullOrWhitespace(databaseId); + this._containerId = Throw.IfNullOrWhitespace(containerId); + this._container = this._cosmosClient.GetContainer(databaseId, containerId); this._ownsClient = false; if (serializedStoreState.ValueKind is JsonValueKind.Object) { var state = JsonSerializer.Deserialize(serializedStoreState, jsonSerializerOptions); - if (state?.ConversationIdentifier is { } conversationId && state.DatabaseIdentifier is { } databaseId && state.ContainerIdentifier is { } containerId) + if (state?.ConversationIdentifier is { } conversationId) { this._conversationId = conversationId; - this._databaseId = databaseId; - this._containerId = containerId; - this._container = this._cosmosClient.GetContainer(databaseId, containerId); // Initialize hierarchical partitioning if available in state this._tenantId = state.TenantId; @@ -474,8 +476,6 @@ public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptio var state = new StoreState { ConversationIdentifier = this._conversationId, - DatabaseIdentifier = this.DatabaseId, - ContainerIdentifier = this.ContainerId, TenantId = this._tenantId, UserId = this._userId, UseHierarchicalPartitioning = this._useHierarchicalPartitioning @@ -589,8 +589,6 @@ public void Dispose() private sealed class StoreState { public string ConversationIdentifier { get; set; } = string.Empty; - public string DatabaseIdentifier { get; set; } = string.Empty; - public string ContainerIdentifier { get; set; } = string.Empty; public string? TenantId { get; set; } public string? UserId { get; set; } public bool UseHierarchicalPartitioning { get; set; } diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs index f821d088e7..164225abc6 100644 --- a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs @@ -639,7 +639,7 @@ public async Task SerializeDeserialize_WithHierarchicalPartitioning_ShouldPreser { TypeInfoResolver = new DefaultJsonTypeInfoResolver() }; - using var deserializedStore = new CosmosChatMessageStore(cosmosClient, serializedState, serializerOptions); + using var deserializedStore = new CosmosChatMessageStore(cosmosClient, serializedState, TestDatabaseId, HierarchicalTestContainerId, serializerOptions); // Wait a moment for eventual consistency await Task.Delay(100); From c2aae5f39e754e7cf382b9d7c19eafe6bbd152b7 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Fri, 7 Nov 2025 21:38:03 +0000 Subject: [PATCH 16/40] Remove auto-generation of MessageId --- .../CosmosChatMessageStore.cs | 695 ++++++++++++++++++ .../CosmosChatMessageStore.cs | 4 +- .../CosmosChatMessageStoreTests.cs | 692 +++++++++++++++++ .../CosmosCheckpointStoreTests.cs | 455 ++++++++++++ 4 files changed, 1844 insertions(+), 2 deletions(-) create mode 100644 dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosChatMessageStore.cs create mode 100644 dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs create mode 100644 dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosChatMessageStore.cs new file mode 100644 index 0000000000..4ca653a84c --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosChatMessageStore.cs @@ -0,0 +1,695 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Diagnostics.CodeAnalysis; +using System.Linq; +using System.Text.Json; +using System.Text.Json.Serialization; +using System.Threading; +using System.Threading.Tasks; +using Azure.Core; +using Azure.Identity; +using Microsoft.Azure.Cosmos; +using Microsoft.Extensions.AI; +using Microsoft.Shared.Diagnostics; + +namespace Microsoft.Agents.AI; + +/// +/// Provides a Cosmos DB implementation of the abstract class. +/// +[RequiresUnreferencedCode("The CosmosChatMessageStore uses JSON serialization which is incompatible with trimming.")] +[RequiresDynamicCode("The CosmosChatMessageStore uses JSON serialization which is incompatible with NativeAOT.")] +public sealed class CosmosChatMessageStore : ChatMessageStore, IDisposable +{ + private readonly CosmosClient _cosmosClient; + private readonly Container _container; + private readonly string _conversationId; + private readonly string _databaseId; + private readonly string _containerId; + private readonly bool _ownsClient; + private bool _disposed; + + // Hierarchical partition key support + private readonly string? _tenantId; + private readonly string? _userId; + private readonly PartitionKey _partitionKey; + private readonly bool _useHierarchicalPartitioning; + + /// + /// Cached JSON serializer options for .NET 9.0 compatibility. + /// + private static readonly JsonSerializerOptions s_defaultJsonOptions = CreateDefaultJsonOptions(); + + private static JsonSerializerOptions CreateDefaultJsonOptions() + { + var options = new JsonSerializerOptions(); +#if NET9_0_OR_GREATER + // Configure TypeInfoResolver for .NET 9.0 to enable JSON serialization + options.TypeInfoResolver = new System.Text.Json.Serialization.Metadata.DefaultJsonTypeInfoResolver(); +#endif + return options; + } + + /// + /// Gets or sets the maximum number of messages to return in a single query batch. + /// Default is 100 for optimal performance. + /// + public int MaxItemCount { get; set; } = 100; + + /// + /// Gets or sets the maximum number of items per transactional batch operation. + /// Default is 100, maximum allowed by Cosmos DB is 100. + /// + public int MaxBatchSize { get; set; } = 100; + + /// + /// Gets or sets the Time-To-Live (TTL) in seconds for messages. + /// Default is 86400 seconds (24 hours). Set to null to disable TTL. + /// + public int? MessageTtlSeconds { get; set; } = 86400; + + /// + /// Gets the conversation ID associated with this message store. + /// + public string ConversationId => this._conversationId; + + /// + /// Gets the database ID associated with this message store. + /// + public string DatabaseId => this._databaseId; + + /// + /// Gets the container ID associated with this message store. + /// + public string ContainerId => this._containerId; + + /// + /// Initializes a new instance of the class using a connection string. + /// + /// The Cosmos DB connection string. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// Thrown when any required parameter is null. + /// Thrown when any string parameter is null or whitespace. + public CosmosChatMessageStore(string connectionString, string databaseId, string containerId) + : this(connectionString, databaseId, containerId, Guid.NewGuid().ToString("N")) + { + } + + /// + /// Initializes a new instance of the class using a connection string. + /// + /// The Cosmos DB connection string. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// The unique identifier for this conversation thread. + /// Thrown when any required parameter is null. + /// Thrown when any string parameter is null or whitespace. + public CosmosChatMessageStore(string connectionString, string databaseId, string containerId, string conversationId) + { + this._cosmosClient = new CosmosClient(Throw.IfNullOrWhitespace(connectionString)); + this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); + this._conversationId = Throw.IfNullOrWhitespace(conversationId); + this._databaseId = databaseId; + this._containerId = containerId; + this._ownsClient = true; + + // Initialize simple partitioning mode + this._tenantId = null; + this._userId = null; + this._useHierarchicalPartitioning = false; + this._partitionKey = new PartitionKey(conversationId); + } + + /// + /// Initializes a new instance of the class using TokenCredential for authentication. + /// + /// The Cosmos DB account endpoint URI. + /// The TokenCredential to use for authentication (e.g., DefaultAzureCredential, ManagedIdentityCredential). + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// Thrown when any required parameter is null. + /// Thrown when any string parameter is null or whitespace. + public CosmosChatMessageStore(string accountEndpoint, TokenCredential tokenCredential, string databaseId, string containerId) + : this(accountEndpoint, tokenCredential, databaseId, containerId, Guid.NewGuid().ToString("N")) + { + } + + /// + /// Initializes a new instance of the class using a TokenCredential for authentication. + /// + /// The Cosmos DB account endpoint URI. + /// The TokenCredential to use for authentication (e.g., DefaultAzureCredential, ManagedIdentityCredential). + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// The unique identifier for this conversation thread. + /// Thrown when any required parameter is null. + /// Thrown when any string parameter is null or whitespace. + public CosmosChatMessageStore(string accountEndpoint, TokenCredential tokenCredential, string databaseId, string containerId, string conversationId) + { + this._cosmosClient = new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), Throw.IfNull(tokenCredential)); + this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); + this._conversationId = Throw.IfNullOrWhitespace(conversationId); + this._databaseId = databaseId; + this._containerId = containerId; + this._ownsClient = true; + + // Initialize simple partitioning mode + this._tenantId = null; + this._userId = null; + this._useHierarchicalPartitioning = false; + this._partitionKey = new PartitionKey(conversationId); + } + + /// + /// Initializes a new instance of the class using an existing . + /// + /// The instance to use for Cosmos DB operations. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// Thrown when is null. + /// Thrown when any string parameter is null or whitespace. + public CosmosChatMessageStore(CosmosClient cosmosClient, string databaseId, string containerId) + : this(cosmosClient, databaseId, containerId, Guid.NewGuid().ToString("N")) + { + } + + /// + /// Initializes a new instance of the class using an existing . + /// + /// The instance to use for Cosmos DB operations. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// The unique identifier for this conversation thread. + /// Thrown when is null. + /// Thrown when any string parameter is null or whitespace. + public CosmosChatMessageStore(CosmosClient cosmosClient, string databaseId, string containerId, string conversationId) + { + this._cosmosClient = Throw.IfNull(cosmosClient); + this._ownsClient = false; + + this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); + this._conversationId = Throw.IfNullOrWhitespace(conversationId); + this._databaseId = databaseId; + this._containerId = containerId; + + // Initialize simple partitioning mode + this._tenantId = null; + this._userId = null; + this._useHierarchicalPartitioning = false; + this._partitionKey = new PartitionKey(conversationId); + } + + /// + /// Initializes a new instance of the class using a connection string with hierarchical partition keys. + /// + /// The Cosmos DB connection string. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// The tenant identifier for hierarchical partitioning. + /// The user identifier for hierarchical partitioning. + /// The session identifier for hierarchical partitioning. + /// Thrown when any required parameter is null. + /// Thrown when any string parameter is null or whitespace. + public CosmosChatMessageStore(string connectionString, string databaseId, string containerId, string tenantId, string userId, string sessionId) + { + this._cosmosClient = new CosmosClient(Throw.IfNullOrWhitespace(connectionString)); + this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); + this._conversationId = Throw.IfNullOrWhitespace(sessionId); // Use sessionId as conversationId for compatibility + this._databaseId = databaseId; + this._containerId = containerId; + + // Initialize hierarchical partitioning mode + this._tenantId = Throw.IfNullOrWhitespace(tenantId); + this._userId = Throw.IfNullOrWhitespace(userId); + this._useHierarchicalPartitioning = true; + // Use native hierarchical partition key with PartitionKeyBuilder + this._partitionKey = new PartitionKeyBuilder() + .Add(tenantId) + .Add(userId) + .Add(sessionId) + .Build(); + } + + /// + /// Initializes a new instance of the class using a TokenCredential for authentication with hierarchical partition keys. + /// + /// The Cosmos DB account endpoint URI. + /// The TokenCredential to use for authentication (e.g., DefaultAzureCredential, ManagedIdentityCredential). + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// The tenant identifier for hierarchical partitioning. + /// The user identifier for hierarchical partitioning. + /// The session identifier for hierarchical partitioning. + /// Thrown when any required parameter is null. + /// Thrown when any string parameter is null or whitespace. + public CosmosChatMessageStore(string accountEndpoint, TokenCredential tokenCredential, string databaseId, string containerId, string tenantId, string userId, string sessionId) + { + this._cosmosClient = new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), Throw.IfNull(tokenCredential)); + this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); + this._conversationId = Throw.IfNullOrWhitespace(sessionId); // Use sessionId as conversationId for compatibility + this._databaseId = databaseId; + this._containerId = containerId; + + // Initialize hierarchical partitioning mode + this._tenantId = Throw.IfNullOrWhitespace(tenantId); + this._userId = Throw.IfNullOrWhitespace(userId); + this._useHierarchicalPartitioning = true; + // Use native hierarchical partition key with PartitionKeyBuilder + this._partitionKey = new PartitionKeyBuilder() + .Add(tenantId) + .Add(userId) + .Add(sessionId) + .Build(); + } + + /// + /// Initializes a new instance of the class using an existing with hierarchical partition keys. + /// + /// The instance to use for Cosmos DB operations. + /// The identifier of the Cosmos DB database. + /// The identifier of the Cosmos DB container. + /// The tenant identifier for hierarchical partitioning. + /// The user identifier for hierarchical partitioning. + /// The session identifier for hierarchical partitioning. + /// Thrown when is null. + /// Thrown when any string parameter is null or whitespace. + public CosmosChatMessageStore(CosmosClient cosmosClient, string databaseId, string containerId, string tenantId, string userId, string sessionId) + { + this._cosmosClient = Throw.IfNull(cosmosClient); + this._ownsClient = false; + + this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); + this._conversationId = Throw.IfNullOrWhitespace(sessionId); // Use sessionId as conversationId for compatibility + this._databaseId = databaseId; + this._containerId = containerId; + + // Initialize hierarchical partitioning mode + this._tenantId = Throw.IfNullOrWhitespace(tenantId); + this._userId = Throw.IfNullOrWhitespace(userId); + this._useHierarchicalPartitioning = true; + // Use native hierarchical partition key with PartitionKeyBuilder + this._partitionKey = new PartitionKeyBuilder() + .Add(tenantId) + .Add(userId) + .Add(sessionId) + .Build(); + } + + /// + /// Initializes a new instance of the class from previously serialized state. + /// + /// A representing the serialized state of the message store. + /// The instance to use for Cosmos DB operations. + /// Optional settings for customizing the JSON deserialization process. + /// Thrown when is null. + /// Thrown when the serialized state cannot be deserialized. + public CosmosChatMessageStore(JsonElement serializedStoreState, CosmosClient cosmosClient, JsonSerializerOptions? jsonSerializerOptions = null) + { + this._cosmosClient = cosmosClient ?? throw new ArgumentNullException(nameof(cosmosClient)); + this._ownsClient = false; + + if (serializedStoreState.ValueKind is JsonValueKind.Object) + { + var state = JsonSerializer.Deserialize(serializedStoreState, jsonSerializerOptions); + if (state?.ConversationIdentifier is { } conversationId && state.DatabaseIdentifier is { } databaseId && state.ContainerIdentifier is { } containerId) + { + this._conversationId = conversationId; + this._databaseId = databaseId; + this._containerId = containerId; + this._container = this._cosmosClient.GetContainer(databaseId, containerId); + + // Initialize hierarchical partitioning if available in state + this._tenantId = state.TenantId; + this._userId = state.UserId; + this._useHierarchicalPartitioning = state.UseHierarchicalPartitioning; + + this._partitionKey = (this._useHierarchicalPartitioning && this._tenantId != null && this._userId != null) + ? new PartitionKeyBuilder() + .Add(this._tenantId) + .Add(this._userId) + .Add(conversationId) + .Build() + : new PartitionKey(conversationId); + + return; + } + } + + throw new ArgumentException("Invalid serialized state", nameof(serializedStoreState)); + } + + /// + public override async Task> GetMessagesAsync(CancellationToken cancellationToken = default) + { +#pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks + if (this._disposed) + { + throw new ObjectDisposedException(GetType().FullName); + } +#pragma warning restore CA1513 + + // Use type discriminator for efficient queries + var query = new QueryDefinition("SELECT * FROM c WHERE c.conversationId = @conversationId AND c.type = @type ORDER BY c.timestamp ASC") + .WithParameter("@conversationId", this._conversationId) + .WithParameter("@type", "ChatMessage"); + + var iterator = this._container.GetItemQueryIterator(query, requestOptions: new QueryRequestOptions + { + PartitionKey = this._partitionKey, + MaxItemCount = this.MaxItemCount // Configurable query performance + }); + + var messages = new List(); + + while (iterator.HasMoreResults) + { + var response = await iterator.ReadNextAsync(cancellationToken).ConfigureAwait(false); + + foreach (var document in response) + { + if (!string.IsNullOrEmpty(document.Message)) + { + var message = JsonSerializer.Deserialize(document.Message, s_defaultJsonOptions); + if (message != null) + { + messages.Add(message); + } + } + } + } + + return messages; + } + + /// + public override async Task AddMessagesAsync(IEnumerable messages, CancellationToken cancellationToken = default) + { + if (messages is null) + { + throw new ArgumentNullException(nameof(messages)); + } + +#pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks + if (this._disposed) + { + throw new ObjectDisposedException(GetType().FullName); + } +#pragma warning restore CA1513 + + var messageList = messages.ToList(); + if (messageList.Count == 0) + { + return; + } + + // Use transactional batch for atomic operations + if (messageList.Count > 1) + { + await this.AddMessagesInBatchAsync(messageList, cancellationToken).ConfigureAwait(false); + } + else + { + await this.AddSingleMessageAsync(messageList[0], cancellationToken).ConfigureAwait(false); + } + } + + /// + /// Adds multiple messages using transactional batch operations for atomicity. + /// + private async Task AddMessagesInBatchAsync(List messages, CancellationToken cancellationToken) + { + var currentTimestamp = DateTimeOffset.UtcNow.ToUnixTimeSeconds(); + + // Process messages in optimal batch sizes + for (int i = 0; i < messages.Count; i += this.MaxBatchSize) + { + var batchMessages = messages.Skip(i).Take(this.MaxBatchSize).ToList(); + await this.ExecuteBatchOperationAsync(batchMessages, currentTimestamp, cancellationToken).ConfigureAwait(false); + } + } + + /// + /// Executes a single batch operation with retry logic and enhanced error handling. + /// + private async Task ExecuteBatchOperationAsync(List messages, long timestamp, CancellationToken cancellationToken) + { + var batch = this._container.CreateTransactionalBatch(this._partitionKey); + + foreach (var message in messages) + { + var document = this.CreateMessageDocument(message, timestamp); + batch.CreateItem(document); + } + + try + { + var response = await batch.ExecuteAsync(cancellationToken).ConfigureAwait(false); + if (!response.IsSuccessStatusCode) + { + throw new InvalidOperationException($"Batch operation failed with status: {response.StatusCode}. Details: {response.ErrorMessage}"); + } + } + catch (CosmosException ex) when (ex.StatusCode == System.Net.HttpStatusCode.RequestEntityTooLarge) + { + // If batch is too large, split into smaller batches + if (messages.Count == 1) + { + // Can't split further, use single operation + await this.AddSingleMessageAsync(messages[0], cancellationToken).ConfigureAwait(false); + return; + } + + // Split the batch in half and retry + var midpoint = messages.Count / 2; + var firstHalf = messages.Take(midpoint).ToList(); + var secondHalf = messages.Skip(midpoint).ToList(); + + await this.ExecuteBatchOperationAsync(firstHalf, timestamp, cancellationToken).ConfigureAwait(false); + await this.ExecuteBatchOperationAsync(secondHalf, timestamp, cancellationToken).ConfigureAwait(false); + } + catch (CosmosException ex) when (ex.StatusCode == (System.Net.HttpStatusCode)429) // TooManyRequests + { + // Handle rate limiting with exponential backoff + await Task.Delay(TimeSpan.FromMilliseconds(ex.RetryAfter?.TotalMilliseconds ?? 1000), cancellationToken).ConfigureAwait(false); + await this.ExecuteBatchOperationAsync(messages, timestamp, cancellationToken).ConfigureAwait(false); + } + } + + /// + /// Adds a single message to the store. + /// + private async Task AddSingleMessageAsync(ChatMessage message, CancellationToken cancellationToken) + { + var document = this.CreateMessageDocument(message, DateTimeOffset.UtcNow.ToUnixTimeSeconds()); + await this._container.CreateItemAsync(document, this._partitionKey, cancellationToken: cancellationToken).ConfigureAwait(false); + } + + /// + /// Creates a message document with enhanced metadata. + /// + private CosmosMessageDocument CreateMessageDocument(ChatMessage message, long timestamp) + { + return new CosmosMessageDocument + { + Id = Guid.NewGuid().ToString(), + ConversationId = this._conversationId, + Timestamp = timestamp, + MessageId = message.MessageId ?? Guid.NewGuid().ToString(), + Role = message.Role.Value ?? "unknown", + Message = JsonSerializer.Serialize(message, s_defaultJsonOptions), + Type = "ChatMessage", // Type discriminator + Ttl = this.MessageTtlSeconds, // Configurable TTL + // Include hierarchical metadata when using hierarchical partitioning + TenantId = this._useHierarchicalPartitioning ? this._tenantId : null, + UserId = this._useHierarchicalPartitioning ? this._userId : null, + SessionId = this._useHierarchicalPartitioning ? this._conversationId : null + }; + } + + /// + public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) + { +#pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks + if (this._disposed) + { + throw new ObjectDisposedException(GetType().FullName); + } +#pragma warning restore CA1513 + + var state = new StoreState + { + ConversationIdentifier = this._conversationId, + DatabaseIdentifier = this.DatabaseId, + ContainerIdentifier = this.ContainerId, + TenantId = this._tenantId, + UserId = this._userId, + UseHierarchicalPartitioning = this._useHierarchicalPartitioning + }; + + var options = jsonSerializerOptions ?? s_defaultJsonOptions; + return JsonSerializer.SerializeToElement(state, options); + } + + /// + /// Gets the count of messages in this conversation. + /// This is an additional utility method beyond the base contract. + /// + /// The cancellation token. + /// The number of messages in the conversation. + public async Task GetMessageCountAsync(CancellationToken cancellationToken = default) + { +#pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks + if (this._disposed) + { + throw new ObjectDisposedException(GetType().FullName); + } +#pragma warning restore CA1513 + + // Efficient count query + var query = new QueryDefinition("SELECT VALUE COUNT(1) FROM c WHERE c.conversationId = @conversationId AND c.Type = @type") + .WithParameter("@conversationId", this._conversationId) + .WithParameter("@type", "ChatMessage"); + + var iterator = this._container.GetItemQueryIterator(query, requestOptions: new QueryRequestOptions + { + PartitionKey = this._partitionKey + }); + + if (iterator.HasMoreResults) + { + var response = await iterator.ReadNextAsync(cancellationToken).ConfigureAwait(false); + return response.FirstOrDefault(); + } + + return 0; + } + + /// + /// Deletes all messages in this conversation. + /// This is an additional utility method beyond the base contract. + /// + /// The cancellation token. + /// The number of messages deleted. + public async Task ClearMessagesAsync(CancellationToken cancellationToken = default) + { +#pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks + if (this._disposed) + { + throw new ObjectDisposedException(GetType().FullName); + } +#pragma warning restore CA1513 + + // Batch delete for efficiency + var query = new QueryDefinition("SELECT VALUE c.id FROM c WHERE c.conversationId = @conversationId AND c.Type = @type") + .WithParameter("@conversationId", this._conversationId) + .WithParameter("@type", "ChatMessage"); + + var iterator = this._container.GetItemQueryIterator(query, requestOptions: new QueryRequestOptions + { + PartitionKey = new PartitionKey(this._conversationId), + MaxItemCount = this.MaxItemCount + }); + + var deletedCount = 0; + var partitionKey = new PartitionKey(this._conversationId); + + while (iterator.HasMoreResults) + { + var response = await iterator.ReadNextAsync(cancellationToken).ConfigureAwait(false); + var batch = this._container.CreateTransactionalBatch(partitionKey); + var batchItemCount = 0; + + foreach (var itemId in response) + { + if (!string.IsNullOrEmpty(itemId)) + { + batch.DeleteItem(itemId); + batchItemCount++; + deletedCount++; + } + } + + if (batchItemCount > 0) + { + await batch.ExecuteAsync(cancellationToken).ConfigureAwait(false); + } + } + + return deletedCount; + } + + /// + public void Dispose() + { + if (!this._disposed) + { + if (this._ownsClient) + { + this._cosmosClient?.Dispose(); + } + this._disposed = true; + } + } + + private sealed class StoreState + { + public string ConversationIdentifier { get; set; } = string.Empty; + public string DatabaseIdentifier { get; set; } = string.Empty; + public string ContainerIdentifier { get; set; } = string.Empty; + public string? TenantId { get; set; } + public string? UserId { get; set; } + public bool UseHierarchicalPartitioning { get; set; } + } + + /// + /// Represents a document stored in Cosmos DB for chat messages. + /// + [SuppressMessage("Performance", "CA1812:Avoid uninstantiated internal classes", Justification = "Instantiated by Cosmos DB operations")] + private sealed class CosmosMessageDocument + { + [Newtonsoft.Json.JsonProperty("id")] + public string Id { get; set; } = string.Empty; + + [Newtonsoft.Json.JsonProperty("conversationId")] + public string ConversationId { get; set; } = string.Empty; + + [Newtonsoft.Json.JsonProperty("timestamp")] + public long Timestamp { get; set; } + + [Newtonsoft.Json.JsonProperty("messageId")] + public string MessageId { get; set; } = string.Empty; + + [Newtonsoft.Json.JsonProperty("role")] + public string Role { get; set; } = string.Empty; + + [Newtonsoft.Json.JsonProperty("message")] + public string Message { get; set; } = string.Empty; + + [Newtonsoft.Json.JsonProperty("type")] + public string Type { get; set; } = string.Empty; + + [Newtonsoft.Json.JsonProperty("ttl")] + public int? Ttl { get; set; } + + /// + /// Tenant ID for hierarchical partitioning scenarios (optional). + /// + [Newtonsoft.Json.JsonProperty("tenantId")] + public string? TenantId { get; set; } + + /// + /// User ID for hierarchical partitioning scenarios (optional). + /// + [Newtonsoft.Json.JsonProperty("userId")] + public string? UserId { get; set; } + + /// + /// Session ID for hierarchical partitioning scenarios (same as ConversationId for compatibility). + /// + [Newtonsoft.Json.JsonProperty("sessionId")] + public string? SessionId { get; set; } + } +} \ No newline at end of file diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs index 796a213ab1..9066ddffee 100644 --- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs @@ -451,7 +451,7 @@ private CosmosMessageDocument CreateMessageDocument(ChatMessage message, long ti Id = Guid.NewGuid().ToString(), ConversationId = this._conversationId, Timestamp = timestamp, - MessageId = message.MessageId ?? Guid.NewGuid().ToString(), + MessageId = message.MessageId, Role = message.Role.Value ?? "unknown", Message = JsonSerializer.Serialize(message, s_defaultJsonOptions), Type = "ChatMessage", // Type discriminator @@ -610,7 +610,7 @@ private sealed class CosmosMessageDocument public long Timestamp { get; set; } [Newtonsoft.Json.JsonProperty("messageId")] - public string MessageId { get; set; } = string.Empty; + public string? MessageId { get; set; } [Newtonsoft.Json.JsonProperty("role")] public string Role { get; set; } = string.Empty; diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs new file mode 100644 index 0000000000..9f8daf4bdb --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs @@ -0,0 +1,692 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Linq; +using System.Text.Json; +using System.Text.Json.Serialization.Metadata; +using System.Threading.Tasks; +using Azure.Core; +using Azure.Identity; +using Microsoft.Azure.Cosmos; +using Microsoft.Extensions.AI; +using Microsoft.Agents.AI; +using Xunit; + +namespace Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests; + +/// +/// Contains tests for . +/// +/// Test Modes: +/// - Default Mode: Cleans up all test data after each test run (deletes database) +/// - Preserve Mode: Keeps containers and data for inspection in Cosmos DB Emulator Data Explorer +/// +/// To enable Preserve Mode, set environment variable: COSMOS_PRESERVE_CONTAINERS=true +/// Example: $env:COSMOS_PRESERVE_CONTAINERS="true"; dotnet test +/// +/// In Preserve Mode, you can view the data in Cosmos DB Emulator Data Explorer at: +/// https://localhost:8081/_explorer/index.html +/// Database: AgentFrameworkTests +/// Container: ChatMessages +/// +/// Environment Variable Reference: +/// | Variable | Values | Description | +/// |----------|--------|-------------| +/// | COSMOS_PRESERVE_CONTAINERS | true / false | Controls whether to preserve test data after completion | +/// +/// Usage Examples: +/// - Run all tests in preserve mode: $env:COSMOS_PRESERVE_CONTAINERS="true"; dotnet test tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/ +/// - Run specific test category in preserve mode: $env:COSMOS_PRESERVE_CONTAINERS="true"; dotnet test tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/ --filter "Category=CosmosDB" +/// - Reset to cleanup mode: $env:COSMOS_PRESERVE_CONTAINERS=""; dotnet test tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/ +/// +[Collection("CosmosDB")] +public sealed class CosmosChatMessageStoreTests : IAsyncLifetime, IDisposable +{ + // Cosmos DB Emulator connection settings + private const string EmulatorEndpoint = "https://localhost:8081"; + private const string EmulatorKey = "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw=="; + private const string TestContainerId = "ChatMessages"; + private const string HierarchicalTestContainerId = "HierarchicalChatMessages"; + // Use unique database ID per test class instance to avoid conflicts +#pragma warning disable CA1802 // Use literals where appropriate + private static readonly string TestDatabaseId = $"AgentFrameworkTests-ChatStore-{Guid.NewGuid():N}"; +#pragma warning restore CA1802 + + private string _connectionString = string.Empty; + private bool _emulatorAvailable; + private bool _preserveContainer; + private CosmosClient? _setupClient; // Only used for test setup/cleanup + + public async Task InitializeAsync() + { + // Check environment variable to determine if we should preserve containers + // Set COSMOS_PRESERVE_CONTAINERS=true to keep containers and data for inspection + this._preserveContainer = string.Equals(Environment.GetEnvironmentVariable("COSMOS_PRESERVE_CONTAINERS"), "true", StringComparison.OrdinalIgnoreCase); + + this._connectionString = $"AccountEndpoint={EmulatorEndpoint};AccountKey={EmulatorKey}"; + + try + { + // Only create CosmosClient for test setup - the actual tests will use connection string constructors + this._setupClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); + + // Test connection by attempting to create database + var databaseResponse = await this._setupClient.CreateDatabaseIfNotExistsAsync(TestDatabaseId); + + // Create container for simple partitioning tests + await databaseResponse.Database.CreateContainerIfNotExistsAsync( + TestContainerId, + "/conversationId", + throughput: 400); + + // Create container for hierarchical partitioning tests with hierarchical partition key + var hierarchicalContainerProperties = new ContainerProperties(HierarchicalTestContainerId, new List { "/tenantId", "/userId", "/sessionId" }); + await databaseResponse.Database.CreateContainerIfNotExistsAsync( + hierarchicalContainerProperties, + throughput: 400); + + this._emulatorAvailable = true; + } + catch (Exception) + { + // Emulator not available, tests will be skipped + this._emulatorAvailable = false; + this._setupClient?.Dispose(); + this._setupClient = null; + } + } + + public async Task DisposeAsync() + { + if (this._setupClient != null && this._emulatorAvailable) + { + try + { + if (this._preserveContainer) + { + // Preserve mode: Don't delete the database/container, keep data for inspection + // This allows viewing data in the Cosmos DB Emulator Data Explorer + // No cleanup needed - data persists for debugging + } + else + { + // Clean mode: Delete the test database and all data + var database = this._setupClient.GetDatabase(TestDatabaseId); + await database.DeleteAsync(); + } + } + catch (Exception ex) + { + // Ignore cleanup errors during test teardown + Console.WriteLine($"Warning: Cleanup failed: {ex.Message}"); + } + finally + { + this._setupClient.Dispose(); + } + } + } + + public void Dispose() + { + this._setupClient?.Dispose(); + GC.SuppressFinalize(this); + } + + private void SkipIfEmulatorNotAvailable() + { + if (!this._emulatorAvailable) + { + Assert.Fail("Cosmos DB Emulator is not available. Start the emulator to run these tests."); + } + } + + #region Constructor Tests + + [Fact] + [Trait("Category", "CosmosDB")] + public void Constructor_WithConnectionString_ShouldCreateInstance() + { + // Arrange & Act + this.SkipIfEmulatorNotAvailable(); + + // Act + using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, "test-conversation"); + + // Assert + Assert.NotNull(store); + Assert.Equal("test-conversation", store.ConversationId); + Assert.Equal(TestDatabaseId, store.DatabaseId); + Assert.Equal(TestContainerId, store.ContainerId); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public void Constructor_WithConnectionStringNoConversationId_ShouldCreateInstance() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + + // Act + using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId); + + // Assert + Assert.NotNull(store); + Assert.NotNull(store.ConversationId); + Assert.Equal(TestDatabaseId, store.DatabaseId); + Assert.Equal(TestContainerId, store.ContainerId); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public void Constructor_WithNullConnectionString_ShouldThrowArgumentException() + { + // Arrange & Act & Assert + Assert.Throws(() => + new CosmosChatMessageStore((string)null!, TestDatabaseId, TestContainerId, "test-conversation")); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public void Constructor_WithEmptyConversationId_ShouldThrowArgumentException() + { + // Arrange & Act & Assert + this.SkipIfEmulatorNotAvailable(); + + Assert.Throws(() => + new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, "")); + } + + #endregion + + #region AddMessagesAsync Tests + + [Fact] + [Trait("Category", "CosmosDB")] + public async Task AddMessagesAsync_WithSingleMessage_ShouldAddMessageAsync() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + var conversationId = Guid.NewGuid().ToString(); + using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, conversationId); + var message = new ChatMessage(ChatRole.User, "Hello, world!"); + + // Act + await store.AddMessagesAsync([message]); + + // Wait a moment for eventual consistency + await Task.Delay(100); + + // Assert + var messages = await store.GetMessagesAsync(); + var messageList = messages.ToList(); + + // Simple assertion - if this fails, we know the deserialization is the issue + if (messageList.Count == 0) + { + // Let's check if we can find ANY items in the container for this conversation + var directQuery = new QueryDefinition("SELECT VALUE COUNT(1) FROM c WHERE c.conversationId = @conversationId") + .WithParameter("@conversationId", conversationId); + var countIterator = this._setupClient!.GetDatabase(TestDatabaseId).GetContainer(TestContainerId) + .GetItemQueryIterator(directQuery, requestOptions: new QueryRequestOptions + { + PartitionKey = new PartitionKey(conversationId) + }); + + var countResponse = await countIterator.ReadNextAsync(); + var count = countResponse.FirstOrDefault(); + + // Debug: Let's see what the raw query returns + var rawQuery = new QueryDefinition("SELECT * FROM c WHERE c.conversationId = @conversationId") + .WithParameter("@conversationId", conversationId); + var rawIterator = this._setupClient!.GetDatabase(TestDatabaseId).GetContainer(TestContainerId) + .GetItemQueryIterator(rawQuery, requestOptions: new QueryRequestOptions + { + PartitionKey = new PartitionKey(conversationId) + }); + + List rawResults = new List(); + while (rawIterator.HasMoreResults) + { + var rawResponse = await rawIterator.ReadNextAsync(); + rawResults.AddRange(rawResponse); + } + + string rawJson = rawResults.Count > 0 ? Newtonsoft.Json.JsonConvert.SerializeObject(rawResults[0], Newtonsoft.Json.Formatting.Indented) : "null"; + Assert.Fail($"GetMessagesAsync returned 0 messages, but direct count query found {count} items for conversation {conversationId}. Raw document: {rawJson}"); + } + + Assert.Single(messageList); + Assert.Equal("Hello, world!", messageList[0].Text); + Assert.Equal(ChatRole.User, messageList[0].Role); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public async Task AddMessagesAsync_WithMultipleMessages_ShouldAddAllMessagesAsync() + { + // Arrange + SkipIfEmulatorNotAvailable(); + var conversationId = Guid.NewGuid().ToString(); + using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, conversationId); + var messages = new[] + { + new ChatMessage(ChatRole.User, "First message"), + new ChatMessage(ChatRole.Assistant, "Second message"), + new ChatMessage(ChatRole.User, "Third message") + }; + + // Act + await store.AddMessagesAsync(messages); + + // Assert + var retrievedMessages = await store.GetMessagesAsync(); + var messageList = retrievedMessages.ToList(); + Assert.Equal(3, messageList.Count); + Assert.Equal("First message", messageList[0].Text); + Assert.Equal("Second message", messageList[1].Text); + Assert.Equal("Third message", messageList[2].Text); + } + + #endregion + + #region GetMessagesAsync Tests + + [Fact] + [Trait("Category", "CosmosDB")] + public async Task GetMessagesAsync_WithNoMessages_ShouldReturnEmptyAsync() + { + // Arrange + SkipIfEmulatorNotAvailable(); + using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, Guid.NewGuid().ToString()); + + // Act + var messages = await store.GetMessagesAsync(); + + // Assert + Assert.Empty(messages); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public async Task GetMessagesAsync_WithConversationIsolation_ShouldOnlyReturnMessagesForConversationAsync() + { + // Arrange + SkipIfEmulatorNotAvailable(); + var conversation1 = Guid.NewGuid().ToString(); + var conversation2 = Guid.NewGuid().ToString(); + + using var store1 = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, conversation1); + using var store2 = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, conversation2); + + await store1.AddMessagesAsync([new ChatMessage(ChatRole.User, "Message for conversation 1")]); + await store2.AddMessagesAsync([new ChatMessage(ChatRole.User, "Message for conversation 2")]); + + // Act + var messages1 = await store1.GetMessagesAsync(); + var messages2 = await store2.GetMessagesAsync(); + + // Assert + var messageList1 = messages1.ToList(); + var messageList2 = messages2.ToList(); + Assert.Single(messageList1); + Assert.Single(messageList2); + Assert.Equal("Message for conversation 1", messageList1[0].Text); + Assert.Equal("Message for conversation 2", messageList2[0].Text); + } + + #endregion + + #region Integration Tests + + [Fact] + [Trait("Category", "CosmosDB")] + public async Task FullWorkflow_AddAndGet_ShouldWorkCorrectlyAsync() + { + // Arrange + SkipIfEmulatorNotAvailable(); + var conversationId = $"test-conversation-{Guid.NewGuid():N}"; // Use unique conversation ID + using var originalStore = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, conversationId); + + var messages = new[] + { + new ChatMessage(ChatRole.System, "You are a helpful assistant."), + new ChatMessage(ChatRole.User, "Hello!"), + new ChatMessage(ChatRole.Assistant, "Hi there! How can I help you today?"), + new ChatMessage(ChatRole.User, "What's the weather like?"), + new ChatMessage(ChatRole.Assistant, "I'm sorry, I don't have access to current weather data.") + }; + + // Act 1: Add messages + await originalStore.AddMessagesAsync(messages); + + // Act 2: Verify messages were added + var retrievedMessages = await originalStore.GetMessagesAsync(); + var retrievedList = retrievedMessages.ToList(); + Assert.Equal(5, retrievedList.Count); + + // Act 3: Create new store instance for same conversation (test persistence) + using var newStore = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, conversationId); + var persistedMessages = await newStore.GetMessagesAsync(); + var persistedList = persistedMessages.ToList(); + + // Assert final state + Assert.Equal(5, persistedList.Count); + Assert.Equal("You are a helpful assistant.", persistedList[0].Text); + Assert.Equal("Hello!", persistedList[1].Text); + Assert.Equal("Hi there! How can I help you today?", persistedList[2].Text); + Assert.Equal("What's the weather like?", persistedList[3].Text); + Assert.Equal("I'm sorry, I don't have access to current weather data.", persistedList[4].Text); + } + + #endregion + + #region Disposal Tests + + [Fact] + [Trait("Category", "CosmosDB")] + public void Dispose_AfterUse_ShouldNotThrow() + { + // Arrange + SkipIfEmulatorNotAvailable(); + var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, Guid.NewGuid().ToString()); + + // Act & Assert + store.Dispose(); // Should not throw + } + + [Fact] + [Trait("Category", "CosmosDB")] + public void Dispose_MultipleCalls_ShouldNotThrow() + { + // Arrange + SkipIfEmulatorNotAvailable(); + var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, Guid.NewGuid().ToString()); + + // Act & Assert + store.Dispose(); // First call + store.Dispose(); // Second call - should not throw + } + + #endregion + + #region Hierarchical Partitioning Tests + + [Fact] + [Trait("Category", "CosmosDB")] + public void Constructor_WithHierarchicalConnectionString_ShouldCreateInstance() + { + // Arrange & Act + this.SkipIfEmulatorNotAvailable(); + + // Act + using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789"); + + // Assert + Assert.NotNull(store); + Assert.Equal("session-789", store.ConversationId); + Assert.Equal(TestDatabaseId, store.DatabaseId); + Assert.Equal(HierarchicalTestContainerId, store.ContainerId); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public void Constructor_WithHierarchicalEndpoint_ShouldCreateInstance() + { + // Arrange & Act + this.SkipIfEmulatorNotAvailable(); + + // Act + TokenCredential credential = new DefaultAzureCredential(); + using var store = new CosmosChatMessageStore(EmulatorEndpoint, credential, TestDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789"); + + // Assert + Assert.NotNull(store); + Assert.Equal("session-789", store.ConversationId); + Assert.Equal(TestDatabaseId, store.DatabaseId); + Assert.Equal(HierarchicalTestContainerId, store.ContainerId); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public void Constructor_WithHierarchicalCosmosClient_ShouldCreateInstance() + { + // Arrange & Act + this.SkipIfEmulatorNotAvailable(); + + using var cosmosClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); + using var store = new CosmosChatMessageStore(cosmosClient, TestDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789"); + + // Assert + Assert.NotNull(store); + Assert.Equal("session-789", store.ConversationId); + Assert.Equal(TestDatabaseId, store.DatabaseId); + Assert.Equal(HierarchicalTestContainerId, store.ContainerId); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public void Constructor_WithHierarchicalNullTenantId_ShouldThrowArgumentException() + { + // Arrange & Act & Assert + this.SkipIfEmulatorNotAvailable(); + + Assert.Throws(() => + new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, null!, "user-456", "session-789")); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public void Constructor_WithHierarchicalEmptyUserId_ShouldThrowArgumentException() + { + // Arrange & Act & Assert + this.SkipIfEmulatorNotAvailable(); + + Assert.Throws(() => + new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, "tenant-123", "", "session-789")); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public void Constructor_WithHierarchicalWhitespaceSessionId_ShouldThrowArgumentException() + { + // Arrange & Act & Assert + this.SkipIfEmulatorNotAvailable(); + + Assert.Throws(() => + new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", " ")); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public async Task AddMessagesAsync_WithHierarchicalPartitioning_ShouldAddMessageWithMetadataAsync() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + const string TenantId = "tenant-123"; + const string UserId = "user-456"; + const string SessionId = "session-789"; + // Test hierarchical partitioning constructor with connection string + using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, TenantId, UserId, SessionId); + var message = new ChatMessage(ChatRole.User, "Hello from hierarchical partitioning!"); + + // Act + await store.AddMessagesAsync([message]); + + // Wait a moment for eventual consistency + await Task.Delay(100); + + // Assert + var messages = await store.GetMessagesAsync(); + var messageList = messages.ToList(); + + Assert.Single(messageList); + Assert.Equal("Hello from hierarchical partitioning!", messageList[0].Text); + Assert.Equal(ChatRole.User, messageList[0].Role); + + // Verify that the document is stored with hierarchical partitioning metadata + var directQuery = new QueryDefinition("SELECT * FROM c WHERE c.conversationId = @conversationId AND c.type = @type") + .WithParameter("@conversationId", SessionId) + .WithParameter("@type", "ChatMessage"); + + var iterator = this._setupClient!.GetDatabase(TestDatabaseId).GetContainer(HierarchicalTestContainerId) + .GetItemQueryIterator(directQuery, requestOptions: new QueryRequestOptions + { + PartitionKey = new PartitionKeyBuilder().Add(TenantId).Add(UserId).Add(SessionId).Build() + }); + + var response = await iterator.ReadNextAsync(); + var document = response.FirstOrDefault(); + + Assert.NotNull(document); + // The document should have hierarchical metadata + Assert.Equal(SessionId, (string)document!.conversationId); + Assert.Equal(TenantId, (string)document!.tenantId); + Assert.Equal(UserId, (string)document!.userId); + Assert.Equal(SessionId, (string)document!.sessionId); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public async Task AddMessagesAsync_WithHierarchicalMultipleMessages_ShouldAddAllMessagesAsync() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + const string TenantId = "tenant-batch"; + const string UserId = "user-batch"; + const string SessionId = "session-batch"; + // Test hierarchical partitioning constructor with connection string + using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, TenantId, UserId, SessionId); + var messages = new[] + { + new ChatMessage(ChatRole.User, "First hierarchical message"), + new ChatMessage(ChatRole.Assistant, "Second hierarchical message"), + new ChatMessage(ChatRole.User, "Third hierarchical message") + }; + + // Act + await store.AddMessagesAsync(messages); + + // Wait a moment for eventual consistency + await Task.Delay(100); + + // Assert + var retrievedMessages = await store.GetMessagesAsync(); + var messageList = retrievedMessages.ToList(); + + Assert.Equal(3, messageList.Count); + Assert.Equal("First hierarchical message", messageList[0].Text); + Assert.Equal("Second hierarchical message", messageList[1].Text); + Assert.Equal("Third hierarchical message", messageList[2].Text); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public async Task GetMessagesAsync_WithHierarchicalPartitionIsolation_ShouldIsolateMessagesByUserIdAsync() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + const string TenantId = "tenant-isolation"; + const string UserId1 = "user-1"; + const string UserId2 = "user-2"; + const string SessionId = "session-isolation"; + + // Different userIds create different hierarchical partitions, providing proper isolation + using var store1 = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, TenantId, UserId1, SessionId); + using var store2 = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, TenantId, UserId2, SessionId); + + // Add messages to both stores + await store1.AddMessagesAsync([new ChatMessage(ChatRole.User, "Message from user 1")]); + await store2.AddMessagesAsync([new ChatMessage(ChatRole.User, "Message from user 2")]); + + // Wait a moment for eventual consistency + await Task.Delay(100); + + // Act & Assert + var messages1 = await store1.GetMessagesAsync(); + var messageList1 = messages1.ToList(); + + var messages2 = await store2.GetMessagesAsync(); + var messageList2 = messages2.ToList(); + + // With true hierarchical partitioning, each user sees only their own messages + Assert.Single(messageList1); + Assert.Single(messageList2); + Assert.Equal("Message from user 1", messageList1[0].Text); + Assert.Equal("Message from user 2", messageList2[0].Text); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public async Task SerializeDeserialize_WithHierarchicalPartitioning_ShouldPreserveStateAsync() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + const string TenantId = "tenant-serialize"; + const string UserId = "user-serialize"; + const string SessionId = "session-serialize"; + + using var originalStore = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, TenantId, UserId, SessionId); + await originalStore.AddMessagesAsync([new ChatMessage(ChatRole.User, "Test serialization message")]); + + // Act - Serialize the store state + var serializedState = originalStore.Serialize(); + + // Create a new store from the serialized state + using var cosmosClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); + var serializerOptions = new JsonSerializerOptions + { + TypeInfoResolver = new DefaultJsonTypeInfoResolver() + }; + using var deserializedStore = new CosmosChatMessageStore(serializedState, cosmosClient, serializerOptions); + + // Wait a moment for eventual consistency + await Task.Delay(100); + + // Assert - The deserialized store should have the same functionality + var messages = await deserializedStore.GetMessagesAsync(); + var messageList = messages.ToList(); + + Assert.Single(messageList); + Assert.Equal("Test serialization message", messageList[0].Text); + Assert.Equal(SessionId, deserializedStore.ConversationId); + Assert.Equal(TestDatabaseId, deserializedStore.DatabaseId); + Assert.Equal(HierarchicalTestContainerId, deserializedStore.ContainerId); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public async Task HierarchicalAndSimplePartitioning_ShouldCoexistAsync() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + const string SessionId = "coexist-session"; + + // Create simple store using simple partitioning container and hierarchical store using hierarchical container + using var simpleStore = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, SessionId); + using var hierarchicalStore = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, "tenant-coexist", "user-coexist", SessionId); + + // Add messages to both + await simpleStore.AddMessagesAsync([new ChatMessage(ChatRole.User, "Simple partitioning message")]); + await hierarchicalStore.AddMessagesAsync([new ChatMessage(ChatRole.User, "Hierarchical partitioning message")]); + + // Wait a moment for eventual consistency + await Task.Delay(100); + + // Act & Assert + var simpleMessages = await simpleStore.GetMessagesAsync(); + var simpleMessageList = simpleMessages.ToList(); + + var hierarchicalMessages = await hierarchicalStore.GetMessagesAsync(); + var hierarchicalMessageList = hierarchicalMessages.ToList(); + + // Each should only see its own messages since they use different containers + Assert.Single(simpleMessageList); + Assert.Single(hierarchicalMessageList); + Assert.Equal("Simple partitioning message", simpleMessageList[0].Text); + Assert.Equal("Hierarchical partitioning message", hierarchicalMessageList[0].Text); + } + + #endregion +} diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs new file mode 100644 index 0000000000..0df3cecb18 --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs @@ -0,0 +1,455 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Linq; +using System.Text.Json; +using System.Threading.Tasks; +using Microsoft.Azure.Cosmos; +using Microsoft.Agents.AI.Workflows; +using Microsoft.Agents.AI.Workflows.Checkpointing; +using Xunit; + +namespace Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests; + +/// +/// Contains tests for . +/// +/// Test Modes: +/// - Default Mode: Cleans up all test data after each test run (deletes database) +/// - Preserve Mode: Keeps containers and data for inspection in Cosmos DB Emulator Data Explorer +/// +/// To enable Preserve Mode, set environment variable: COSMOS_PRESERVE_CONTAINERS=true +/// Example: $env:COSMOS_PRESERVE_CONTAINERS="true"; dotnet test +/// +/// In Preserve Mode, you can view the data in Cosmos DB Emulator Data Explorer at: +/// https://localhost:8081/_explorer/index.html +/// Database: AgentFrameworkTests +/// Container: Checkpoints +/// +[Collection("CosmosDB")] +public class CosmosCheckpointStoreTests : IAsyncLifetime, IDisposable +{ + // Cosmos DB Emulator connection settings + private const string EmulatorEndpoint = "https://localhost:8081"; + private const string EmulatorKey = "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw=="; + private const string TestContainerId = "Checkpoints"; + // Use unique database ID per test class instance to avoid conflicts +#pragma warning disable CA1802 // Use literals where appropriate + private static readonly string TestDatabaseId = $"AgentFrameworkTests-CheckpointStore-{Guid.NewGuid():N}"; +#pragma warning restore CA1802 + + private string _connectionString = string.Empty; + private CosmosClient? _cosmosClient; + private Database? _database; + private Container? _container; + private bool _emulatorAvailable; + private bool _preserveContainer; + + // JsonSerializerOptions configured for .NET 9+ compatibility + private static readonly JsonSerializerOptions s_jsonOptions = CreateJsonOptions(); + + private static JsonSerializerOptions CreateJsonOptions() + { + var options = new JsonSerializerOptions(); +#if NET9_0_OR_GREATER + options.TypeInfoResolver = new System.Text.Json.Serialization.Metadata.DefaultJsonTypeInfoResolver(); +#endif + return options; + } + + public async Task InitializeAsync() + { + // Check environment variable to determine if we should preserve containers + // Set COSMOS_PRESERVE_CONTAINERS=true to keep containers and data for inspection + _preserveContainer = string.Equals(Environment.GetEnvironmentVariable("COSMOS_PRESERVE_CONTAINERS"), "true", StringComparison.OrdinalIgnoreCase); + + _connectionString = $"AccountEndpoint={EmulatorEndpoint};AccountKey={EmulatorKey}"; + + try + { + _cosmosClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); + + // Test connection by attempting to create database + _database = await _cosmosClient.CreateDatabaseIfNotExistsAsync(TestDatabaseId); + _container = await _database.CreateContainerIfNotExistsAsync( + TestContainerId, + "/runId", + throughput: 400); + + _emulatorAvailable = true; + } + catch (Exception ex) when (!(ex is OutOfMemoryException || ex is StackOverflowException || ex is AccessViolationException)) + { + // Emulator not available, tests will be skipped + _emulatorAvailable = false; + _cosmosClient?.Dispose(); + _cosmosClient = null; + } + } + + public async Task DisposeAsync() + { + if (_cosmosClient != null && _emulatorAvailable) + { + try + { + if (_preserveContainer) + { + // Preserve mode: Don't delete the database/container, keep data for inspection + // This allows viewing data in the Cosmos DB Emulator Data Explorer + // No cleanup needed - data persists for debugging + } + else + { + // Clean mode: Delete the test database and all data + await _database!.DeleteAsync(); + } + } + catch (Exception ex) + { + // Ignore cleanup errors, but log for diagnostics + Console.WriteLine($"[DisposeAsync] Cleanup error: {ex.Message}\n{ex.StackTrace}"); + } + finally + { + _cosmosClient.Dispose(); + } + } + } + + private void SkipIfEmulatorNotAvailable() + { + if (!_emulatorAvailable) + { + // For now let's just fail to see if tests work + Assert.Fail("Cosmos DB Emulator is not available. Start the emulator to run these tests."); + } + } + + #region Constructor Tests + + [Fact] + public void Constructor_WithCosmosClient_SetsProperties() + { + // Arrange + SkipIfEmulatorNotAvailable(); + + // Act + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + + // Assert + Assert.Equal(TestDatabaseId, store.DatabaseId); + Assert.Equal(TestContainerId, store.ContainerId); + } + + [Fact] + public void Constructor_WithConnectionString_SetsProperties() + { + // Arrange + SkipIfEmulatorNotAvailable(); + + // Act + using var store = new CosmosCheckpointStore(_connectionString, TestDatabaseId, TestContainerId); + + // Assert + Assert.Equal(TestDatabaseId, store.DatabaseId); + Assert.Equal(TestContainerId, store.ContainerId); + } + + [Fact] + public void Constructor_WithNullCosmosClient_ThrowsArgumentNullException() + { + // Act & Assert + Assert.Throws(() => + new CosmosCheckpointStore((CosmosClient)null!, TestDatabaseId, TestContainerId)); + } + + [Fact] + public void Constructor_WithNullConnectionString_ThrowsArgumentException() + { + // Act & Assert + Assert.Throws(() => + new CosmosCheckpointStore((string)null!, TestDatabaseId, TestContainerId)); + } + + #endregion + + #region Checkpoint Operations Tests + + [Fact] + public async Task CreateCheckpointAsync_NewCheckpoint_CreatesSuccessfully() + { + SkipIfEmulatorNotAvailable(); + + // Arrange + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + var runId = Guid.NewGuid().ToString(); + var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test checkpoint" }, s_jsonOptions); + + // Act + var checkpointInfo = await store.CreateCheckpointAsync(runId, checkpointValue); + + // Assert + Assert.NotNull(checkpointInfo); + Assert.Equal(runId, checkpointInfo.RunId); + Assert.NotNull(checkpointInfo.CheckpointId); + Assert.NotEmpty(checkpointInfo.CheckpointId); + } + + [Fact] + public async Task RetrieveCheckpointAsync_ExistingCheckpoint_ReturnsCorrectValue() + { + SkipIfEmulatorNotAvailable(); + + // Arrange + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + var runId = Guid.NewGuid().ToString(); + var originalData = new { message = "Hello, World!", timestamp = DateTimeOffset.UtcNow }; + var checkpointValue = JsonSerializer.SerializeToElement(originalData, s_jsonOptions); + + // Act + var checkpointInfo = await store.CreateCheckpointAsync(runId, checkpointValue); + var retrievedValue = await store.RetrieveCheckpointAsync(runId, checkpointInfo); + + // Assert + Assert.Equal(JsonValueKind.Object, retrievedValue.ValueKind); + Assert.True(retrievedValue.TryGetProperty("message", out var messageProp)); + Assert.Equal("Hello, World!", messageProp.GetString()); + } + + [Fact] + public async Task RetrieveCheckpointAsync_NonExistentCheckpoint_ThrowsInvalidOperationException() + { + SkipIfEmulatorNotAvailable(); + + // Arrange + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + var runId = Guid.NewGuid().ToString(); + var fakeCheckpointInfo = new CheckpointInfo(runId, "nonexistent-checkpoint"); + + // Act & Assert + await Assert.ThrowsAsync(() => + store.RetrieveCheckpointAsync(runId, fakeCheckpointInfo).AsTask()); + } + + [Fact] + public async Task RetrieveIndexAsync_EmptyStore_ReturnsEmptyCollection() + { + SkipIfEmulatorNotAvailable(); + + // Arrange + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + var runId = Guid.NewGuid().ToString(); + + // Act + var index = await store.RetrieveIndexAsync(runId); + + // Assert + Assert.NotNull(index); + Assert.Empty(index); + } + + [Fact] + public async Task RetrieveIndexAsync_WithCheckpoints_ReturnsAllCheckpoints() + { + SkipIfEmulatorNotAvailable(); + + // Arrange + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + var runId = Guid.NewGuid().ToString(); + var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); + + // Create multiple checkpoints + var checkpoint1 = await store.CreateCheckpointAsync(runId, checkpointValue); + var checkpoint2 = await store.CreateCheckpointAsync(runId, checkpointValue); + var checkpoint3 = await store.CreateCheckpointAsync(runId, checkpointValue); + + // Act + var index = (await store.RetrieveIndexAsync(runId)).ToList(); + + // Assert + Assert.Equal(3, index.Count); + Assert.Contains(index, c => c.CheckpointId == checkpoint1.CheckpointId); + Assert.Contains(index, c => c.CheckpointId == checkpoint2.CheckpointId); + Assert.Contains(index, c => c.CheckpointId == checkpoint3.CheckpointId); + } + + [Fact] + public async Task CreateCheckpointAsync_WithParent_CreatesHierarchy() + { + SkipIfEmulatorNotAvailable(); + + // Arrange + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + var runId = Guid.NewGuid().ToString(); + var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); + + // Act + var parentCheckpoint = await store.CreateCheckpointAsync(runId, checkpointValue); + var childCheckpoint = await store.CreateCheckpointAsync(runId, checkpointValue, parentCheckpoint); + + // Assert + Assert.NotEqual(parentCheckpoint.CheckpointId, childCheckpoint.CheckpointId); + Assert.Equal(runId, parentCheckpoint.RunId); + Assert.Equal(runId, childCheckpoint.RunId); + } + + [Fact] + public async Task RetrieveIndexAsync_WithParentFilter_ReturnsFilteredResults() + { + SkipIfEmulatorNotAvailable(); + + // Arrange + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + var runId = Guid.NewGuid().ToString(); + var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); + + // Create parent and child checkpoints + var parent = await store.CreateCheckpointAsync(runId, checkpointValue); + var child1 = await store.CreateCheckpointAsync(runId, checkpointValue, parent); + var child2 = await store.CreateCheckpointAsync(runId, checkpointValue, parent); + + // Create an orphan checkpoint + var orphan = await store.CreateCheckpointAsync(runId, checkpointValue); + + // Act + var allCheckpoints = (await store.RetrieveIndexAsync(runId)).ToList(); + var childrenOfParent = (await store.RetrieveIndexAsync(runId, parent)).ToList(); + + // Assert + Assert.Equal(4, allCheckpoints.Count); // parent + 2 children + orphan + Assert.Equal(2, childrenOfParent.Count); // only children + + Assert.Contains(childrenOfParent, c => c.CheckpointId == child1.CheckpointId); + Assert.Contains(childrenOfParent, c => c.CheckpointId == child2.CheckpointId); + Assert.DoesNotContain(childrenOfParent, c => c.CheckpointId == parent.CheckpointId); + Assert.DoesNotContain(childrenOfParent, c => c.CheckpointId == orphan.CheckpointId); + } + + #endregion + + #region Run Isolation Tests + + [Fact] + public async Task CheckpointOperations_DifferentRuns_IsolatesData() + { + SkipIfEmulatorNotAvailable(); + + // Arrange + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + var runId1 = Guid.NewGuid().ToString(); + var runId2 = Guid.NewGuid().ToString(); + var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); + + // Act + var checkpoint1 = await store.CreateCheckpointAsync(runId1, checkpointValue); + var checkpoint2 = await store.CreateCheckpointAsync(runId2, checkpointValue); + + var index1 = (await store.RetrieveIndexAsync(runId1)).ToList(); + var index2 = (await store.RetrieveIndexAsync(runId2)).ToList(); + + // Assert + Assert.Single(index1); + Assert.Single(index2); + Assert.Equal(checkpoint1.CheckpointId, index1[0].CheckpointId); + Assert.Equal(checkpoint2.CheckpointId, index2[0].CheckpointId); + Assert.NotEqual(checkpoint1.CheckpointId, checkpoint2.CheckpointId); + } + + #endregion + + #region Error Handling Tests + + [Fact] + public async Task CreateCheckpointAsync_WithNullRunId_ThrowsArgumentException() + { + SkipIfEmulatorNotAvailable(); + + // Arrange + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); + + // Act & Assert + await Assert.ThrowsAsync(() => + store.CreateCheckpointAsync(null!, checkpointValue).AsTask()); + } + + [Fact] + public async Task CreateCheckpointAsync_WithEmptyRunId_ThrowsArgumentException() + { + SkipIfEmulatorNotAvailable(); + + // Arrange + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); + + // Act & Assert + await Assert.ThrowsAsync(() => + store.CreateCheckpointAsync("", checkpointValue).AsTask()); + } + + [Fact] + public async Task RetrieveCheckpointAsync_WithNullCheckpointInfo_ThrowsArgumentNullException() + { + SkipIfEmulatorNotAvailable(); + + // Arrange + using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + var runId = Guid.NewGuid().ToString(); + + // Act & Assert + await Assert.ThrowsAsync(() => + store.RetrieveCheckpointAsync(runId, null!).AsTask()); + } + + #endregion + + #region Disposal Tests + + [Fact] + public async Task Dispose_AfterDisposal_ThrowsObjectDisposedException() + { + SkipIfEmulatorNotAvailable(); + + // Arrange + var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); + + // Act + store.Dispose(); + + // Assert + await Assert.ThrowsAsync(() => + store.CreateCheckpointAsync("test-run", checkpointValue).AsTask()); + } + + [Fact] + public void Dispose_MultipleCalls_DoesNotThrow() + { + SkipIfEmulatorNotAvailable(); + + // Arrange + var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + + // Act & Assert (should not throw) + store.Dispose(); + store.Dispose(); + store.Dispose(); + } + + #endregion + + public void Dispose() + { + Dispose(true); + GC.SuppressFinalize(this); + } + + protected virtual void Dispose(bool disposing) + { + if (disposing) + { + this._cosmosClient?.Dispose(); + } + } +} From c5ad371cb6e08180743433141cbcc6848afcb275 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Fri, 7 Nov 2025 21:48:34 +0000 Subject: [PATCH 17/40] Optimize AddMessagesAsync to avoid enumeration when possible --- .../CosmosChatMessageStore.cs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs index 9066ddffee..6d92b74727 100644 --- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs @@ -353,7 +353,7 @@ public override async Task AddMessagesAsync(IEnumerable messages, C } #pragma warning restore CA1513 - var messageList = messages.ToList(); + var messageList = messages as IReadOnlyCollection ?? messages.ToList(); if (messageList.Count == 0) { return; @@ -366,14 +366,14 @@ public override async Task AddMessagesAsync(IEnumerable messages, C } else { - await this.AddSingleMessageAsync(messageList[0], cancellationToken).ConfigureAwait(false); + await this.AddSingleMessageAsync(messageList.First(), cancellationToken).ConfigureAwait(false); } } /// /// Adds multiple messages using transactional batch operations for atomicity. /// - private async Task AddMessagesInBatchAsync(List messages, CancellationToken cancellationToken) + private async Task AddMessagesInBatchAsync(IReadOnlyCollection messages, CancellationToken cancellationToken) { var currentTimestamp = DateTimeOffset.UtcNow.ToUnixTimeSeconds(); From ba719f12d391f3970e3fb7b52bf9b29971e0aba0 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Fri, 7 Nov 2025 22:07:21 +0000 Subject: [PATCH 18/40] Add MaxMessagesToRetrieve to limit context window --- .../CosmosChatMessageStore.cs | 30 ++++++++- .../CosmosChatMessageStoreTests.cs | 67 +++++++++++++++++++ 2 files changed, 95 insertions(+), 2 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs index 6d92b74727..2823f9c770 100644 --- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs @@ -64,6 +64,13 @@ private static JsonSerializerOptions CreateDefaultJsonOptions() /// public int MaxBatchSize { get; set; } = 100; + /// + /// Gets or sets the maximum number of messages to retrieve from the store. + /// This helps prevent exceeding LLM context windows in long conversations. + /// Default is null (no limit). When set, only the most recent messages are returned. + /// + public int? MaxMessagesToRetrieve { get; set; } + /// /// Gets or sets the Time-To-Live (TTL) in seconds for messages. /// Default is 86400 seconds (24 hours). Set to null to disable TTL. @@ -305,8 +312,9 @@ public override async Task> GetMessagesAsync(Cancellati } #pragma warning restore CA1513 - // Use type discriminator for efficient queries - var query = new QueryDefinition("SELECT * FROM c WHERE c.conversationId = @conversationId AND c.type = @type ORDER BY c.timestamp ASC") + // Fetch most recent messages in descending order when limit is set, then reverse to ascending + var orderDirection = this.MaxMessagesToRetrieve.HasValue ? "DESC" : "ASC"; + var query = new QueryDefinition($"SELECT * FROM c WHERE c.conversationId = @conversationId AND c.type = @type ORDER BY c.timestamp {orderDirection}") .WithParameter("@conversationId", this._conversationId) .WithParameter("@type", "ChatMessage"); @@ -317,6 +325,7 @@ public override async Task> GetMessagesAsync(Cancellati }); var messages = new List(); + int messageCount = 0; while (iterator.HasMoreResults) { @@ -324,15 +333,32 @@ public override async Task> GetMessagesAsync(Cancellati foreach (var document in response) { + if (this.MaxMessagesToRetrieve.HasValue && messageCount >= this.MaxMessagesToRetrieve.Value) + { + break; + } + if (!string.IsNullOrEmpty(document.Message)) { var message = JsonSerializer.Deserialize(document.Message, s_defaultJsonOptions); if (message != null) { messages.Add(message); + messageCount++; } } } + + if (this.MaxMessagesToRetrieve.HasValue && messageCount >= this.MaxMessagesToRetrieve.Value) + { + break; + } + } + + // If we fetched in descending order (most recent first), reverse to ascending order + if (this.MaxMessagesToRetrieve.HasValue) + { + messages.Reverse(); } return messages; diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs index 164225abc6..8a8ba1e445 100644 --- a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs @@ -688,5 +688,72 @@ public async Task HierarchicalAndSimplePartitioning_ShouldCoexistAsync() Assert.Equal("Hierarchical partitioning message", hierarchicalMessageList[0].Text); } + [Fact] + [Trait("Category", "CosmosDB")] + public async Task MaxMessagesToRetrieve_ShouldLimitAndReturnMostRecentAsync() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + const string ConversationId = "max-messages-test"; + + using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, ConversationId); + + // Add 10 messages + var messages = new List(); + for (int i = 1; i <= 10; i++) + { + messages.Add(new ChatMessage(ChatRole.User, $"Message {i}")); + await Task.Delay(10); // Small delay to ensure different timestamps + } + await store.AddMessagesAsync(messages); + + // Wait for eventual consistency + await Task.Delay(100); + + // Act - Set max to 5 and retrieve + store.MaxMessagesToRetrieve = 5; + var retrievedMessages = await store.GetMessagesAsync(); + var messageList = retrievedMessages.ToList(); + + // Assert - Should get the 5 most recent messages (6-10) in ascending order + Assert.Equal(5, messageList.Count); + Assert.Equal("Message 6", messageList[0].Text); + Assert.Equal("Message 7", messageList[1].Text); + Assert.Equal("Message 8", messageList[2].Text); + Assert.Equal("Message 9", messageList[3].Text); + Assert.Equal("Message 10", messageList[4].Text); + } + + [Fact] + [Trait("Category", "CosmosDB")] + public async Task MaxMessagesToRetrieve_Null_ShouldReturnAllMessagesAsync() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + const string ConversationId = "max-messages-null-test"; + + using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, ConversationId); + + // Add 10 messages + var messages = new List(); + for (int i = 1; i <= 10; i++) + { + messages.Add(new ChatMessage(ChatRole.User, $"Message {i}")); + } + await store.AddMessagesAsync(messages); + + // Wait for eventual consistency + await Task.Delay(100); + + // Act - No limit set (default null) + var retrievedMessages = await store.GetMessagesAsync(); + var messageList = retrievedMessages.ToList(); + + // Assert - Should get all 10 messages + Assert.Equal(10, messageList.Count); + Assert.Equal("Message 1", messageList[0].Text); + Assert.Equal("Message 10", messageList[9].Text); + } + #endregion } From ba431e540dd48deafe58a1ffbd32c509355f0176 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Fri, 7 Nov 2025 22:14:10 +0000 Subject: [PATCH 19/40] Make Role nullable instead of defaulting --- .../Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs index 2823f9c770..7c61007c29 100644 --- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs @@ -478,7 +478,7 @@ private CosmosMessageDocument CreateMessageDocument(ChatMessage message, long ti ConversationId = this._conversationId, Timestamp = timestamp, MessageId = message.MessageId, - Role = message.Role.Value ?? "unknown", + Role = message.Role.Value, Message = JsonSerializer.Serialize(message, s_defaultJsonOptions), Type = "ChatMessage", // Type discriminator Ttl = this.MessageTtlSeconds, // Configurable TTL @@ -639,7 +639,7 @@ private sealed class CosmosMessageDocument public string? MessageId { get; set; } [Newtonsoft.Json.JsonProperty("role")] - public string Role { get; set; } = string.Empty; + public string? Role { get; set; } [Newtonsoft.Json.JsonProperty("message")] public string Message { get; set; } = string.Empty; From 03e76211d30e11f2f23b31fe2062105c465efa6a Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Fri, 7 Nov 2025 23:17:13 +0000 Subject: [PATCH 20/40] Fix net472 build without rebasing 19 commits --- .../Extensions/AgentProviderExtensions.cs | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Extensions/AgentProviderExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Extensions/AgentProviderExtensions.cs index 5b6bbbc297..e594edd0f0 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Extensions/AgentProviderExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Extensions/AgentProviderExtensions.cs @@ -4,20 +4,24 @@ using System.Linq; using System.Threading; using System.Threading.Tasks; +#if NET9_0_OR_GREATER using Azure.AI.Agents.Persistent; +#endif using Microsoft.Extensions.AI; namespace Microsoft.Agents.AI.Workflows.Declarative.Extensions; internal static class AgentProviderExtensions { - private static readonly HashSet s_failureStatus = +#if NET9_0_OR_GREATER + private static readonly HashSet s_failureStatus = [ - Azure.AI.Agents.Persistent.RunStatus.Failed, - Azure.AI.Agents.Persistent.RunStatus.Cancelled, - Azure.AI.Agents.Persistent.RunStatus.Cancelling, - Azure.AI.Agents.Persistent.RunStatus.Expired, + global::Azure.AI.Agents.Persistent.RunStatus.Failed, + global::Azure.AI.Agents.Persistent.RunStatus.Cancelled, + global::Azure.AI.Agents.Persistent.RunStatus.Cancelling, + global::Azure.AI.Agents.Persistent.RunStatus.Expired, ]; +#endif public static async ValueTask InvokeAgentAsync( this WorkflowAgentProvider agentProvider, @@ -60,12 +64,14 @@ inputMessages is not null ? updates.Add(update); +#if NET9_0_OR_GREATER if (update.RawRepresentation is ChatResponseUpdate chatUpdate && chatUpdate.RawRepresentation is RunUpdate runUpdate && s_failureStatus.Contains(runUpdate.Value.Status)) { throw new DeclarativeActionException($"Unexpected failure invoking agent, run {runUpdate.Value.Status}: {agent.Name ?? agent.Id} [{runUpdate.Value.Id}/{conversationId}]"); } +#endif if (autoSend) { From 8d0fa99b133b1b2614008b3e637271d0b28e290b Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Sat, 8 Nov 2025 11:29:40 +0000 Subject: [PATCH 21/40] Add Cosmos DB emulator to CI workflow --- .github/workflows/dotnet-build-and-test.yml | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/.github/workflows/dotnet-build-and-test.yml b/.github/workflows/dotnet-build-and-test.yml index 5abfe2a879..71fa302398 100644 --- a/.github/workflows/dotnet-build-and-test.yml +++ b/.github/workflows/dotnet-build-and-test.yml @@ -143,6 +143,15 @@ jobs: shell: bash run: echo "github.event_name:${{ github.event_name }} matrix.integration-tests:${{ matrix.integration-tests }} github.event.action:${{ github.event.action }} github.event.pull_request.merged:${{ github.event.pull_request.merged }}" + # Start Cosmos DB Emulator for Cosmos-based integration tests (only on Windows) + - name: Start Azure Cosmos DB Emulator + if: runner.os == 'Windows' + shell: pwsh + run: | + Write-Host "Launching Azure Cosmos DB Emulator" + Import-Module "$env:ProgramFiles\Azure Cosmos DB Emulator\PSModules\Microsoft.Azure.CosmosDB.Emulator" + Start-CosmosDbEmulator -NoUI -Key "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==" + - name: Azure CLI Login if: github.event_name != 'pull_request' && matrix.integration-tests uses: azure/login@v2 @@ -168,6 +177,9 @@ jobs: fi done env: + # Cosmos DB Emulator connection settings + COSMOSDB_ENDPOINT: https://localhost:8081 + COSMOSDB_KEY: C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw== # OpenAI Models OpenAI__ApiKey: ${{ secrets.OPENAI__APIKEY }} OpenAI__ChatModelId: ${{ vars.OPENAI__CHATMODELID }} From 823b59a7b7b00a843f971167726457886c90a6fd Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Sat, 8 Nov 2025 11:43:05 +0000 Subject: [PATCH 22/40] Fix Cosmos DB emulator tests: use Skip.If instead of Assert.Fail and start emulator before unit tests --- .github/workflows/dotnet-build-and-test.yml | 18 +++++++++--------- .../CosmosChatMessageStoreTests.cs | 5 +---- .../CosmosCheckpointStoreTests.cs | 6 +----- 3 files changed, 11 insertions(+), 18 deletions(-) diff --git a/.github/workflows/dotnet-build-and-test.yml b/.github/workflows/dotnet-build-and-test.yml index 71fa302398..ef2aa240e0 100644 --- a/.github/workflows/dotnet-build-and-test.yml +++ b/.github/workflows/dotnet-build-and-test.yml @@ -123,6 +123,15 @@ jobs: popd rm -rf "$TEMP_DIR" + # Start Cosmos DB Emulator for Cosmos-based unit tests (only on Windows) + - name: Start Azure Cosmos DB Emulator + if: runner.os == 'Windows' + shell: pwsh + run: | + Write-Host "Launching Azure Cosmos DB Emulator" + Import-Module "$env:ProgramFiles\Azure Cosmos DB Emulator\PSModules\Microsoft.Azure.CosmosDB.Emulator" + Start-CosmosDbEmulator -NoUI -Key "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==" + - name: Run Unit Tests Windows shell: bash run: | @@ -143,15 +152,6 @@ jobs: shell: bash run: echo "github.event_name:${{ github.event_name }} matrix.integration-tests:${{ matrix.integration-tests }} github.event.action:${{ github.event.action }} github.event.pull_request.merged:${{ github.event.pull_request.merged }}" - # Start Cosmos DB Emulator for Cosmos-based integration tests (only on Windows) - - name: Start Azure Cosmos DB Emulator - if: runner.os == 'Windows' - shell: pwsh - run: | - Write-Host "Launching Azure Cosmos DB Emulator" - Import-Module "$env:ProgramFiles\Azure Cosmos DB Emulator\PSModules\Microsoft.Azure.CosmosDB.Emulator" - Start-CosmosDbEmulator -NoUI -Key "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==" - - name: Azure CLI Login if: github.event_name != 'pull_request' && matrix.integration-tests uses: azure/login@v2 diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs index 8a8ba1e445..cde3ffbc85 100644 --- a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs @@ -136,10 +136,7 @@ public void Dispose() private void SkipIfEmulatorNotAvailable() { - if (!this._emulatorAvailable) - { - Assert.Fail("Cosmos DB Emulator is not available. Start the emulator to run these tests."); - } + Skip.If(!this._emulatorAvailable, "Cosmos DB Emulator is not available. Start the emulator to run these tests."); } #region Constructor Tests diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs index e3bebce494..b6962484fb 100644 --- a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs @@ -119,11 +119,7 @@ public async Task DisposeAsync() private void SkipIfEmulatorNotAvailable() { - if (!_emulatorAvailable) - { - // For now let's just fail to see if tests work - Assert.Fail("Cosmos DB Emulator is not available. Start the emulator to run these tests."); - } + Skip.If(!this._emulatorAvailable, "Cosmos DB Emulator is not available. Start the emulator to run these tests."); } #region Constructor Tests From ecc3498fb995c25552d947acb1ac06500184ebaf Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Sat, 8 Nov 2025 12:47:25 +0000 Subject: [PATCH 23/40] Replace Skip.If() with conditional return to fix compilation --- .../CosmosChatMessageStoreTests.cs | 12 ++++++++---- .../CosmosCheckpointStoreTests.cs | 12 ++++++++---- 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs index cde3ffbc85..145e35a1d9 100644 --- a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs @@ -134,10 +134,14 @@ public void Dispose() GC.SuppressFinalize(this); } - private void SkipIfEmulatorNotAvailable() - { - Skip.If(!this._emulatorAvailable, "Cosmos DB Emulator is not available. Start the emulator to run these tests."); - } + private void SkipIfEmulatorNotAvailable() + { + if (!this._emulatorAvailable) + { + // Skip test if emulator is not available (e.g., on Linux CI runners) + return; + } + } #region Constructor Tests diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs index b6962484fb..a4a7b8b6f2 100644 --- a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs @@ -117,10 +117,14 @@ public async Task DisposeAsync() } } - private void SkipIfEmulatorNotAvailable() - { - Skip.If(!this._emulatorAvailable, "Cosmos DB Emulator is not available. Start the emulator to run these tests."); - } + private void SkipIfEmulatorNotAvailable() + { + if (!this._emulatorAvailable) + { + // Skip test if emulator is not available (e.g., on Linux CI runners) + return; + } + } #region Constructor Tests From 7a574c65c6059c24e09eec553617d159fb1b2a8d Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Sat, 8 Nov 2025 13:40:42 +0000 Subject: [PATCH 24/40] Use env var to skip Cosmos tests on non-Windows CI --- .github/workflows/dotnet-build-and-test.yml | 1 + .../CosmosChatMessageStoreTests.cs | 8 ++++++-- .../CosmosCheckpointStoreTests.cs | 8 ++++++-- 3 files changed, 13 insertions(+), 4 deletions(-) diff --git a/.github/workflows/dotnet-build-and-test.yml b/.github/workflows/dotnet-build-and-test.yml index ef2aa240e0..0cf6597f2f 100644 --- a/.github/workflows/dotnet-build-and-test.yml +++ b/.github/workflows/dotnet-build-and-test.yml @@ -131,6 +131,7 @@ jobs: Write-Host "Launching Azure Cosmos DB Emulator" Import-Module "$env:ProgramFiles\Azure Cosmos DB Emulator\PSModules\Microsoft.Azure.CosmosDB.Emulator" Start-CosmosDbEmulator -NoUI -Key "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==" + echo "COSMOS_EMULATOR_AVAILABLE=true" >> $env:GITHUB_ENV - name: Run Unit Tests Windows shell: bash diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs index 145e35a1d9..7b47917ead 100644 --- a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs @@ -136,9 +136,13 @@ public void Dispose() private void SkipIfEmulatorNotAvailable() { - if (!this._emulatorAvailable) + // In CI: Skip if COSMOS_EMULATOR_AVAILABLE is not set to "true" + // Locally: Skip if emulator connection check failed + var ciEmulatorAvailable = string.Equals(Environment.GetEnvironmentVariable("COSMOS_EMULATOR_AVAILABLE"), "true", StringComparison.OrdinalIgnoreCase); + + if (!ciEmulatorAvailable && !this._emulatorAvailable) { - // Skip test if emulator is not available (e.g., on Linux CI runners) + // Skip test if emulator is not available return; } } diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs index a4a7b8b6f2..20a0667c11 100644 --- a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs @@ -119,9 +119,13 @@ public async Task DisposeAsync() private void SkipIfEmulatorNotAvailable() { - if (!this._emulatorAvailable) + // In CI: Skip if COSMOS_EMULATOR_AVAILABLE is not set to "true" + // Locally: Skip if emulator connection check failed + var ciEmulatorAvailable = string.Equals(Environment.GetEnvironmentVariable("COSMOS_EMULATOR_AVAILABLE"), "true", StringComparison.OrdinalIgnoreCase); + + if (!ciEmulatorAvailable && !this._emulatorAvailable) { - // Skip test if emulator is not available (e.g., on Linux CI runners) + // Skip test if emulator is not available return; } } From 8f0a3ba8e36c9261c32c59a857dad75b1af4b571 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Sat, 8 Nov 2025 14:34:03 +0000 Subject: [PATCH 25/40] Add Xunit.SkippableFact package to properly skip Cosmos tests on Linux --- dotnet/Directory.Packages.props | 1 + .../CosmosChatMessageStoreTests.cs | 6 +----- .../CosmosCheckpointStoreTests.cs | 6 +----- .../Microsoft.Agents.AI.CosmosNoSql.UnitTests.csproj | 1 + 4 files changed, 4 insertions(+), 10 deletions(-) diff --git a/dotnet/Directory.Packages.props b/dotnet/Directory.Packages.props index 7ae6e2bd6d..d714be3428 100644 --- a/dotnet/Directory.Packages.props +++ b/dotnet/Directory.Packages.props @@ -113,6 +113,7 @@ + diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs index 7b47917ead..e693658886 100644 --- a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs @@ -140,11 +140,7 @@ private void SkipIfEmulatorNotAvailable() // Locally: Skip if emulator connection check failed var ciEmulatorAvailable = string.Equals(Environment.GetEnvironmentVariable("COSMOS_EMULATOR_AVAILABLE"), "true", StringComparison.OrdinalIgnoreCase); - if (!ciEmulatorAvailable && !this._emulatorAvailable) - { - // Skip test if emulator is not available - return; - } + Xunit.Skip.If(!ciEmulatorAvailable && !this._emulatorAvailable, "Cosmos DB Emulator is not available"); } #region Constructor Tests diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs index 20a0667c11..ace3ed4839 100644 --- a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs @@ -123,11 +123,7 @@ private void SkipIfEmulatorNotAvailable() // Locally: Skip if emulator connection check failed var ciEmulatorAvailable = string.Equals(Environment.GetEnvironmentVariable("COSMOS_EMULATOR_AVAILABLE"), "true", StringComparison.OrdinalIgnoreCase); - if (!ciEmulatorAvailable && !this._emulatorAvailable) - { - // Skip test if emulator is not available - return; - } + Xunit.Skip.If(!ciEmulatorAvailable && !this._emulatorAvailable, "Cosmos DB Emulator is not available"); } #region Constructor Tests diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/Microsoft.Agents.AI.CosmosNoSql.UnitTests.csproj b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/Microsoft.Agents.AI.CosmosNoSql.UnitTests.csproj index d75ef79f99..d2050fb52f 100644 --- a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/Microsoft.Agents.AI.CosmosNoSql.UnitTests.csproj +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/Microsoft.Agents.AI.CosmosNoSql.UnitTests.csproj @@ -18,6 +18,7 @@ + From 76ba62b6b1d748fe5a90fa1fe99d39c00de2dfbc Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Sat, 8 Nov 2025 14:52:03 +0000 Subject: [PATCH 26/40] Change [Fact] to [SkippableFact] for proper test skipping behavior --- .../CosmosChatMessageStoreTests.cs | 48 +++++++++---------- .../CosmosCheckpointStoreTests.cs | 34 ++++++------- 2 files changed, 41 insertions(+), 41 deletions(-) diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs index e693658886..51795e3059 100644 --- a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs @@ -145,7 +145,7 @@ private void SkipIfEmulatorNotAvailable() #region Constructor Tests - [Fact] + [SkippableFact] [Trait("Category", "CosmosDB")] public void Constructor_WithConnectionString_ShouldCreateInstance() { @@ -162,7 +162,7 @@ public void Constructor_WithConnectionString_ShouldCreateInstance() Assert.Equal(TestContainerId, store.ContainerId); } - [Fact] + [SkippableFact] [Trait("Category", "CosmosDB")] public void Constructor_WithConnectionStringNoConversationId_ShouldCreateInstance() { @@ -179,7 +179,7 @@ public void Constructor_WithConnectionStringNoConversationId_ShouldCreateInstanc Assert.Equal(TestContainerId, store.ContainerId); } - [Fact] + [SkippableFact] [Trait("Category", "CosmosDB")] public void Constructor_WithNullConnectionString_ShouldThrowArgumentException() { @@ -188,7 +188,7 @@ public void Constructor_WithNullConnectionString_ShouldThrowArgumentException() new CosmosChatMessageStore((string)null!, TestDatabaseId, TestContainerId, "test-conversation")); } - [Fact] + [SkippableFact] [Trait("Category", "CosmosDB")] public void Constructor_WithEmptyConversationId_ShouldThrowArgumentException() { @@ -203,7 +203,7 @@ public void Constructor_WithEmptyConversationId_ShouldThrowArgumentException() #region AddMessagesAsync Tests - [Fact] + [SkippableFact] [Trait("Category", "CosmosDB")] public async Task AddMessagesAsync_WithSingleMessage_ShouldAddMessageAsync() { @@ -263,7 +263,7 @@ public async Task AddMessagesAsync_WithSingleMessage_ShouldAddMessageAsync() Assert.Equal(ChatRole.User, messageList[0].Role); } - [Fact] + [SkippableFact] [Trait("Category", "CosmosDB")] public async Task AddMessagesAsync_WithMultipleMessages_ShouldAddAllMessagesAsync() { @@ -294,7 +294,7 @@ public async Task AddMessagesAsync_WithMultipleMessages_ShouldAddAllMessagesAsyn #region GetMessagesAsync Tests - [Fact] + [SkippableFact] [Trait("Category", "CosmosDB")] public async Task GetMessagesAsync_WithNoMessages_ShouldReturnEmptyAsync() { @@ -309,7 +309,7 @@ public async Task GetMessagesAsync_WithNoMessages_ShouldReturnEmptyAsync() Assert.Empty(messages); } - [Fact] + [SkippableFact] [Trait("Category", "CosmosDB")] public async Task GetMessagesAsync_WithConversationIsolation_ShouldOnlyReturnMessagesForConversationAsync() { @@ -341,7 +341,7 @@ public async Task GetMessagesAsync_WithConversationIsolation_ShouldOnlyReturnMes #region Integration Tests - [Fact] + [SkippableFact] [Trait("Category", "CosmosDB")] public async Task FullWorkflow_AddAndGet_ShouldWorkCorrectlyAsync() { @@ -385,7 +385,7 @@ public async Task FullWorkflow_AddAndGet_ShouldWorkCorrectlyAsync() #region Disposal Tests - [Fact] + [SkippableFact] [Trait("Category", "CosmosDB")] public void Dispose_AfterUse_ShouldNotThrow() { @@ -397,7 +397,7 @@ public void Dispose_AfterUse_ShouldNotThrow() store.Dispose(); // Should not throw } - [Fact] + [SkippableFact] [Trait("Category", "CosmosDB")] public void Dispose_MultipleCalls_ShouldNotThrow() { @@ -414,7 +414,7 @@ public void Dispose_MultipleCalls_ShouldNotThrow() #region Hierarchical Partitioning Tests - [Fact] + [SkippableFact] [Trait("Category", "CosmosDB")] public void Constructor_WithHierarchicalConnectionString_ShouldCreateInstance() { @@ -431,7 +431,7 @@ public void Constructor_WithHierarchicalConnectionString_ShouldCreateInstance() Assert.Equal(HierarchicalTestContainerId, store.ContainerId); } - [Fact] + [SkippableFact] [Trait("Category", "CosmosDB")] public void Constructor_WithHierarchicalEndpoint_ShouldCreateInstance() { @@ -449,7 +449,7 @@ public void Constructor_WithHierarchicalEndpoint_ShouldCreateInstance() Assert.Equal(HierarchicalTestContainerId, store.ContainerId); } - [Fact] + [SkippableFact] [Trait("Category", "CosmosDB")] public void Constructor_WithHierarchicalCosmosClient_ShouldCreateInstance() { @@ -466,7 +466,7 @@ public void Constructor_WithHierarchicalCosmosClient_ShouldCreateInstance() Assert.Equal(HierarchicalTestContainerId, store.ContainerId); } - [Fact] + [SkippableFact] [Trait("Category", "CosmosDB")] public void Constructor_WithHierarchicalNullTenantId_ShouldThrowArgumentException() { @@ -477,7 +477,7 @@ public void Constructor_WithHierarchicalNullTenantId_ShouldThrowArgumentExceptio new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, null!, "user-456", "session-789")); } - [Fact] + [SkippableFact] [Trait("Category", "CosmosDB")] public void Constructor_WithHierarchicalEmptyUserId_ShouldThrowArgumentException() { @@ -488,7 +488,7 @@ public void Constructor_WithHierarchicalEmptyUserId_ShouldThrowArgumentException new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, "tenant-123", "", "session-789")); } - [Fact] + [SkippableFact] [Trait("Category", "CosmosDB")] public void Constructor_WithHierarchicalWhitespaceSessionId_ShouldThrowArgumentException() { @@ -499,7 +499,7 @@ public void Constructor_WithHierarchicalWhitespaceSessionId_ShouldThrowArgumentE new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", " ")); } - [Fact] + [SkippableFact] [Trait("Category", "CosmosDB")] public async Task AddMessagesAsync_WithHierarchicalPartitioning_ShouldAddMessageWithMetadataAsync() { @@ -548,7 +548,7 @@ public async Task AddMessagesAsync_WithHierarchicalPartitioning_ShouldAddMessage Assert.Equal(SessionId, (string)document!.sessionId); } - [Fact] + [SkippableFact] [Trait("Category", "CosmosDB")] public async Task AddMessagesAsync_WithHierarchicalMultipleMessages_ShouldAddAllMessagesAsync() { @@ -582,7 +582,7 @@ public async Task AddMessagesAsync_WithHierarchicalMultipleMessages_ShouldAddAll Assert.Equal("Third hierarchical message", messageList[2].Text); } - [Fact] + [SkippableFact] [Trait("Category", "CosmosDB")] public async Task GetMessagesAsync_WithHierarchicalPartitionIsolation_ShouldIsolateMessagesByUserIdAsync() { @@ -618,7 +618,7 @@ public async Task GetMessagesAsync_WithHierarchicalPartitionIsolation_ShouldIsol Assert.Equal("Message from user 2", messageList2[0].Text); } - [Fact] + [SkippableFact] [Trait("Category", "CosmosDB")] public async Task SerializeDeserialize_WithHierarchicalPartitioning_ShouldPreserveStateAsync() { @@ -656,7 +656,7 @@ public async Task SerializeDeserialize_WithHierarchicalPartitioning_ShouldPreser Assert.Equal(HierarchicalTestContainerId, deserializedStore.ContainerId); } - [Fact] + [SkippableFact] [Trait("Category", "CosmosDB")] public async Task HierarchicalAndSimplePartitioning_ShouldCoexistAsync() { @@ -689,7 +689,7 @@ public async Task HierarchicalAndSimplePartitioning_ShouldCoexistAsync() Assert.Equal("Hierarchical partitioning message", hierarchicalMessageList[0].Text); } - [Fact] + [SkippableFact] [Trait("Category", "CosmosDB")] public async Task MaxMessagesToRetrieve_ShouldLimitAndReturnMostRecentAsync() { @@ -725,7 +725,7 @@ public async Task MaxMessagesToRetrieve_ShouldLimitAndReturnMostRecentAsync() Assert.Equal("Message 10", messageList[4].Text); } - [Fact] + [SkippableFact] [Trait("Category", "CosmosDB")] public async Task MaxMessagesToRetrieve_Null_ShouldReturnAllMessagesAsync() { diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs index ace3ed4839..224cdc28bc 100644 --- a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs @@ -128,7 +128,7 @@ private void SkipIfEmulatorNotAvailable() #region Constructor Tests - [Fact] + [SkippableFact] public void Constructor_WithCosmosClient_SetsProperties() { // Arrange @@ -142,7 +142,7 @@ public void Constructor_WithCosmosClient_SetsProperties() Assert.Equal(TestContainerId, store.ContainerId); } - [Fact] + [SkippableFact] public void Constructor_WithConnectionString_SetsProperties() { // Arrange @@ -156,7 +156,7 @@ public void Constructor_WithConnectionString_SetsProperties() Assert.Equal(TestContainerId, store.ContainerId); } - [Fact] + [SkippableFact] public void Constructor_WithNullCosmosClient_ThrowsArgumentNullException() { // Act & Assert @@ -164,7 +164,7 @@ public void Constructor_WithNullCosmosClient_ThrowsArgumentNullException() new CosmosCheckpointStore((CosmosClient)null!, TestDatabaseId, TestContainerId)); } - [Fact] + [SkippableFact] public void Constructor_WithNullConnectionString_ThrowsArgumentException() { // Act & Assert @@ -176,7 +176,7 @@ public void Constructor_WithNullConnectionString_ThrowsArgumentException() #region Checkpoint Operations Tests - [Fact] + [SkippableFact] public async Task CreateCheckpointAsync_NewCheckpoint_CreatesSuccessfully() { SkipIfEmulatorNotAvailable(); @@ -196,7 +196,7 @@ public async Task CreateCheckpointAsync_NewCheckpoint_CreatesSuccessfully() Assert.NotEmpty(checkpointInfo.CheckpointId); } - [Fact] + [SkippableFact] public async Task RetrieveCheckpointAsync_ExistingCheckpoint_ReturnsCorrectValue() { SkipIfEmulatorNotAvailable(); @@ -217,7 +217,7 @@ public async Task RetrieveCheckpointAsync_ExistingCheckpoint_ReturnsCorrectValue Assert.Equal("Hello, World!", messageProp.GetString()); } - [Fact] + [SkippableFact] public async Task RetrieveCheckpointAsync_NonExistentCheckpoint_ThrowsInvalidOperationException() { SkipIfEmulatorNotAvailable(); @@ -232,7 +232,7 @@ await Assert.ThrowsAsync(() => store.RetrieveCheckpointAsync(runId, fakeCheckpointInfo).AsTask()); } - [Fact] + [SkippableFact] public async Task RetrieveIndexAsync_EmptyStore_ReturnsEmptyCollection() { SkipIfEmulatorNotAvailable(); @@ -249,7 +249,7 @@ public async Task RetrieveIndexAsync_EmptyStore_ReturnsEmptyCollection() Assert.Empty(index); } - [Fact] + [SkippableFact] public async Task RetrieveIndexAsync_WithCheckpoints_ReturnsAllCheckpoints() { SkipIfEmulatorNotAvailable(); @@ -274,7 +274,7 @@ public async Task RetrieveIndexAsync_WithCheckpoints_ReturnsAllCheckpoints() Assert.Contains(index, c => c.CheckpointId == checkpoint3.CheckpointId); } - [Fact] + [SkippableFact] public async Task CreateCheckpointAsync_WithParent_CreatesHierarchy() { SkipIfEmulatorNotAvailable(); @@ -294,7 +294,7 @@ public async Task CreateCheckpointAsync_WithParent_CreatesHierarchy() Assert.Equal(runId, childCheckpoint.RunId); } - [Fact] + [SkippableFact] public async Task RetrieveIndexAsync_WithParentFilter_ReturnsFilteredResults() { SkipIfEmulatorNotAvailable(); @@ -330,7 +330,7 @@ public async Task RetrieveIndexAsync_WithParentFilter_ReturnsFilteredResults() #region Run Isolation Tests - [Fact] + [SkippableFact] public async Task CheckpointOperations_DifferentRuns_IsolatesData() { SkipIfEmulatorNotAvailable(); @@ -360,7 +360,7 @@ public async Task CheckpointOperations_DifferentRuns_IsolatesData() #region Error Handling Tests - [Fact] + [SkippableFact] public async Task CreateCheckpointAsync_WithNullRunId_ThrowsArgumentException() { SkipIfEmulatorNotAvailable(); @@ -374,7 +374,7 @@ await Assert.ThrowsAsync(() => store.CreateCheckpointAsync(null!, checkpointValue).AsTask()); } - [Fact] + [SkippableFact] public async Task CreateCheckpointAsync_WithEmptyRunId_ThrowsArgumentException() { SkipIfEmulatorNotAvailable(); @@ -388,7 +388,7 @@ await Assert.ThrowsAsync(() => store.CreateCheckpointAsync("", checkpointValue).AsTask()); } - [Fact] + [SkippableFact] public async Task RetrieveCheckpointAsync_WithNullCheckpointInfo_ThrowsArgumentNullException() { SkipIfEmulatorNotAvailable(); @@ -406,7 +406,7 @@ await Assert.ThrowsAsync(() => #region Disposal Tests - [Fact] + [SkippableFact] public async Task Dispose_AfterDisposal_ThrowsObjectDisposedException() { SkipIfEmulatorNotAvailable(); @@ -423,7 +423,7 @@ await Assert.ThrowsAsync(() => store.CreateCheckpointAsync("test-run", checkpointValue).AsTask()); } - [Fact] + [SkippableFact] public void Dispose_MultipleCalls_DoesNotThrow() { SkipIfEmulatorNotAvailable(); From 4862ca449224b4aa36a6d11e7808d40c3ac310c6 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Sat, 8 Nov 2025 15:31:03 +0000 Subject: [PATCH 27/40] Remove stale Microsoft.Agents.AI.Abstractions.CosmosNoSql directory --- .../CosmosChatMessageStore.cs | 695 ------------------ 1 file changed, 695 deletions(-) delete mode 100644 dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosChatMessageStore.cs diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosChatMessageStore.cs deleted file mode 100644 index 4ca653a84c..0000000000 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions.CosmosNoSql/CosmosChatMessageStore.cs +++ /dev/null @@ -1,695 +0,0 @@ -// Copyright (c) Microsoft. All rights reserved. - -using System; -using System.Collections.Generic; -using System.Diagnostics.CodeAnalysis; -using System.Linq; -using System.Text.Json; -using System.Text.Json.Serialization; -using System.Threading; -using System.Threading.Tasks; -using Azure.Core; -using Azure.Identity; -using Microsoft.Azure.Cosmos; -using Microsoft.Extensions.AI; -using Microsoft.Shared.Diagnostics; - -namespace Microsoft.Agents.AI; - -/// -/// Provides a Cosmos DB implementation of the abstract class. -/// -[RequiresUnreferencedCode("The CosmosChatMessageStore uses JSON serialization which is incompatible with trimming.")] -[RequiresDynamicCode("The CosmosChatMessageStore uses JSON serialization which is incompatible with NativeAOT.")] -public sealed class CosmosChatMessageStore : ChatMessageStore, IDisposable -{ - private readonly CosmosClient _cosmosClient; - private readonly Container _container; - private readonly string _conversationId; - private readonly string _databaseId; - private readonly string _containerId; - private readonly bool _ownsClient; - private bool _disposed; - - // Hierarchical partition key support - private readonly string? _tenantId; - private readonly string? _userId; - private readonly PartitionKey _partitionKey; - private readonly bool _useHierarchicalPartitioning; - - /// - /// Cached JSON serializer options for .NET 9.0 compatibility. - /// - private static readonly JsonSerializerOptions s_defaultJsonOptions = CreateDefaultJsonOptions(); - - private static JsonSerializerOptions CreateDefaultJsonOptions() - { - var options = new JsonSerializerOptions(); -#if NET9_0_OR_GREATER - // Configure TypeInfoResolver for .NET 9.0 to enable JSON serialization - options.TypeInfoResolver = new System.Text.Json.Serialization.Metadata.DefaultJsonTypeInfoResolver(); -#endif - return options; - } - - /// - /// Gets or sets the maximum number of messages to return in a single query batch. - /// Default is 100 for optimal performance. - /// - public int MaxItemCount { get; set; } = 100; - - /// - /// Gets or sets the maximum number of items per transactional batch operation. - /// Default is 100, maximum allowed by Cosmos DB is 100. - /// - public int MaxBatchSize { get; set; } = 100; - - /// - /// Gets or sets the Time-To-Live (TTL) in seconds for messages. - /// Default is 86400 seconds (24 hours). Set to null to disable TTL. - /// - public int? MessageTtlSeconds { get; set; } = 86400; - - /// - /// Gets the conversation ID associated with this message store. - /// - public string ConversationId => this._conversationId; - - /// - /// Gets the database ID associated with this message store. - /// - public string DatabaseId => this._databaseId; - - /// - /// Gets the container ID associated with this message store. - /// - public string ContainerId => this._containerId; - - /// - /// Initializes a new instance of the class using a connection string. - /// - /// The Cosmos DB connection string. - /// The identifier of the Cosmos DB database. - /// The identifier of the Cosmos DB container. - /// Thrown when any required parameter is null. - /// Thrown when any string parameter is null or whitespace. - public CosmosChatMessageStore(string connectionString, string databaseId, string containerId) - : this(connectionString, databaseId, containerId, Guid.NewGuid().ToString("N")) - { - } - - /// - /// Initializes a new instance of the class using a connection string. - /// - /// The Cosmos DB connection string. - /// The identifier of the Cosmos DB database. - /// The identifier of the Cosmos DB container. - /// The unique identifier for this conversation thread. - /// Thrown when any required parameter is null. - /// Thrown when any string parameter is null or whitespace. - public CosmosChatMessageStore(string connectionString, string databaseId, string containerId, string conversationId) - { - this._cosmosClient = new CosmosClient(Throw.IfNullOrWhitespace(connectionString)); - this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); - this._conversationId = Throw.IfNullOrWhitespace(conversationId); - this._databaseId = databaseId; - this._containerId = containerId; - this._ownsClient = true; - - // Initialize simple partitioning mode - this._tenantId = null; - this._userId = null; - this._useHierarchicalPartitioning = false; - this._partitionKey = new PartitionKey(conversationId); - } - - /// - /// Initializes a new instance of the class using TokenCredential for authentication. - /// - /// The Cosmos DB account endpoint URI. - /// The TokenCredential to use for authentication (e.g., DefaultAzureCredential, ManagedIdentityCredential). - /// The identifier of the Cosmos DB database. - /// The identifier of the Cosmos DB container. - /// Thrown when any required parameter is null. - /// Thrown when any string parameter is null or whitespace. - public CosmosChatMessageStore(string accountEndpoint, TokenCredential tokenCredential, string databaseId, string containerId) - : this(accountEndpoint, tokenCredential, databaseId, containerId, Guid.NewGuid().ToString("N")) - { - } - - /// - /// Initializes a new instance of the class using a TokenCredential for authentication. - /// - /// The Cosmos DB account endpoint URI. - /// The TokenCredential to use for authentication (e.g., DefaultAzureCredential, ManagedIdentityCredential). - /// The identifier of the Cosmos DB database. - /// The identifier of the Cosmos DB container. - /// The unique identifier for this conversation thread. - /// Thrown when any required parameter is null. - /// Thrown when any string parameter is null or whitespace. - public CosmosChatMessageStore(string accountEndpoint, TokenCredential tokenCredential, string databaseId, string containerId, string conversationId) - { - this._cosmosClient = new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), Throw.IfNull(tokenCredential)); - this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); - this._conversationId = Throw.IfNullOrWhitespace(conversationId); - this._databaseId = databaseId; - this._containerId = containerId; - this._ownsClient = true; - - // Initialize simple partitioning mode - this._tenantId = null; - this._userId = null; - this._useHierarchicalPartitioning = false; - this._partitionKey = new PartitionKey(conversationId); - } - - /// - /// Initializes a new instance of the class using an existing . - /// - /// The instance to use for Cosmos DB operations. - /// The identifier of the Cosmos DB database. - /// The identifier of the Cosmos DB container. - /// Thrown when is null. - /// Thrown when any string parameter is null or whitespace. - public CosmosChatMessageStore(CosmosClient cosmosClient, string databaseId, string containerId) - : this(cosmosClient, databaseId, containerId, Guid.NewGuid().ToString("N")) - { - } - - /// - /// Initializes a new instance of the class using an existing . - /// - /// The instance to use for Cosmos DB operations. - /// The identifier of the Cosmos DB database. - /// The identifier of the Cosmos DB container. - /// The unique identifier for this conversation thread. - /// Thrown when is null. - /// Thrown when any string parameter is null or whitespace. - public CosmosChatMessageStore(CosmosClient cosmosClient, string databaseId, string containerId, string conversationId) - { - this._cosmosClient = Throw.IfNull(cosmosClient); - this._ownsClient = false; - - this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); - this._conversationId = Throw.IfNullOrWhitespace(conversationId); - this._databaseId = databaseId; - this._containerId = containerId; - - // Initialize simple partitioning mode - this._tenantId = null; - this._userId = null; - this._useHierarchicalPartitioning = false; - this._partitionKey = new PartitionKey(conversationId); - } - - /// - /// Initializes a new instance of the class using a connection string with hierarchical partition keys. - /// - /// The Cosmos DB connection string. - /// The identifier of the Cosmos DB database. - /// The identifier of the Cosmos DB container. - /// The tenant identifier for hierarchical partitioning. - /// The user identifier for hierarchical partitioning. - /// The session identifier for hierarchical partitioning. - /// Thrown when any required parameter is null. - /// Thrown when any string parameter is null or whitespace. - public CosmosChatMessageStore(string connectionString, string databaseId, string containerId, string tenantId, string userId, string sessionId) - { - this._cosmosClient = new CosmosClient(Throw.IfNullOrWhitespace(connectionString)); - this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); - this._conversationId = Throw.IfNullOrWhitespace(sessionId); // Use sessionId as conversationId for compatibility - this._databaseId = databaseId; - this._containerId = containerId; - - // Initialize hierarchical partitioning mode - this._tenantId = Throw.IfNullOrWhitespace(tenantId); - this._userId = Throw.IfNullOrWhitespace(userId); - this._useHierarchicalPartitioning = true; - // Use native hierarchical partition key with PartitionKeyBuilder - this._partitionKey = new PartitionKeyBuilder() - .Add(tenantId) - .Add(userId) - .Add(sessionId) - .Build(); - } - - /// - /// Initializes a new instance of the class using a TokenCredential for authentication with hierarchical partition keys. - /// - /// The Cosmos DB account endpoint URI. - /// The TokenCredential to use for authentication (e.g., DefaultAzureCredential, ManagedIdentityCredential). - /// The identifier of the Cosmos DB database. - /// The identifier of the Cosmos DB container. - /// The tenant identifier for hierarchical partitioning. - /// The user identifier for hierarchical partitioning. - /// The session identifier for hierarchical partitioning. - /// Thrown when any required parameter is null. - /// Thrown when any string parameter is null or whitespace. - public CosmosChatMessageStore(string accountEndpoint, TokenCredential tokenCredential, string databaseId, string containerId, string tenantId, string userId, string sessionId) - { - this._cosmosClient = new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), Throw.IfNull(tokenCredential)); - this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); - this._conversationId = Throw.IfNullOrWhitespace(sessionId); // Use sessionId as conversationId for compatibility - this._databaseId = databaseId; - this._containerId = containerId; - - // Initialize hierarchical partitioning mode - this._tenantId = Throw.IfNullOrWhitespace(tenantId); - this._userId = Throw.IfNullOrWhitespace(userId); - this._useHierarchicalPartitioning = true; - // Use native hierarchical partition key with PartitionKeyBuilder - this._partitionKey = new PartitionKeyBuilder() - .Add(tenantId) - .Add(userId) - .Add(sessionId) - .Build(); - } - - /// - /// Initializes a new instance of the class using an existing with hierarchical partition keys. - /// - /// The instance to use for Cosmos DB operations. - /// The identifier of the Cosmos DB database. - /// The identifier of the Cosmos DB container. - /// The tenant identifier for hierarchical partitioning. - /// The user identifier for hierarchical partitioning. - /// The session identifier for hierarchical partitioning. - /// Thrown when is null. - /// Thrown when any string parameter is null or whitespace. - public CosmosChatMessageStore(CosmosClient cosmosClient, string databaseId, string containerId, string tenantId, string userId, string sessionId) - { - this._cosmosClient = Throw.IfNull(cosmosClient); - this._ownsClient = false; - - this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); - this._conversationId = Throw.IfNullOrWhitespace(sessionId); // Use sessionId as conversationId for compatibility - this._databaseId = databaseId; - this._containerId = containerId; - - // Initialize hierarchical partitioning mode - this._tenantId = Throw.IfNullOrWhitespace(tenantId); - this._userId = Throw.IfNullOrWhitespace(userId); - this._useHierarchicalPartitioning = true; - // Use native hierarchical partition key with PartitionKeyBuilder - this._partitionKey = new PartitionKeyBuilder() - .Add(tenantId) - .Add(userId) - .Add(sessionId) - .Build(); - } - - /// - /// Initializes a new instance of the class from previously serialized state. - /// - /// A representing the serialized state of the message store. - /// The instance to use for Cosmos DB operations. - /// Optional settings for customizing the JSON deserialization process. - /// Thrown when is null. - /// Thrown when the serialized state cannot be deserialized. - public CosmosChatMessageStore(JsonElement serializedStoreState, CosmosClient cosmosClient, JsonSerializerOptions? jsonSerializerOptions = null) - { - this._cosmosClient = cosmosClient ?? throw new ArgumentNullException(nameof(cosmosClient)); - this._ownsClient = false; - - if (serializedStoreState.ValueKind is JsonValueKind.Object) - { - var state = JsonSerializer.Deserialize(serializedStoreState, jsonSerializerOptions); - if (state?.ConversationIdentifier is { } conversationId && state.DatabaseIdentifier is { } databaseId && state.ContainerIdentifier is { } containerId) - { - this._conversationId = conversationId; - this._databaseId = databaseId; - this._containerId = containerId; - this._container = this._cosmosClient.GetContainer(databaseId, containerId); - - // Initialize hierarchical partitioning if available in state - this._tenantId = state.TenantId; - this._userId = state.UserId; - this._useHierarchicalPartitioning = state.UseHierarchicalPartitioning; - - this._partitionKey = (this._useHierarchicalPartitioning && this._tenantId != null && this._userId != null) - ? new PartitionKeyBuilder() - .Add(this._tenantId) - .Add(this._userId) - .Add(conversationId) - .Build() - : new PartitionKey(conversationId); - - return; - } - } - - throw new ArgumentException("Invalid serialized state", nameof(serializedStoreState)); - } - - /// - public override async Task> GetMessagesAsync(CancellationToken cancellationToken = default) - { -#pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks - if (this._disposed) - { - throw new ObjectDisposedException(GetType().FullName); - } -#pragma warning restore CA1513 - - // Use type discriminator for efficient queries - var query = new QueryDefinition("SELECT * FROM c WHERE c.conversationId = @conversationId AND c.type = @type ORDER BY c.timestamp ASC") - .WithParameter("@conversationId", this._conversationId) - .WithParameter("@type", "ChatMessage"); - - var iterator = this._container.GetItemQueryIterator(query, requestOptions: new QueryRequestOptions - { - PartitionKey = this._partitionKey, - MaxItemCount = this.MaxItemCount // Configurable query performance - }); - - var messages = new List(); - - while (iterator.HasMoreResults) - { - var response = await iterator.ReadNextAsync(cancellationToken).ConfigureAwait(false); - - foreach (var document in response) - { - if (!string.IsNullOrEmpty(document.Message)) - { - var message = JsonSerializer.Deserialize(document.Message, s_defaultJsonOptions); - if (message != null) - { - messages.Add(message); - } - } - } - } - - return messages; - } - - /// - public override async Task AddMessagesAsync(IEnumerable messages, CancellationToken cancellationToken = default) - { - if (messages is null) - { - throw new ArgumentNullException(nameof(messages)); - } - -#pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks - if (this._disposed) - { - throw new ObjectDisposedException(GetType().FullName); - } -#pragma warning restore CA1513 - - var messageList = messages.ToList(); - if (messageList.Count == 0) - { - return; - } - - // Use transactional batch for atomic operations - if (messageList.Count > 1) - { - await this.AddMessagesInBatchAsync(messageList, cancellationToken).ConfigureAwait(false); - } - else - { - await this.AddSingleMessageAsync(messageList[0], cancellationToken).ConfigureAwait(false); - } - } - - /// - /// Adds multiple messages using transactional batch operations for atomicity. - /// - private async Task AddMessagesInBatchAsync(List messages, CancellationToken cancellationToken) - { - var currentTimestamp = DateTimeOffset.UtcNow.ToUnixTimeSeconds(); - - // Process messages in optimal batch sizes - for (int i = 0; i < messages.Count; i += this.MaxBatchSize) - { - var batchMessages = messages.Skip(i).Take(this.MaxBatchSize).ToList(); - await this.ExecuteBatchOperationAsync(batchMessages, currentTimestamp, cancellationToken).ConfigureAwait(false); - } - } - - /// - /// Executes a single batch operation with retry logic and enhanced error handling. - /// - private async Task ExecuteBatchOperationAsync(List messages, long timestamp, CancellationToken cancellationToken) - { - var batch = this._container.CreateTransactionalBatch(this._partitionKey); - - foreach (var message in messages) - { - var document = this.CreateMessageDocument(message, timestamp); - batch.CreateItem(document); - } - - try - { - var response = await batch.ExecuteAsync(cancellationToken).ConfigureAwait(false); - if (!response.IsSuccessStatusCode) - { - throw new InvalidOperationException($"Batch operation failed with status: {response.StatusCode}. Details: {response.ErrorMessage}"); - } - } - catch (CosmosException ex) when (ex.StatusCode == System.Net.HttpStatusCode.RequestEntityTooLarge) - { - // If batch is too large, split into smaller batches - if (messages.Count == 1) - { - // Can't split further, use single operation - await this.AddSingleMessageAsync(messages[0], cancellationToken).ConfigureAwait(false); - return; - } - - // Split the batch in half and retry - var midpoint = messages.Count / 2; - var firstHalf = messages.Take(midpoint).ToList(); - var secondHalf = messages.Skip(midpoint).ToList(); - - await this.ExecuteBatchOperationAsync(firstHalf, timestamp, cancellationToken).ConfigureAwait(false); - await this.ExecuteBatchOperationAsync(secondHalf, timestamp, cancellationToken).ConfigureAwait(false); - } - catch (CosmosException ex) when (ex.StatusCode == (System.Net.HttpStatusCode)429) // TooManyRequests - { - // Handle rate limiting with exponential backoff - await Task.Delay(TimeSpan.FromMilliseconds(ex.RetryAfter?.TotalMilliseconds ?? 1000), cancellationToken).ConfigureAwait(false); - await this.ExecuteBatchOperationAsync(messages, timestamp, cancellationToken).ConfigureAwait(false); - } - } - - /// - /// Adds a single message to the store. - /// - private async Task AddSingleMessageAsync(ChatMessage message, CancellationToken cancellationToken) - { - var document = this.CreateMessageDocument(message, DateTimeOffset.UtcNow.ToUnixTimeSeconds()); - await this._container.CreateItemAsync(document, this._partitionKey, cancellationToken: cancellationToken).ConfigureAwait(false); - } - - /// - /// Creates a message document with enhanced metadata. - /// - private CosmosMessageDocument CreateMessageDocument(ChatMessage message, long timestamp) - { - return new CosmosMessageDocument - { - Id = Guid.NewGuid().ToString(), - ConversationId = this._conversationId, - Timestamp = timestamp, - MessageId = message.MessageId ?? Guid.NewGuid().ToString(), - Role = message.Role.Value ?? "unknown", - Message = JsonSerializer.Serialize(message, s_defaultJsonOptions), - Type = "ChatMessage", // Type discriminator - Ttl = this.MessageTtlSeconds, // Configurable TTL - // Include hierarchical metadata when using hierarchical partitioning - TenantId = this._useHierarchicalPartitioning ? this._tenantId : null, - UserId = this._useHierarchicalPartitioning ? this._userId : null, - SessionId = this._useHierarchicalPartitioning ? this._conversationId : null - }; - } - - /// - public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) - { -#pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks - if (this._disposed) - { - throw new ObjectDisposedException(GetType().FullName); - } -#pragma warning restore CA1513 - - var state = new StoreState - { - ConversationIdentifier = this._conversationId, - DatabaseIdentifier = this.DatabaseId, - ContainerIdentifier = this.ContainerId, - TenantId = this._tenantId, - UserId = this._userId, - UseHierarchicalPartitioning = this._useHierarchicalPartitioning - }; - - var options = jsonSerializerOptions ?? s_defaultJsonOptions; - return JsonSerializer.SerializeToElement(state, options); - } - - /// - /// Gets the count of messages in this conversation. - /// This is an additional utility method beyond the base contract. - /// - /// The cancellation token. - /// The number of messages in the conversation. - public async Task GetMessageCountAsync(CancellationToken cancellationToken = default) - { -#pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks - if (this._disposed) - { - throw new ObjectDisposedException(GetType().FullName); - } -#pragma warning restore CA1513 - - // Efficient count query - var query = new QueryDefinition("SELECT VALUE COUNT(1) FROM c WHERE c.conversationId = @conversationId AND c.Type = @type") - .WithParameter("@conversationId", this._conversationId) - .WithParameter("@type", "ChatMessage"); - - var iterator = this._container.GetItemQueryIterator(query, requestOptions: new QueryRequestOptions - { - PartitionKey = this._partitionKey - }); - - if (iterator.HasMoreResults) - { - var response = await iterator.ReadNextAsync(cancellationToken).ConfigureAwait(false); - return response.FirstOrDefault(); - } - - return 0; - } - - /// - /// Deletes all messages in this conversation. - /// This is an additional utility method beyond the base contract. - /// - /// The cancellation token. - /// The number of messages deleted. - public async Task ClearMessagesAsync(CancellationToken cancellationToken = default) - { -#pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks - if (this._disposed) - { - throw new ObjectDisposedException(GetType().FullName); - } -#pragma warning restore CA1513 - - // Batch delete for efficiency - var query = new QueryDefinition("SELECT VALUE c.id FROM c WHERE c.conversationId = @conversationId AND c.Type = @type") - .WithParameter("@conversationId", this._conversationId) - .WithParameter("@type", "ChatMessage"); - - var iterator = this._container.GetItemQueryIterator(query, requestOptions: new QueryRequestOptions - { - PartitionKey = new PartitionKey(this._conversationId), - MaxItemCount = this.MaxItemCount - }); - - var deletedCount = 0; - var partitionKey = new PartitionKey(this._conversationId); - - while (iterator.HasMoreResults) - { - var response = await iterator.ReadNextAsync(cancellationToken).ConfigureAwait(false); - var batch = this._container.CreateTransactionalBatch(partitionKey); - var batchItemCount = 0; - - foreach (var itemId in response) - { - if (!string.IsNullOrEmpty(itemId)) - { - batch.DeleteItem(itemId); - batchItemCount++; - deletedCount++; - } - } - - if (batchItemCount > 0) - { - await batch.ExecuteAsync(cancellationToken).ConfigureAwait(false); - } - } - - return deletedCount; - } - - /// - public void Dispose() - { - if (!this._disposed) - { - if (this._ownsClient) - { - this._cosmosClient?.Dispose(); - } - this._disposed = true; - } - } - - private sealed class StoreState - { - public string ConversationIdentifier { get; set; } = string.Empty; - public string DatabaseIdentifier { get; set; } = string.Empty; - public string ContainerIdentifier { get; set; } = string.Empty; - public string? TenantId { get; set; } - public string? UserId { get; set; } - public bool UseHierarchicalPartitioning { get; set; } - } - - /// - /// Represents a document stored in Cosmos DB for chat messages. - /// - [SuppressMessage("Performance", "CA1812:Avoid uninstantiated internal classes", Justification = "Instantiated by Cosmos DB operations")] - private sealed class CosmosMessageDocument - { - [Newtonsoft.Json.JsonProperty("id")] - public string Id { get; set; } = string.Empty; - - [Newtonsoft.Json.JsonProperty("conversationId")] - public string ConversationId { get; set; } = string.Empty; - - [Newtonsoft.Json.JsonProperty("timestamp")] - public long Timestamp { get; set; } - - [Newtonsoft.Json.JsonProperty("messageId")] - public string MessageId { get; set; } = string.Empty; - - [Newtonsoft.Json.JsonProperty("role")] - public string Role { get; set; } = string.Empty; - - [Newtonsoft.Json.JsonProperty("message")] - public string Message { get; set; } = string.Empty; - - [Newtonsoft.Json.JsonProperty("type")] - public string Type { get; set; } = string.Empty; - - [Newtonsoft.Json.JsonProperty("ttl")] - public int? Ttl { get; set; } - - /// - /// Tenant ID for hierarchical partitioning scenarios (optional). - /// - [Newtonsoft.Json.JsonProperty("tenantId")] - public string? TenantId { get; set; } - - /// - /// User ID for hierarchical partitioning scenarios (optional). - /// - [Newtonsoft.Json.JsonProperty("userId")] - public string? UserId { get; set; } - - /// - /// Session ID for hierarchical partitioning scenarios (same as ConversationId for compatibility). - /// - [Newtonsoft.Json.JsonProperty("sessionId")] - public string? SessionId { get; set; } - } -} \ No newline at end of file From b5675d648c3ff9e603862a942349362a9b0e9ae2 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Sat, 8 Nov 2025 16:15:42 +0000 Subject: [PATCH 28/40] Fix code formatting: add braces, this. qualifications, and final newlines --- .../CosmosChatMessageStore.cs | 13 +++-- .../CosmosCheckpointStore.cs | 50 +++++++++++++------ .../CosmosDBChatExtensions.cs | 35 ++++++++++++- .../CosmosDBWorkflowExtensions.cs | 50 ++++++++++++++++++- 4 files changed, 123 insertions(+), 25 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs index 7c61007c29..bd002b6439 100644 --- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs @@ -5,7 +5,6 @@ using System.Diagnostics.CodeAnalysis; using System.Linq; using System.Text.Json; -using System.Text.Json.Serialization; using System.Threading; using System.Threading.Tasks; using Azure.Core; @@ -308,7 +307,7 @@ public override async Task> GetMessagesAsync(Cancellati #pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks if (this._disposed) { - throw new ObjectDisposedException(GetType().FullName); + throw new ObjectDisposedException(this.GetType().FullName); } #pragma warning restore CA1513 @@ -375,7 +374,7 @@ public override async Task AddMessagesAsync(IEnumerable messages, C #pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks if (this._disposed) { - throw new ObjectDisposedException(GetType().FullName); + throw new ObjectDisposedException(this.GetType().FullName); } #pragma warning restore CA1513 @@ -495,7 +494,7 @@ public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptio #pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks if (this._disposed) { - throw new ObjectDisposedException(GetType().FullName); + throw new ObjectDisposedException(this.GetType().FullName); } #pragma warning restore CA1513 @@ -522,7 +521,7 @@ public async Task GetMessageCountAsync(CancellationToken cancellationToken #pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks if (this._disposed) { - throw new ObjectDisposedException(GetType().FullName); + throw new ObjectDisposedException(this.GetType().FullName); } #pragma warning restore CA1513 @@ -556,7 +555,7 @@ public async Task ClearMessagesAsync(CancellationToken cancellationToken = #pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks if (this._disposed) { - throw new ObjectDisposedException(GetType().FullName); + throw new ObjectDisposedException(this.GetType().FullName); } #pragma warning restore CA1513 @@ -668,4 +667,4 @@ private sealed class CosmosMessageDocument [Newtonsoft.Json.JsonProperty("sessionId")] public string? SessionId { get; set; } } -} \ No newline at end of file +} diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosCheckpointStore.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosCheckpointStore.cs index daae028d05..b051a6501c 100644 --- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosCheckpointStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosCheckpointStore.cs @@ -40,9 +40,9 @@ public CosmosCheckpointStore(string connectionString, string databaseId, string { var cosmosClientOptions = new CosmosClientOptions(); - _cosmosClient = new CosmosClient(Throw.IfNullOrWhitespace(connectionString), cosmosClientOptions); - _container = _cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); - _ownsClient = true; + this._cosmosClient = new CosmosClient(Throw.IfNullOrWhitespace(connectionString), cosmosClientOptions); + this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); + this._ownsClient = true; } /// @@ -88,21 +88,26 @@ public CosmosCheckpointStore(CosmosClient cosmosClient, string databaseId, strin /// /// Gets the identifier of the Cosmos DB database. /// - public string DatabaseId => _container.Database.Id; + public string DatabaseId => this._container.Database.Id; /// /// Gets the identifier of the Cosmos DB container. /// - public string ContainerId => _container.Id; + public string ContainerId => this._container.Id; /// public override async ValueTask CreateCheckpointAsync(string runId, JsonElement value, CheckpointInfo? parent = null) { if (string.IsNullOrWhiteSpace(runId)) + { throw new ArgumentException("Cannot be null or whitespace", nameof(runId)); + } + #pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks - if (_disposed) - throw new ObjectDisposedException(GetType().FullName); + if (this._disposed) + { + throw new ObjectDisposedException(this.GetType().FullName); + } #pragma warning restore CA1513 var checkpointId = Guid.NewGuid().ToString("N"); @@ -118,7 +123,7 @@ public override async ValueTask CreateCheckpointAsync(string run Timestamp = DateTimeOffset.UtcNow.ToUnixTimeSeconds() }; - await _container.CreateItemAsync(document, new PartitionKey(runId)).ConfigureAwait(false); + await this._container.CreateItemAsync(document, new PartitionKey(runId)).ConfigureAwait(false); return checkpointInfo; } @@ -126,19 +131,27 @@ public override async ValueTask CreateCheckpointAsync(string run public override async ValueTask RetrieveCheckpointAsync(string runId, CheckpointInfo key) { if (string.IsNullOrWhiteSpace(runId)) + { throw new ArgumentException("Cannot be null or whitespace", nameof(runId)); + } + if (key is null) + { throw new ArgumentNullException(nameof(key)); + } + #pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks - if (_disposed) - throw new ObjectDisposedException(GetType().FullName); + if (this._disposed) + { + throw new ObjectDisposedException(this.GetType().FullName); + } #pragma warning restore CA1513 var id = $"{runId}_{key.CheckpointId}"; try { - var response = await _container.ReadItemAsync(id, new PartitionKey(runId)).ConfigureAwait(false); + var response = await this._container.ReadItemAsync(id, new PartitionKey(runId)).ConfigureAwait(false); using var document = JsonDocument.Parse(response.Resource.Value.ToString()); return document.RootElement.Clone(); } @@ -152,10 +165,15 @@ public override async ValueTask RetrieveCheckpointAsync(string runI public override async ValueTask> RetrieveIndexAsync(string runId, CheckpointInfo? withParent = null) { if (string.IsNullOrWhiteSpace(runId)) + { throw new ArgumentException("Cannot be null or whitespace", nameof(runId)); + } + #pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks - if (_disposed) - throw new ObjectDisposedException(GetType().FullName); + if (this._disposed) + { + throw new ObjectDisposedException(this.GetType().FullName); + } #pragma warning restore CA1513 QueryDefinition query = withParent == null @@ -165,7 +183,7 @@ public override async ValueTask> RetrieveIndexAsync( .WithParameter("@runId", runId) .WithParameter("@parentCheckpointId", withParent.CheckpointId); - var iterator = _container.GetItemQueryIterator(query); + var iterator = this._container.GetItemQueryIterator(query); var checkpoints = new List(); while (iterator.HasMoreResults) @@ -227,7 +245,7 @@ internal sealed class CosmosCheckpointDocument /// /// Represents the result of a checkpoint query. /// - [System.Diagnostics.CodeAnalysis.SuppressMessage("Performance", "CA1812:Avoid uninstantiated internal classes", Justification = "Instantiated by Cosmos DB query deserialization")] + [SuppressMessage("Performance", "CA1812:Avoid uninstantiated internal classes", Justification = "Instantiated by Cosmos DB query deserialization")] private sealed class CheckpointQueryResult { public string RunId { get; set; } = string.Empty; @@ -259,4 +277,4 @@ public CosmosCheckpointStore(CosmosClient cosmosClient, string databaseId, strin : base(cosmosClient, databaseId, containerId) { } -} \ No newline at end of file +} diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBChatExtensions.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBChatExtensions.cs index 9b9c76bd34..f20a6982cb 100644 --- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBChatExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBChatExtensions.cs @@ -31,13 +31,24 @@ public static ChatClientAgentOptions WithCosmosDBMessageStore( string containerId) { if (options is null) + { throw new ArgumentNullException(nameof(options)); + } + if (string.IsNullOrWhiteSpace(connectionString)) + { throw new ArgumentException("Cannot be null or whitespace", nameof(connectionString)); + } + if (string.IsNullOrWhiteSpace(databaseId)) + { throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); + } + if (string.IsNullOrWhiteSpace(containerId)) + { throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); + } options.ChatMessageStoreFactory = context => new CosmosChatMessageStore(connectionString, databaseId, containerId); return options; @@ -62,13 +73,24 @@ public static ChatClientAgentOptions WithCosmosDBMessageStoreUsingManagedIdentit string containerId) { if (options is null) + { throw new ArgumentNullException(nameof(options)); + } + if (string.IsNullOrWhiteSpace(accountEndpoint)) + { throw new ArgumentException("Cannot be null or whitespace", nameof(accountEndpoint)); + } + if (string.IsNullOrWhiteSpace(databaseId)) + { throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); + } + if (string.IsNullOrWhiteSpace(containerId)) + { throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); + } options.ChatMessageStoreFactory = context => new CosmosChatMessageStore(accountEndpoint, new DefaultAzureCredential(), databaseId, containerId); return options; @@ -93,15 +115,26 @@ public static ChatClientAgentOptions WithCosmosDBMessageStore( string containerId) { if (options is null) + { throw new ArgumentNullException(nameof(options)); + } + if (cosmosClient is null) + { throw new ArgumentNullException(nameof(cosmosClient)); + } + if (string.IsNullOrWhiteSpace(databaseId)) + { throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); + } + if (string.IsNullOrWhiteSpace(containerId)) + { throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); + } options.ChatMessageStoreFactory = context => new CosmosChatMessageStore(cosmosClient, databaseId, containerId); return options; } -} \ No newline at end of file +} diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBWorkflowExtensions.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBWorkflowExtensions.cs index cc55a7081b..013a13a266 100644 --- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBWorkflowExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBWorkflowExtensions.cs @@ -30,11 +30,19 @@ public static CosmosCheckpointStore CreateCheckpointStore( string containerId) { if (string.IsNullOrWhiteSpace(connectionString)) + { throw new ArgumentException("Cannot be null or whitespace", nameof(connectionString)); + } + if (string.IsNullOrWhiteSpace(databaseId)) + { throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); + } + if (string.IsNullOrWhiteSpace(containerId)) + { throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); + } return new CosmosCheckpointStore(connectionString, databaseId, containerId); } @@ -55,11 +63,19 @@ public static CosmosCheckpointStore CreateCheckpointStoreUsingManagedIdentity( string containerId) { if (string.IsNullOrWhiteSpace(accountEndpoint)) + { throw new ArgumentException("Cannot be null or whitespace", nameof(accountEndpoint)); + } + if (string.IsNullOrWhiteSpace(databaseId)) + { throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); + } + if (string.IsNullOrWhiteSpace(containerId)) + { throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); + } return new CosmosCheckpointStore(accountEndpoint, new DefaultAzureCredential(), databaseId, containerId); } @@ -81,11 +97,19 @@ public static CosmosCheckpointStore CreateCheckpointStore( string containerId) { if (cosmosClient is null) + { throw new ArgumentNullException(nameof(cosmosClient)); + } + if (string.IsNullOrWhiteSpace(databaseId)) + { throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); + } + if (string.IsNullOrWhiteSpace(containerId)) + { throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); + } return new CosmosCheckpointStore(cosmosClient, databaseId, containerId); } @@ -107,11 +131,19 @@ public static CosmosCheckpointStore CreateCheckpointStore( string containerId) { if (string.IsNullOrWhiteSpace(connectionString)) + { throw new ArgumentException("Cannot be null or whitespace", nameof(connectionString)); + } + if (string.IsNullOrWhiteSpace(databaseId)) + { throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); + } + if (string.IsNullOrWhiteSpace(containerId)) + { throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); + } return new CosmosCheckpointStore(connectionString, databaseId, containerId); } @@ -133,11 +165,19 @@ public static CosmosCheckpointStore CreateCheckpointStoreUsingManagedIdentity string containerId) { if (string.IsNullOrWhiteSpace(accountEndpoint)) + { throw new ArgumentException("Cannot be null or whitespace", nameof(accountEndpoint)); + } + if (string.IsNullOrWhiteSpace(databaseId)) + { throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); + } + if (string.IsNullOrWhiteSpace(containerId)) + { throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); + } return new CosmosCheckpointStore(accountEndpoint, new DefaultAzureCredential(), databaseId, containerId); } @@ -160,12 +200,20 @@ public static CosmosCheckpointStore CreateCheckpointStore( string containerId) { if (cosmosClient is null) + { throw new ArgumentNullException(nameof(cosmosClient)); + } + if (string.IsNullOrWhiteSpace(databaseId)) + { throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); + } + if (string.IsNullOrWhiteSpace(containerId)) + { throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); + } return new CosmosCheckpointStore(cosmosClient, databaseId, containerId); } -} \ No newline at end of file +} From e8c0e271235448f3e74d2e367d65639a3a9b2d81 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Sat, 8 Nov 2025 16:29:33 +0000 Subject: [PATCH 29/40] Fix file encoding to UTF-8 with BOM, fix import ordering, and remove unnecessary using directives --- .../CosmosChatMessageStore.cs | 2 +- .../Microsoft.Agents.AI.CosmosNoSql/CosmosCheckpointStore.cs | 3 +-- .../CosmosDBChatExtensions.cs | 2 +- .../CosmosDBWorkflowExtensions.cs | 5 ++--- 4 files changed, 5 insertions(+), 7 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs index bd002b6439..94314764a5 100644 --- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs @@ -1,4 +1,4 @@ -// Copyright (c) Microsoft. All rights reserved. +// Copyright (c) Microsoft. All rights reserved. using System; using System.Collections.Generic; diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosCheckpointStore.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosCheckpointStore.cs index b051a6501c..62987b1dfc 100644 --- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosCheckpointStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosCheckpointStore.cs @@ -1,4 +1,4 @@ -// Copyright (c) Microsoft. All rights reserved. +// Copyright (c) Microsoft. All rights reserved. using System; using System.Collections.Generic; @@ -7,7 +7,6 @@ using System.Text.Json; using System.Threading.Tasks; using Azure.Core; -using Azure.Identity; using Microsoft.Azure.Cosmos; using Microsoft.Shared.Diagnostics; using Newtonsoft.Json; diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBChatExtensions.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBChatExtensions.cs index f20a6982cb..5878aab5a6 100644 --- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBChatExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBChatExtensions.cs @@ -1,4 +1,4 @@ -// Copyright (c) Microsoft. All rights reserved. +// Copyright (c) Microsoft. All rights reserved. using System; using System.Diagnostics.CodeAnalysis; diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBWorkflowExtensions.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBWorkflowExtensions.cs index 013a13a266..9d8bc52e68 100644 --- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBWorkflowExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBWorkflowExtensions.cs @@ -1,11 +1,10 @@ -// Copyright (c) Microsoft. All rights reserved. +// Copyright (c) Microsoft. All rights reserved. using System; using System.Diagnostics.CodeAnalysis; -using System.Text.Json; using Azure.Identity; -using Microsoft.Azure.Cosmos; using Microsoft.Agents.AI.Workflows.Checkpointing; +using Microsoft.Azure.Cosmos; namespace Microsoft.Agents.AI.Workflows; From ec0a45493033d44241c7367656f65ecb10722db1 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Sat, 8 Nov 2025 16:47:13 +0000 Subject: [PATCH 30/40] Convert backing fields to auto-properties and remove Azure.Identity using directive --- .../CosmosChatMessageStore.cs | 40 +++++++++---------- 1 file changed, 18 insertions(+), 22 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs index 94314764a5..8abb2fca17 100644 --- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs @@ -1,4 +1,4 @@ -// Copyright (c) Microsoft. All rights reserved. +// Copyright (c) Microsoft. All rights reserved. using System; using System.Collections.Generic; @@ -8,7 +8,6 @@ using System.Threading; using System.Threading.Tasks; using Azure.Core; -using Azure.Identity; using Microsoft.Azure.Cosmos; using Microsoft.Extensions.AI; using Microsoft.Shared.Diagnostics; @@ -24,9 +23,6 @@ public sealed class CosmosChatMessageStore : ChatMessageStore, IDisposable { private readonly CosmosClient _cosmosClient; private readonly Container _container; - private readonly string _conversationId; - private readonly string _databaseId; - private readonly string _containerId; private readonly bool _ownsClient; private bool _disposed; @@ -79,17 +75,17 @@ private static JsonSerializerOptions CreateDefaultJsonOptions() /// /// Gets the conversation ID associated with this message store. /// - public string ConversationId => this._conversationId; + public string ConversationId { get; init; } /// /// Gets the database ID associated with this message store. /// - public string DatabaseId => this._databaseId; + public string DatabaseId { get; init; } /// /// Gets the container ID associated with this message store. /// - public string ContainerId => this._containerId; + public string ContainerId { get; init; } /// /// Internal primary constructor used by all public constructors. @@ -105,9 +101,9 @@ internal CosmosChatMessageStore(CosmosClient cosmosClient, string databaseId, st { this._cosmosClient = Throw.IfNull(cosmosClient); this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); - this._conversationId = Throw.IfNullOrWhitespace(conversationId); - this._databaseId = databaseId; - this._containerId = containerId; + this.ConversationId = Throw.IfNullOrWhitespace(conversationId); + this.DatabaseId = databaseId; + this.ContainerId = containerId; this._ownsClient = ownsClient; // Initialize partitioning mode @@ -269,8 +265,8 @@ public CosmosChatMessageStore(CosmosClient cosmosClient, string databaseId, stri public CosmosChatMessageStore(CosmosClient cosmosClient, JsonElement serializedStoreState, string databaseId, string containerId, JsonSerializerOptions? jsonSerializerOptions = null) { this._cosmosClient = cosmosClient ?? throw new ArgumentNullException(nameof(cosmosClient)); - this._databaseId = Throw.IfNullOrWhitespace(databaseId); - this._containerId = Throw.IfNullOrWhitespace(containerId); + this.DatabaseId = Throw.IfNullOrWhitespace(databaseId); + this.ContainerId = Throw.IfNullOrWhitespace(containerId); this._container = this._cosmosClient.GetContainer(databaseId, containerId); this._ownsClient = false; @@ -279,7 +275,7 @@ public CosmosChatMessageStore(CosmosClient cosmosClient, JsonElement serializedS var state = JsonSerializer.Deserialize(serializedStoreState, jsonSerializerOptions); if (state?.ConversationIdentifier is { } conversationId) { - this._conversationId = conversationId; + this.ConversationId = conversationId; // Initialize hierarchical partitioning if available in state this._tenantId = state.TenantId; @@ -314,7 +310,7 @@ public override async Task> GetMessagesAsync(Cancellati // Fetch most recent messages in descending order when limit is set, then reverse to ascending var orderDirection = this.MaxMessagesToRetrieve.HasValue ? "DESC" : "ASC"; var query = new QueryDefinition($"SELECT * FROM c WHERE c.conversationId = @conversationId AND c.type = @type ORDER BY c.timestamp {orderDirection}") - .WithParameter("@conversationId", this._conversationId) + .WithParameter("@conversationId", this.ConversationId) .WithParameter("@type", "ChatMessage"); var iterator = this._container.GetItemQueryIterator(query, requestOptions: new QueryRequestOptions @@ -474,7 +470,7 @@ private CosmosMessageDocument CreateMessageDocument(ChatMessage message, long ti return new CosmosMessageDocument { Id = Guid.NewGuid().ToString(), - ConversationId = this._conversationId, + ConversationId = this.ConversationId, Timestamp = timestamp, MessageId = message.MessageId, Role = message.Role.Value, @@ -484,7 +480,7 @@ private CosmosMessageDocument CreateMessageDocument(ChatMessage message, long ti // Include hierarchical metadata when using hierarchical partitioning TenantId = this._useHierarchicalPartitioning ? this._tenantId : null, UserId = this._useHierarchicalPartitioning ? this._userId : null, - SessionId = this._useHierarchicalPartitioning ? this._conversationId : null + SessionId = this._useHierarchicalPartitioning ? this.ConversationId : null }; } @@ -500,7 +496,7 @@ public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptio var state = new StoreState { - ConversationIdentifier = this._conversationId, + ConversationIdentifier = this.ConversationId, TenantId = this._tenantId, UserId = this._userId, UseHierarchicalPartitioning = this._useHierarchicalPartitioning @@ -527,7 +523,7 @@ public async Task GetMessageCountAsync(CancellationToken cancellationToken // Efficient count query var query = new QueryDefinition("SELECT VALUE COUNT(1) FROM c WHERE c.conversationId = @conversationId AND c.Type = @type") - .WithParameter("@conversationId", this._conversationId) + .WithParameter("@conversationId", this.ConversationId) .WithParameter("@type", "ChatMessage"); var iterator = this._container.GetItemQueryIterator(query, requestOptions: new QueryRequestOptions @@ -561,17 +557,17 @@ public async Task ClearMessagesAsync(CancellationToken cancellationToken = // Batch delete for efficiency var query = new QueryDefinition("SELECT VALUE c.id FROM c WHERE c.conversationId = @conversationId AND c.Type = @type") - .WithParameter("@conversationId", this._conversationId) + .WithParameter("@conversationId", this.ConversationId) .WithParameter("@type", "ChatMessage"); var iterator = this._container.GetItemQueryIterator(query, requestOptions: new QueryRequestOptions { - PartitionKey = new PartitionKey(this._conversationId), + PartitionKey = new PartitionKey(this.ConversationId), MaxItemCount = this.MaxItemCount }); var deletedCount = 0; - var partitionKey = new PartitionKey(this._conversationId); + var partitionKey = new PartitionKey(this.ConversationId); while (iterator.HasMoreResults) { From 0b49065fc1d296cfd3d2fb5eb67758349c61deb2 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Sat, 8 Nov 2025 16:57:37 +0000 Subject: [PATCH 31/40] Fix CosmosChatMessageStore.cs encoding back to UTF-8 with BOM --- .../Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs index 8abb2fca17..2fb0377bc1 100644 --- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs @@ -1,4 +1,4 @@ -// Copyright (c) Microsoft. All rights reserved. +// Copyright (c) Microsoft. All rights reserved. using System; using System.Collections.Generic; From e6f4b79d6b9049ced4dc57b5e072a11ebd62aa68 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Sat, 8 Nov 2025 17:40:06 +0000 Subject: [PATCH 32/40] Fix test file formatting: indentation, encoding, imports, this. qualifications, naming conventions, and simplify new expressions --- .../CosmosChatMessageStoreTests.cs | 44 +++--- .../CosmosCheckpointStoreTests.cs | 131 +++++++++--------- .../CosmosDBCollectionFixture.cs | 2 +- 3 files changed, 88 insertions(+), 89 deletions(-) diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs index 51795e3059..9d486a1bfe 100644 --- a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs @@ -1,4 +1,4 @@ -// Copyright (c) Microsoft. All rights reserved. +// Copyright (c) Microsoft. All rights reserved. using System; using System.Collections.Generic; @@ -8,9 +8,9 @@ using System.Threading.Tasks; using Azure.Core; using Azure.Identity; +using Microsoft.Agents.AI; using Microsoft.Azure.Cosmos; using Microsoft.Extensions.AI; -using Microsoft.Agents.AI; using Xunit; namespace Microsoft.Agents.AI.CosmosNoSql.UnitTests; @@ -44,7 +44,7 @@ namespace Microsoft.Agents.AI.CosmosNoSql.UnitTests; public sealed class CosmosChatMessageStoreTests : IAsyncLifetime, IDisposable { // Cosmos DB Emulator connection settings - private const string EmulatorEndpoint = "https://localhost:8081"; + private const string s_emulatorEndpoint = "https://localhost:8081"; private const string EmulatorKey = "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw=="; private const string TestContainerId = "ChatMessages"; private const string HierarchicalTestContainerId = "HierarchicalChatMessages"; @@ -64,12 +64,12 @@ public async Task InitializeAsync() // Set COSMOS_PRESERVE_CONTAINERS=true to keep containers and data for inspection this._preserveContainer = string.Equals(Environment.GetEnvironmentVariable("COSMOS_PRESERVE_CONTAINERS"), "true", StringComparison.OrdinalIgnoreCase); - this._connectionString = $"AccountEndpoint={EmulatorEndpoint};AccountKey={EmulatorKey}"; + this._connectionString = $"AccountEndpoint={s_emulatorEndpoint};AccountKey={EmulatorKey}"; try { // Only create CosmosClient for test setup - the actual tests will use connection string constructors - this._setupClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); + this._setupClient = new CosmosClient(s_emulatorEndpoint, EmulatorKey); // Test connection by attempting to create database var databaseResponse = await this._setupClient.CreateDatabaseIfNotExistsAsync(TestDatabaseId); @@ -134,14 +134,14 @@ public void Dispose() GC.SuppressFinalize(this); } - private void SkipIfEmulatorNotAvailable() - { - // In CI: Skip if COSMOS_EMULATOR_AVAILABLE is not set to "true" - // Locally: Skip if emulator connection check failed - var ciEmulatorAvailable = string.Equals(Environment.GetEnvironmentVariable("COSMOS_EMULATOR_AVAILABLE"), "true", StringComparison.OrdinalIgnoreCase); + private void SkipIfEmulatorNotAvailable() + { + // In CI: Skip if COSMOS_EMULATOR_AVAILABLE is not set to "true" + // Locally: Skip if emulator connection check failed + var ciEmulatorAvailable = string.Equals(Environment.GetEnvironmentVariable("COSMOS_EMULATOR_AVAILABLE"), "true", StringComparison.OrdinalIgnoreCase); - Xunit.Skip.If(!ciEmulatorAvailable && !this._emulatorAvailable, "Cosmos DB Emulator is not available"); - } + Xunit.Skip.If(!ciEmulatorAvailable && !this._emulatorAvailable, "Cosmos DB Emulator is not available"); + } #region Constructor Tests @@ -247,7 +247,7 @@ public async Task AddMessagesAsync_WithSingleMessage_ShouldAddMessageAsync() PartitionKey = new PartitionKey(conversationId) }); - List rawResults = new List(); + List rawResults = new(); while (rawIterator.HasMoreResults) { var rawResponse = await rawIterator.ReadNextAsync(); @@ -268,7 +268,7 @@ public async Task AddMessagesAsync_WithSingleMessage_ShouldAddMessageAsync() public async Task AddMessagesAsync_WithMultipleMessages_ShouldAddAllMessagesAsync() { // Arrange - SkipIfEmulatorNotAvailable(); + this.SkipIfEmulatorNotAvailable(); var conversationId = Guid.NewGuid().ToString(); using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, conversationId); var messages = new[] @@ -299,7 +299,7 @@ public async Task AddMessagesAsync_WithMultipleMessages_ShouldAddAllMessagesAsyn public async Task GetMessagesAsync_WithNoMessages_ShouldReturnEmptyAsync() { // Arrange - SkipIfEmulatorNotAvailable(); + this.SkipIfEmulatorNotAvailable(); using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, Guid.NewGuid().ToString()); // Act @@ -314,7 +314,7 @@ public async Task GetMessagesAsync_WithNoMessages_ShouldReturnEmptyAsync() public async Task GetMessagesAsync_WithConversationIsolation_ShouldOnlyReturnMessagesForConversationAsync() { // Arrange - SkipIfEmulatorNotAvailable(); + this.SkipIfEmulatorNotAvailable(); var conversation1 = Guid.NewGuid().ToString(); var conversation2 = Guid.NewGuid().ToString(); @@ -346,7 +346,7 @@ public async Task GetMessagesAsync_WithConversationIsolation_ShouldOnlyReturnMes public async Task FullWorkflow_AddAndGet_ShouldWorkCorrectlyAsync() { // Arrange - SkipIfEmulatorNotAvailable(); + this.SkipIfEmulatorNotAvailable(); var conversationId = $"test-conversation-{Guid.NewGuid():N}"; // Use unique conversation ID using var originalStore = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, conversationId); @@ -390,7 +390,7 @@ public async Task FullWorkflow_AddAndGet_ShouldWorkCorrectlyAsync() public void Dispose_AfterUse_ShouldNotThrow() { // Arrange - SkipIfEmulatorNotAvailable(); + this.SkipIfEmulatorNotAvailable(); var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, Guid.NewGuid().ToString()); // Act & Assert @@ -402,7 +402,7 @@ public void Dispose_AfterUse_ShouldNotThrow() public void Dispose_MultipleCalls_ShouldNotThrow() { // Arrange - SkipIfEmulatorNotAvailable(); + this.SkipIfEmulatorNotAvailable(); var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, Guid.NewGuid().ToString()); // Act & Assert @@ -440,7 +440,7 @@ public void Constructor_WithHierarchicalEndpoint_ShouldCreateInstance() // Act TokenCredential credential = new DefaultAzureCredential(); - using var store = new CosmosChatMessageStore(EmulatorEndpoint, credential, TestDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789"); + using var store = new CosmosChatMessageStore(s_emulatorEndpoint, credential, TestDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789"); // Assert Assert.NotNull(store); @@ -456,7 +456,7 @@ public void Constructor_WithHierarchicalCosmosClient_ShouldCreateInstance() // Arrange & Act this.SkipIfEmulatorNotAvailable(); - using var cosmosClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); + using var cosmosClient = new CosmosClient(s_emulatorEndpoint, EmulatorKey); using var store = new CosmosChatMessageStore(cosmosClient, TestDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789"); // Assert @@ -635,7 +635,7 @@ public async Task SerializeDeserialize_WithHierarchicalPartitioning_ShouldPreser var serializedState = originalStore.Serialize(); // Create a new store from the serialized state - using var cosmosClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); + using var cosmosClient = new CosmosClient(s_emulatorEndpoint, EmulatorKey); var serializerOptions = new JsonSerializerOptions { TypeInfoResolver = new DefaultJsonTypeInfoResolver() diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs index 224cdc28bc..da820e4fa1 100644 --- a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs @@ -1,12 +1,12 @@ -// Copyright (c) Microsoft. All rights reserved. +// Copyright (c) Microsoft. All rights reserved. using System; using System.Linq; using System.Text.Json; using System.Threading.Tasks; -using Microsoft.Azure.Cosmos; using Microsoft.Agents.AI.Workflows; using Microsoft.Agents.AI.Workflows.Checkpointing; +using Microsoft.Azure.Cosmos; using Xunit; namespace Microsoft.Agents.AI.CosmosNoSql.UnitTests; @@ -30,7 +30,7 @@ namespace Microsoft.Agents.AI.CosmosNoSql.UnitTests; public class CosmosCheckpointStoreTests : IAsyncLifetime, IDisposable { // Cosmos DB Emulator connection settings - private const string EmulatorEndpoint = "https://localhost:8081"; + private const string s_emulatorEndpoint = "https://localhost:8081"; private const string EmulatorKey = "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw=="; private const string TestContainerId = "Checkpoints"; // Use unique database ID per test class instance to avoid conflicts @@ -41,7 +41,6 @@ public class CosmosCheckpointStoreTests : IAsyncLifetime, IDisposable private string _connectionString = string.Empty; private CosmosClient? _cosmosClient; private Database? _database; - private Container? _container; private bool _emulatorAvailable; private bool _preserveContainer; @@ -61,39 +60,39 @@ public async Task InitializeAsync() { // Check environment variable to determine if we should preserve containers // Set COSMOS_PRESERVE_CONTAINERS=true to keep containers and data for inspection - _preserveContainer = string.Equals(Environment.GetEnvironmentVariable("COSMOS_PRESERVE_CONTAINERS"), "true", StringComparison.OrdinalIgnoreCase); + this._preserveContainer = string.Equals(Environment.GetEnvironmentVariable("COSMOS_PRESERVE_CONTAINERS"), "true", StringComparison.OrdinalIgnoreCase); - _connectionString = $"AccountEndpoint={EmulatorEndpoint};AccountKey={EmulatorKey}"; + this._connectionString = $"AccountEndpoint={s_emulatorEndpoint};AccountKey={EmulatorKey}"; try { - _cosmosClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); + this._cosmosClient = new CosmosClient(s_emulatorEndpoint, EmulatorKey); // Test connection by attempting to create database - _database = await _cosmosClient.CreateDatabaseIfNotExistsAsync(TestDatabaseId); - _container = await _database.CreateContainerIfNotExistsAsync( + this._database = await this._cosmosClient.CreateDatabaseIfNotExistsAsync(TestDatabaseId); + await this._database.CreateContainerIfNotExistsAsync( TestContainerId, "/runId", throughput: 400); - _emulatorAvailable = true; + this._emulatorAvailable = true; } catch (Exception ex) when (!(ex is OutOfMemoryException || ex is StackOverflowException || ex is AccessViolationException)) { // Emulator not available, tests will be skipped - _emulatorAvailable = false; - _cosmosClient?.Dispose(); - _cosmosClient = null; + this._emulatorAvailable = false; + this._cosmosClient?.Dispose(); + this._cosmosClient = null; } } public async Task DisposeAsync() { - if (_cosmosClient != null && _emulatorAvailable) + if (this._cosmosClient != null && this._emulatorAvailable) { try { - if (_preserveContainer) + if (this._preserveContainer) { // Preserve mode: Don't delete the database/container, keep data for inspection // This allows viewing data in the Cosmos DB Emulator Data Explorer @@ -102,7 +101,7 @@ public async Task DisposeAsync() else { // Clean mode: Delete the test database and all data - await _database!.DeleteAsync(); + await this._database!.DeleteAsync(); } } catch (Exception ex) @@ -112,19 +111,19 @@ public async Task DisposeAsync() } finally { - _cosmosClient.Dispose(); + this._cosmosClient.Dispose(); } } } - private void SkipIfEmulatorNotAvailable() - { - // In CI: Skip if COSMOS_EMULATOR_AVAILABLE is not set to "true" - // Locally: Skip if emulator connection check failed - var ciEmulatorAvailable = string.Equals(Environment.GetEnvironmentVariable("COSMOS_EMULATOR_AVAILABLE"), "true", StringComparison.OrdinalIgnoreCase); + private void SkipIfEmulatorNotAvailable() + { + // In CI: Skip if COSMOS_EMULATOR_AVAILABLE is not set to "true" + // Locally: Skip if emulator connection check failed + var ciEmulatorAvailable = string.Equals(Environment.GetEnvironmentVariable("COSMOS_EMULATOR_AVAILABLE"), "true", StringComparison.OrdinalIgnoreCase); - Xunit.Skip.If(!ciEmulatorAvailable && !this._emulatorAvailable, "Cosmos DB Emulator is not available"); - } + Xunit.Skip.If(!ciEmulatorAvailable && !this._emulatorAvailable, "Cosmos DB Emulator is not available"); + } #region Constructor Tests @@ -132,10 +131,10 @@ private void SkipIfEmulatorNotAvailable() public void Constructor_WithCosmosClient_SetsProperties() { // Arrange - SkipIfEmulatorNotAvailable(); + this.SkipIfEmulatorNotAvailable(); // Act - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(this._cosmosClient!, TestDatabaseId, TestContainerId); // Assert Assert.Equal(TestDatabaseId, store.DatabaseId); @@ -146,10 +145,10 @@ public void Constructor_WithCosmosClient_SetsProperties() public void Constructor_WithConnectionString_SetsProperties() { // Arrange - SkipIfEmulatorNotAvailable(); + this.SkipIfEmulatorNotAvailable(); // Act - using var store = new CosmosCheckpointStore(_connectionString, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(this._connectionString, TestDatabaseId, TestContainerId); // Assert Assert.Equal(TestDatabaseId, store.DatabaseId); @@ -177,12 +176,12 @@ public void Constructor_WithNullConnectionString_ThrowsArgumentException() #region Checkpoint Operations Tests [SkippableFact] - public async Task CreateCheckpointAsync_NewCheckpoint_CreatesSuccessfully() + public async Task CreateCheckpointAsync_NewCheckpoint_CreatesSuccessfullyAsync() { - SkipIfEmulatorNotAvailable(); + this.SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(this._cosmosClient!, TestDatabaseId, TestContainerId); var runId = Guid.NewGuid().ToString(); var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test checkpoint" }, s_jsonOptions); @@ -197,12 +196,12 @@ public async Task CreateCheckpointAsync_NewCheckpoint_CreatesSuccessfully() } [SkippableFact] - public async Task RetrieveCheckpointAsync_ExistingCheckpoint_ReturnsCorrectValue() + public async Task RetrieveCheckpointAsync_ExistingCheckpoint_ReturnsCorrectValueAsync() { - SkipIfEmulatorNotAvailable(); + this.SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(this._cosmosClient!, TestDatabaseId, TestContainerId); var runId = Guid.NewGuid().ToString(); var originalData = new { message = "Hello, World!", timestamp = DateTimeOffset.UtcNow }; var checkpointValue = JsonSerializer.SerializeToElement(originalData, s_jsonOptions); @@ -218,12 +217,12 @@ public async Task RetrieveCheckpointAsync_ExistingCheckpoint_ReturnsCorrectValue } [SkippableFact] - public async Task RetrieveCheckpointAsync_NonExistentCheckpoint_ThrowsInvalidOperationException() + public async Task RetrieveCheckpointAsync_NonExistentCheckpoint_ThrowsInvalidOperationExceptionAsync() { - SkipIfEmulatorNotAvailable(); + this.SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(this._cosmosClient!, TestDatabaseId, TestContainerId); var runId = Guid.NewGuid().ToString(); var fakeCheckpointInfo = new CheckpointInfo(runId, "nonexistent-checkpoint"); @@ -233,12 +232,12 @@ await Assert.ThrowsAsync(() => } [SkippableFact] - public async Task RetrieveIndexAsync_EmptyStore_ReturnsEmptyCollection() + public async Task RetrieveIndexAsync_EmptyStore_ReturnsEmptyCollectionAsync() { - SkipIfEmulatorNotAvailable(); + this.SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(this._cosmosClient!, TestDatabaseId, TestContainerId); var runId = Guid.NewGuid().ToString(); // Act @@ -250,12 +249,12 @@ public async Task RetrieveIndexAsync_EmptyStore_ReturnsEmptyCollection() } [SkippableFact] - public async Task RetrieveIndexAsync_WithCheckpoints_ReturnsAllCheckpoints() + public async Task RetrieveIndexAsync_WithCheckpoints_ReturnsAllCheckpointsAsync() { - SkipIfEmulatorNotAvailable(); + this.SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(this._cosmosClient!, TestDatabaseId, TestContainerId); var runId = Guid.NewGuid().ToString(); var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); @@ -275,12 +274,12 @@ public async Task RetrieveIndexAsync_WithCheckpoints_ReturnsAllCheckpoints() } [SkippableFact] - public async Task CreateCheckpointAsync_WithParent_CreatesHierarchy() + public async Task CreateCheckpointAsync_WithParent_CreatesHierarchyAsync() { - SkipIfEmulatorNotAvailable(); + this.SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(this._cosmosClient!, TestDatabaseId, TestContainerId); var runId = Guid.NewGuid().ToString(); var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); @@ -295,12 +294,12 @@ public async Task CreateCheckpointAsync_WithParent_CreatesHierarchy() } [SkippableFact] - public async Task RetrieveIndexAsync_WithParentFilter_ReturnsFilteredResults() + public async Task RetrieveIndexAsync_WithParentFilter_ReturnsFilteredResultsAsync() { - SkipIfEmulatorNotAvailable(); + this.SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(this._cosmosClient!, TestDatabaseId, TestContainerId); var runId = Guid.NewGuid().ToString(); var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); @@ -331,12 +330,12 @@ public async Task RetrieveIndexAsync_WithParentFilter_ReturnsFilteredResults() #region Run Isolation Tests [SkippableFact] - public async Task CheckpointOperations_DifferentRuns_IsolatesData() + public async Task CheckpointOperations_DifferentRuns_IsolatesDataAsync() { - SkipIfEmulatorNotAvailable(); + this.SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(this._cosmosClient!, TestDatabaseId, TestContainerId); var runId1 = Guid.NewGuid().ToString(); var runId2 = Guid.NewGuid().ToString(); var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); @@ -361,12 +360,12 @@ public async Task CheckpointOperations_DifferentRuns_IsolatesData() #region Error Handling Tests [SkippableFact] - public async Task CreateCheckpointAsync_WithNullRunId_ThrowsArgumentException() + public async Task CreateCheckpointAsync_WithNullRunId_ThrowsArgumentExceptionAsync() { - SkipIfEmulatorNotAvailable(); + this.SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(this._cosmosClient!, TestDatabaseId, TestContainerId); var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); // Act & Assert @@ -375,12 +374,12 @@ await Assert.ThrowsAsync(() => } [SkippableFact] - public async Task CreateCheckpointAsync_WithEmptyRunId_ThrowsArgumentException() + public async Task CreateCheckpointAsync_WithEmptyRunId_ThrowsArgumentExceptionAsync() { - SkipIfEmulatorNotAvailable(); + this.SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(this._cosmosClient!, TestDatabaseId, TestContainerId); var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); // Act & Assert @@ -389,12 +388,12 @@ await Assert.ThrowsAsync(() => } [SkippableFact] - public async Task RetrieveCheckpointAsync_WithNullCheckpointInfo_ThrowsArgumentNullException() + public async Task RetrieveCheckpointAsync_WithNullCheckpointInfo_ThrowsArgumentNullExceptionAsync() { - SkipIfEmulatorNotAvailable(); + this.SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(this._cosmosClient!, TestDatabaseId, TestContainerId); var runId = Guid.NewGuid().ToString(); // Act & Assert @@ -407,12 +406,12 @@ await Assert.ThrowsAsync(() => #region Disposal Tests [SkippableFact] - public async Task Dispose_AfterDisposal_ThrowsObjectDisposedException() + public async Task Dispose_AfterDisposal_ThrowsObjectDisposedExceptionAsync() { - SkipIfEmulatorNotAvailable(); + this.SkipIfEmulatorNotAvailable(); // Arrange - var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + var store = new CosmosCheckpointStore(this._cosmosClient!, TestDatabaseId, TestContainerId); var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); // Act @@ -426,10 +425,10 @@ await Assert.ThrowsAsync(() => [SkippableFact] public void Dispose_MultipleCalls_DoesNotThrow() { - SkipIfEmulatorNotAvailable(); + this.SkipIfEmulatorNotAvailable(); // Arrange - var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); + var store = new CosmosCheckpointStore(this._cosmosClient!, TestDatabaseId, TestContainerId); // Act & Assert (should not throw) store.Dispose(); diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosDBCollectionFixture.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosDBCollectionFixture.cs index c2567ed2ac..195c433de5 100644 --- a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosDBCollectionFixture.cs +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosDBCollectionFixture.cs @@ -1,4 +1,4 @@ -// Copyright (c) Microsoft. All rights reserved. +// Copyright (c) Microsoft. All rights reserved. using Xunit; From 1d32377ac58b890c4ecf9b7b3a6b9e379167e894 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Sat, 8 Nov 2025 18:43:51 +0000 Subject: [PATCH 33/40] Fix const field naming violations: Remove s_ prefix from const fields and add this. qualification to Dispose call --- .../CosmosChatMessageStoreTests.cs | 14 +++++++------- .../CosmosCheckpointStoreTests.cs | 10 +++++----- 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs index 9d486a1bfe..c7a295b34a 100644 --- a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs @@ -44,7 +44,7 @@ namespace Microsoft.Agents.AI.CosmosNoSql.UnitTests; public sealed class CosmosChatMessageStoreTests : IAsyncLifetime, IDisposable { // Cosmos DB Emulator connection settings - private const string s_emulatorEndpoint = "https://localhost:8081"; + private const string EmulatorEndpoint = "https://localhost:8081"; private const string EmulatorKey = "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw=="; private const string TestContainerId = "ChatMessages"; private const string HierarchicalTestContainerId = "HierarchicalChatMessages"; @@ -64,12 +64,12 @@ public async Task InitializeAsync() // Set COSMOS_PRESERVE_CONTAINERS=true to keep containers and data for inspection this._preserveContainer = string.Equals(Environment.GetEnvironmentVariable("COSMOS_PRESERVE_CONTAINERS"), "true", StringComparison.OrdinalIgnoreCase); - this._connectionString = $"AccountEndpoint={s_emulatorEndpoint};AccountKey={EmulatorKey}"; + this._connectionString = $"AccountEndpoint={EmulatorEndpoint};AccountKey={EmulatorKey}"; try { // Only create CosmosClient for test setup - the actual tests will use connection string constructors - this._setupClient = new CosmosClient(s_emulatorEndpoint, EmulatorKey); + this._setupClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); // Test connection by attempting to create database var databaseResponse = await this._setupClient.CreateDatabaseIfNotExistsAsync(TestDatabaseId); @@ -440,7 +440,7 @@ public void Constructor_WithHierarchicalEndpoint_ShouldCreateInstance() // Act TokenCredential credential = new DefaultAzureCredential(); - using var store = new CosmosChatMessageStore(s_emulatorEndpoint, credential, TestDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789"); + using var store = new CosmosChatMessageStore(EmulatorEndpoint, credential, TestDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789"); // Assert Assert.NotNull(store); @@ -456,7 +456,7 @@ public void Constructor_WithHierarchicalCosmosClient_ShouldCreateInstance() // Arrange & Act this.SkipIfEmulatorNotAvailable(); - using var cosmosClient = new CosmosClient(s_emulatorEndpoint, EmulatorKey); + using var cosmosClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); using var store = new CosmosChatMessageStore(cosmosClient, TestDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789"); // Assert @@ -635,7 +635,7 @@ public async Task SerializeDeserialize_WithHierarchicalPartitioning_ShouldPreser var serializedState = originalStore.Serialize(); // Create a new store from the serialized state - using var cosmosClient = new CosmosClient(s_emulatorEndpoint, EmulatorKey); + using var cosmosClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); var serializerOptions = new JsonSerializerOptions { TypeInfoResolver = new DefaultJsonTypeInfoResolver() @@ -757,4 +757,4 @@ public async Task MaxMessagesToRetrieve_Null_ShouldReturnAllMessagesAsync() } #endregion -} +} \ No newline at end of file diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs index da820e4fa1..b429f999d5 100644 --- a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs @@ -30,7 +30,7 @@ namespace Microsoft.Agents.AI.CosmosNoSql.UnitTests; public class CosmosCheckpointStoreTests : IAsyncLifetime, IDisposable { // Cosmos DB Emulator connection settings - private const string s_emulatorEndpoint = "https://localhost:8081"; + private const string EmulatorEndpoint = "https://localhost:8081"; private const string EmulatorKey = "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw=="; private const string TestContainerId = "Checkpoints"; // Use unique database ID per test class instance to avoid conflicts @@ -62,11 +62,11 @@ public async Task InitializeAsync() // Set COSMOS_PRESERVE_CONTAINERS=true to keep containers and data for inspection this._preserveContainer = string.Equals(Environment.GetEnvironmentVariable("COSMOS_PRESERVE_CONTAINERS"), "true", StringComparison.OrdinalIgnoreCase); - this._connectionString = $"AccountEndpoint={s_emulatorEndpoint};AccountKey={EmulatorKey}"; + this._connectionString = $"AccountEndpoint={EmulatorEndpoint};AccountKey={EmulatorKey}"; try { - this._cosmosClient = new CosmosClient(s_emulatorEndpoint, EmulatorKey); + this._cosmosClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); // Test connection by attempting to create database this._database = await this._cosmosClient.CreateDatabaseIfNotExistsAsync(TestDatabaseId); @@ -440,7 +440,7 @@ public void Dispose_MultipleCalls_DoesNotThrow() public void Dispose() { - Dispose(true); + this.Dispose(true); GC.SuppressFinalize(this); } @@ -451,4 +451,4 @@ protected virtual void Dispose(bool disposing) this._cosmosClient?.Dispose(); } } -} +} \ No newline at end of file From b40a8cee4b41e7e77f6edb9ec2bdf07953ee4a88 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Sat, 8 Nov 2025 19:29:57 +0000 Subject: [PATCH 34/40] Add local .editorconfig for Cosmos DB tests to suppress IDE0005 false positives from multi-targeting --- .../.editorconfig | 9 +++++++++ 1 file changed, 9 insertions(+) create mode 100644 dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/.editorconfig diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/.editorconfig b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/.editorconfig new file mode 100644 index 0000000000..83e05f582a --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/.editorconfig @@ -0,0 +1,9 @@ +# EditorConfig overrides for Cosmos DB Unit Tests +# Multi-targeting (net472 + net9.0) causes false positives for IDE0005 (unnecessary using directives) + +root = false + +[*.cs] +# Suppress IDE0005 for this project - multi-targeting causes false positives +# These using directives ARE necessary but appear unnecessary in one target framework +dotnet_diagnostic.IDE0005.severity = none From 10aecaada65b09e33d1144116269c335f73022a2 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Sat, 8 Nov 2025 19:48:27 +0000 Subject: [PATCH 35/40] Fix IDE1006 naming violations: Rename TestDatabaseId to s_testDatabaseId and add final newlines --- .../CosmosChatMessageStoreTests.cs | 84 +++++++++---------- .../CosmosCheckpointStoreTests.cs | 44 +++++----- 2 files changed, 64 insertions(+), 64 deletions(-) diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs index c7a295b34a..6b62248a16 100644 --- a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs @@ -50,7 +50,7 @@ public sealed class CosmosChatMessageStoreTests : IAsyncLifetime, IDisposable private const string HierarchicalTestContainerId = "HierarchicalChatMessages"; // Use unique database ID per test class instance to avoid conflicts #pragma warning disable CA1802 // Use literals where appropriate - private static readonly string TestDatabaseId = $"AgentFrameworkTests-ChatStore-{Guid.NewGuid():N}"; + private static readonly string s_testDatabaseId = $"AgentFrameworkTests-ChatStore-{Guid.NewGuid():N}"; #pragma warning restore CA1802 private string _connectionString = string.Empty; @@ -72,7 +72,7 @@ public async Task InitializeAsync() this._setupClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); // Test connection by attempting to create database - var databaseResponse = await this._setupClient.CreateDatabaseIfNotExistsAsync(TestDatabaseId); + var databaseResponse = await this._setupClient.CreateDatabaseIfNotExistsAsync(s_testDatabaseId); // Create container for simple partitioning tests await databaseResponse.Database.CreateContainerIfNotExistsAsync( @@ -112,7 +112,7 @@ public async Task DisposeAsync() else { // Clean mode: Delete the test database and all data - var database = this._setupClient.GetDatabase(TestDatabaseId); + var database = this._setupClient.GetDatabase(s_testDatabaseId); await database.DeleteAsync(); } } @@ -153,12 +153,12 @@ public void Constructor_WithConnectionString_ShouldCreateInstance() this.SkipIfEmulatorNotAvailable(); // Act - using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, "test-conversation"); + using var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, "test-conversation"); // Assert Assert.NotNull(store); Assert.Equal("test-conversation", store.ConversationId); - Assert.Equal(TestDatabaseId, store.DatabaseId); + Assert.Equal(s_testDatabaseId, store.DatabaseId); Assert.Equal(TestContainerId, store.ContainerId); } @@ -170,12 +170,12 @@ public void Constructor_WithConnectionStringNoConversationId_ShouldCreateInstanc this.SkipIfEmulatorNotAvailable(); // Act - using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId); + using var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId); // Assert Assert.NotNull(store); Assert.NotNull(store.ConversationId); - Assert.Equal(TestDatabaseId, store.DatabaseId); + Assert.Equal(s_testDatabaseId, store.DatabaseId); Assert.Equal(TestContainerId, store.ContainerId); } @@ -185,7 +185,7 @@ public void Constructor_WithNullConnectionString_ShouldThrowArgumentException() { // Arrange & Act & Assert Assert.Throws(() => - new CosmosChatMessageStore((string)null!, TestDatabaseId, TestContainerId, "test-conversation")); + new CosmosChatMessageStore((string)null!, s_testDatabaseId, TestContainerId, "test-conversation")); } [SkippableFact] @@ -196,7 +196,7 @@ public void Constructor_WithEmptyConversationId_ShouldThrowArgumentException() this.SkipIfEmulatorNotAvailable(); Assert.Throws(() => - new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, "")); + new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, "")); } #endregion @@ -210,7 +210,7 @@ public async Task AddMessagesAsync_WithSingleMessage_ShouldAddMessageAsync() // Arrange this.SkipIfEmulatorNotAvailable(); var conversationId = Guid.NewGuid().ToString(); - using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, conversationId); + using var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, conversationId); var message = new ChatMessage(ChatRole.User, "Hello, world!"); // Act @@ -229,7 +229,7 @@ public async Task AddMessagesAsync_WithSingleMessage_ShouldAddMessageAsync() // Let's check if we can find ANY items in the container for this conversation var directQuery = new QueryDefinition("SELECT VALUE COUNT(1) FROM c WHERE c.conversationId = @conversationId") .WithParameter("@conversationId", conversationId); - var countIterator = this._setupClient!.GetDatabase(TestDatabaseId).GetContainer(TestContainerId) + var countIterator = this._setupClient!.GetDatabase(s_testDatabaseId).GetContainer(TestContainerId) .GetItemQueryIterator(directQuery, requestOptions: new QueryRequestOptions { PartitionKey = new PartitionKey(conversationId) @@ -241,7 +241,7 @@ public async Task AddMessagesAsync_WithSingleMessage_ShouldAddMessageAsync() // Debug: Let's see what the raw query returns var rawQuery = new QueryDefinition("SELECT * FROM c WHERE c.conversationId = @conversationId") .WithParameter("@conversationId", conversationId); - var rawIterator = this._setupClient!.GetDatabase(TestDatabaseId).GetContainer(TestContainerId) + var rawIterator = this._setupClient!.GetDatabase(s_testDatabaseId).GetContainer(TestContainerId) .GetItemQueryIterator(rawQuery, requestOptions: new QueryRequestOptions { PartitionKey = new PartitionKey(conversationId) @@ -270,7 +270,7 @@ public async Task AddMessagesAsync_WithMultipleMessages_ShouldAddAllMessagesAsyn // Arrange this.SkipIfEmulatorNotAvailable(); var conversationId = Guid.NewGuid().ToString(); - using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, conversationId); + using var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, conversationId); var messages = new[] { new ChatMessage(ChatRole.User, "First message"), @@ -300,7 +300,7 @@ public async Task GetMessagesAsync_WithNoMessages_ShouldReturnEmptyAsync() { // Arrange this.SkipIfEmulatorNotAvailable(); - using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, Guid.NewGuid().ToString()); + using var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, Guid.NewGuid().ToString()); // Act var messages = await store.GetMessagesAsync(); @@ -318,8 +318,8 @@ public async Task GetMessagesAsync_WithConversationIsolation_ShouldOnlyReturnMes var conversation1 = Guid.NewGuid().ToString(); var conversation2 = Guid.NewGuid().ToString(); - using var store1 = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, conversation1); - using var store2 = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, conversation2); + using var store1 = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, conversation1); + using var store2 = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, conversation2); await store1.AddMessagesAsync([new ChatMessage(ChatRole.User, "Message for conversation 1")]); await store2.AddMessagesAsync([new ChatMessage(ChatRole.User, "Message for conversation 2")]); @@ -348,7 +348,7 @@ public async Task FullWorkflow_AddAndGet_ShouldWorkCorrectlyAsync() // Arrange this.SkipIfEmulatorNotAvailable(); var conversationId = $"test-conversation-{Guid.NewGuid():N}"; // Use unique conversation ID - using var originalStore = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, conversationId); + using var originalStore = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, conversationId); var messages = new[] { @@ -368,7 +368,7 @@ public async Task FullWorkflow_AddAndGet_ShouldWorkCorrectlyAsync() Assert.Equal(5, retrievedList.Count); // Act 3: Create new store instance for same conversation (test persistence) - using var newStore = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, conversationId); + using var newStore = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, conversationId); var persistedMessages = await newStore.GetMessagesAsync(); var persistedList = persistedMessages.ToList(); @@ -391,7 +391,7 @@ public void Dispose_AfterUse_ShouldNotThrow() { // Arrange this.SkipIfEmulatorNotAvailable(); - var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, Guid.NewGuid().ToString()); + var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, Guid.NewGuid().ToString()); // Act & Assert store.Dispose(); // Should not throw @@ -403,7 +403,7 @@ public void Dispose_MultipleCalls_ShouldNotThrow() { // Arrange this.SkipIfEmulatorNotAvailable(); - var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, Guid.NewGuid().ToString()); + var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, Guid.NewGuid().ToString()); // Act & Assert store.Dispose(); // First call @@ -422,12 +422,12 @@ public void Constructor_WithHierarchicalConnectionString_ShouldCreateInstance() this.SkipIfEmulatorNotAvailable(); // Act - using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789"); + using var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789"); // Assert Assert.NotNull(store); Assert.Equal("session-789", store.ConversationId); - Assert.Equal(TestDatabaseId, store.DatabaseId); + Assert.Equal(s_testDatabaseId, store.DatabaseId); Assert.Equal(HierarchicalTestContainerId, store.ContainerId); } @@ -440,12 +440,12 @@ public void Constructor_WithHierarchicalEndpoint_ShouldCreateInstance() // Act TokenCredential credential = new DefaultAzureCredential(); - using var store = new CosmosChatMessageStore(EmulatorEndpoint, credential, TestDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789"); + using var store = new CosmosChatMessageStore(EmulatorEndpoint, credential, s_testDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789"); // Assert Assert.NotNull(store); Assert.Equal("session-789", store.ConversationId); - Assert.Equal(TestDatabaseId, store.DatabaseId); + Assert.Equal(s_testDatabaseId, store.DatabaseId); Assert.Equal(HierarchicalTestContainerId, store.ContainerId); } @@ -457,12 +457,12 @@ public void Constructor_WithHierarchicalCosmosClient_ShouldCreateInstance() this.SkipIfEmulatorNotAvailable(); using var cosmosClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); - using var store = new CosmosChatMessageStore(cosmosClient, TestDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789"); + using var store = new CosmosChatMessageStore(cosmosClient, s_testDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789"); // Assert Assert.NotNull(store); Assert.Equal("session-789", store.ConversationId); - Assert.Equal(TestDatabaseId, store.DatabaseId); + Assert.Equal(s_testDatabaseId, store.DatabaseId); Assert.Equal(HierarchicalTestContainerId, store.ContainerId); } @@ -474,7 +474,7 @@ public void Constructor_WithHierarchicalNullTenantId_ShouldThrowArgumentExceptio this.SkipIfEmulatorNotAvailable(); Assert.Throws(() => - new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, null!, "user-456", "session-789")); + new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, null!, "user-456", "session-789")); } [SkippableFact] @@ -485,7 +485,7 @@ public void Constructor_WithHierarchicalEmptyUserId_ShouldThrowArgumentException this.SkipIfEmulatorNotAvailable(); Assert.Throws(() => - new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, "tenant-123", "", "session-789")); + new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, "tenant-123", "", "session-789")); } [SkippableFact] @@ -496,7 +496,7 @@ public void Constructor_WithHierarchicalWhitespaceSessionId_ShouldThrowArgumentE this.SkipIfEmulatorNotAvailable(); Assert.Throws(() => - new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", " ")); + new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", " ")); } [SkippableFact] @@ -509,7 +509,7 @@ public async Task AddMessagesAsync_WithHierarchicalPartitioning_ShouldAddMessage const string UserId = "user-456"; const string SessionId = "session-789"; // Test hierarchical partitioning constructor with connection string - using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, TenantId, UserId, SessionId); + using var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, TenantId, UserId, SessionId); var message = new ChatMessage(ChatRole.User, "Hello from hierarchical partitioning!"); // Act @@ -531,7 +531,7 @@ public async Task AddMessagesAsync_WithHierarchicalPartitioning_ShouldAddMessage .WithParameter("@conversationId", SessionId) .WithParameter("@type", "ChatMessage"); - var iterator = this._setupClient!.GetDatabase(TestDatabaseId).GetContainer(HierarchicalTestContainerId) + var iterator = this._setupClient!.GetDatabase(s_testDatabaseId).GetContainer(HierarchicalTestContainerId) .GetItemQueryIterator(directQuery, requestOptions: new QueryRequestOptions { PartitionKey = new PartitionKeyBuilder().Add(TenantId).Add(UserId).Add(SessionId).Build() @@ -558,7 +558,7 @@ public async Task AddMessagesAsync_WithHierarchicalMultipleMessages_ShouldAddAll const string UserId = "user-batch"; const string SessionId = "session-batch"; // Test hierarchical partitioning constructor with connection string - using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, TenantId, UserId, SessionId); + using var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, TenantId, UserId, SessionId); var messages = new[] { new ChatMessage(ChatRole.User, "First hierarchical message"), @@ -594,8 +594,8 @@ public async Task GetMessagesAsync_WithHierarchicalPartitionIsolation_ShouldIsol const string SessionId = "session-isolation"; // Different userIds create different hierarchical partitions, providing proper isolation - using var store1 = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, TenantId, UserId1, SessionId); - using var store2 = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, TenantId, UserId2, SessionId); + using var store1 = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, TenantId, UserId1, SessionId); + using var store2 = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, TenantId, UserId2, SessionId); // Add messages to both stores await store1.AddMessagesAsync([new ChatMessage(ChatRole.User, "Message from user 1")]); @@ -628,7 +628,7 @@ public async Task SerializeDeserialize_WithHierarchicalPartitioning_ShouldPreser const string UserId = "user-serialize"; const string SessionId = "session-serialize"; - using var originalStore = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, TenantId, UserId, SessionId); + using var originalStore = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, TenantId, UserId, SessionId); await originalStore.AddMessagesAsync([new ChatMessage(ChatRole.User, "Test serialization message")]); // Act - Serialize the store state @@ -640,7 +640,7 @@ public async Task SerializeDeserialize_WithHierarchicalPartitioning_ShouldPreser { TypeInfoResolver = new DefaultJsonTypeInfoResolver() }; - using var deserializedStore = new CosmosChatMessageStore(cosmosClient, serializedState, TestDatabaseId, HierarchicalTestContainerId, serializerOptions); + using var deserializedStore = new CosmosChatMessageStore(cosmosClient, serializedState, s_testDatabaseId, HierarchicalTestContainerId, serializerOptions); // Wait a moment for eventual consistency await Task.Delay(100); @@ -652,7 +652,7 @@ public async Task SerializeDeserialize_WithHierarchicalPartitioning_ShouldPreser Assert.Single(messageList); Assert.Equal("Test serialization message", messageList[0].Text); Assert.Equal(SessionId, deserializedStore.ConversationId); - Assert.Equal(TestDatabaseId, deserializedStore.DatabaseId); + Assert.Equal(s_testDatabaseId, deserializedStore.DatabaseId); Assert.Equal(HierarchicalTestContainerId, deserializedStore.ContainerId); } @@ -665,8 +665,8 @@ public async Task HierarchicalAndSimplePartitioning_ShouldCoexistAsync() const string SessionId = "coexist-session"; // Create simple store using simple partitioning container and hierarchical store using hierarchical container - using var simpleStore = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, SessionId); - using var hierarchicalStore = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, "tenant-coexist", "user-coexist", SessionId); + using var simpleStore = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, SessionId); + using var hierarchicalStore = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, "tenant-coexist", "user-coexist", SessionId); // Add messages to both await simpleStore.AddMessagesAsync([new ChatMessage(ChatRole.User, "Simple partitioning message")]); @@ -697,7 +697,7 @@ public async Task MaxMessagesToRetrieve_ShouldLimitAndReturnMostRecentAsync() this.SkipIfEmulatorNotAvailable(); const string ConversationId = "max-messages-test"; - using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, ConversationId); + using var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, ConversationId); // Add 10 messages var messages = new List(); @@ -733,7 +733,7 @@ public async Task MaxMessagesToRetrieve_Null_ShouldReturnAllMessagesAsync() this.SkipIfEmulatorNotAvailable(); const string ConversationId = "max-messages-null-test"; - using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, ConversationId); + using var store = new CosmosChatMessageStore(this._connectionString, s_testDatabaseId, TestContainerId, ConversationId); // Add 10 messages var messages = new List(); @@ -757,4 +757,4 @@ public async Task MaxMessagesToRetrieve_Null_ShouldReturnAllMessagesAsync() } #endregion -} \ No newline at end of file +} diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs index b429f999d5..dfa1f14221 100644 --- a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs @@ -35,7 +35,7 @@ public class CosmosCheckpointStoreTests : IAsyncLifetime, IDisposable private const string TestContainerId = "Checkpoints"; // Use unique database ID per test class instance to avoid conflicts #pragma warning disable CA1802 // Use literals where appropriate - private static readonly string TestDatabaseId = $"AgentFrameworkTests-CheckpointStore-{Guid.NewGuid():N}"; + private static readonly string s_testDatabaseId = $"AgentFrameworkTests-CheckpointStore-{Guid.NewGuid():N}"; #pragma warning restore CA1802 private string _connectionString = string.Empty; @@ -69,7 +69,7 @@ public async Task InitializeAsync() this._cosmosClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); // Test connection by attempting to create database - this._database = await this._cosmosClient.CreateDatabaseIfNotExistsAsync(TestDatabaseId); + this._database = await this._cosmosClient.CreateDatabaseIfNotExistsAsync(s_testDatabaseId); await this._database.CreateContainerIfNotExistsAsync( TestContainerId, "/runId", @@ -134,10 +134,10 @@ public void Constructor_WithCosmosClient_SetsProperties() this.SkipIfEmulatorNotAvailable(); // Act - using var store = new CosmosCheckpointStore(this._cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId); // Assert - Assert.Equal(TestDatabaseId, store.DatabaseId); + Assert.Equal(s_testDatabaseId, store.DatabaseId); Assert.Equal(TestContainerId, store.ContainerId); } @@ -148,10 +148,10 @@ public void Constructor_WithConnectionString_SetsProperties() this.SkipIfEmulatorNotAvailable(); // Act - using var store = new CosmosCheckpointStore(this._connectionString, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(this._connectionString, s_testDatabaseId, TestContainerId); // Assert - Assert.Equal(TestDatabaseId, store.DatabaseId); + Assert.Equal(s_testDatabaseId, store.DatabaseId); Assert.Equal(TestContainerId, store.ContainerId); } @@ -160,7 +160,7 @@ public void Constructor_WithNullCosmosClient_ThrowsArgumentNullException() { // Act & Assert Assert.Throws(() => - new CosmosCheckpointStore((CosmosClient)null!, TestDatabaseId, TestContainerId)); + new CosmosCheckpointStore((CosmosClient)null!, s_testDatabaseId, TestContainerId)); } [SkippableFact] @@ -168,7 +168,7 @@ public void Constructor_WithNullConnectionString_ThrowsArgumentException() { // Act & Assert Assert.Throws(() => - new CosmosCheckpointStore((string)null!, TestDatabaseId, TestContainerId)); + new CosmosCheckpointStore((string)null!, s_testDatabaseId, TestContainerId)); } #endregion @@ -181,7 +181,7 @@ public async Task CreateCheckpointAsync_NewCheckpoint_CreatesSuccessfullyAsync() this.SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(this._cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId); var runId = Guid.NewGuid().ToString(); var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test checkpoint" }, s_jsonOptions); @@ -201,7 +201,7 @@ public async Task RetrieveCheckpointAsync_ExistingCheckpoint_ReturnsCorrectValue this.SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(this._cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId); var runId = Guid.NewGuid().ToString(); var originalData = new { message = "Hello, World!", timestamp = DateTimeOffset.UtcNow }; var checkpointValue = JsonSerializer.SerializeToElement(originalData, s_jsonOptions); @@ -222,7 +222,7 @@ public async Task RetrieveCheckpointAsync_NonExistentCheckpoint_ThrowsInvalidOpe this.SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(this._cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId); var runId = Guid.NewGuid().ToString(); var fakeCheckpointInfo = new CheckpointInfo(runId, "nonexistent-checkpoint"); @@ -237,7 +237,7 @@ public async Task RetrieveIndexAsync_EmptyStore_ReturnsEmptyCollectionAsync() this.SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(this._cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId); var runId = Guid.NewGuid().ToString(); // Act @@ -254,7 +254,7 @@ public async Task RetrieveIndexAsync_WithCheckpoints_ReturnsAllCheckpointsAsync( this.SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(this._cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId); var runId = Guid.NewGuid().ToString(); var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); @@ -279,7 +279,7 @@ public async Task CreateCheckpointAsync_WithParent_CreatesHierarchyAsync() this.SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(this._cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId); var runId = Guid.NewGuid().ToString(); var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); @@ -299,7 +299,7 @@ public async Task RetrieveIndexAsync_WithParentFilter_ReturnsFilteredResultsAsyn this.SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(this._cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId); var runId = Guid.NewGuid().ToString(); var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); @@ -335,7 +335,7 @@ public async Task CheckpointOperations_DifferentRuns_IsolatesDataAsync() this.SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(this._cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId); var runId1 = Guid.NewGuid().ToString(); var runId2 = Guid.NewGuid().ToString(); var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); @@ -365,7 +365,7 @@ public async Task CreateCheckpointAsync_WithNullRunId_ThrowsArgumentExceptionAsy this.SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(this._cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId); var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); // Act & Assert @@ -379,7 +379,7 @@ public async Task CreateCheckpointAsync_WithEmptyRunId_ThrowsArgumentExceptionAs this.SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(this._cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId); var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); // Act & Assert @@ -393,7 +393,7 @@ public async Task RetrieveCheckpointAsync_WithNullCheckpointInfo_ThrowsArgumentN this.SkipIfEmulatorNotAvailable(); // Arrange - using var store = new CosmosCheckpointStore(this._cosmosClient!, TestDatabaseId, TestContainerId); + using var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId); var runId = Guid.NewGuid().ToString(); // Act & Assert @@ -411,7 +411,7 @@ public async Task Dispose_AfterDisposal_ThrowsObjectDisposedExceptionAsync() this.SkipIfEmulatorNotAvailable(); // Arrange - var store = new CosmosCheckpointStore(this._cosmosClient!, TestDatabaseId, TestContainerId); + var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId); var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); // Act @@ -428,7 +428,7 @@ public void Dispose_MultipleCalls_DoesNotThrow() this.SkipIfEmulatorNotAvailable(); // Arrange - var store = new CosmosCheckpointStore(this._cosmosClient!, TestDatabaseId, TestContainerId); + var store = new CosmosCheckpointStore(this._cosmosClient!, s_testDatabaseId, TestContainerId); // Act & Assert (should not throw) store.Dispose(); @@ -451,4 +451,4 @@ protected virtual void Dispose(bool disposing) this._cosmosClient?.Dispose(); } } -} \ No newline at end of file +} From d3200f79d2ed1a5b8d51bb28f2d9114354d2c3af Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Tue, 11 Nov 2025 15:12:52 +0000 Subject: [PATCH 36/40] Address PR review comments Address Wesley's review comments: - Remove Cosmos DB package references from core projects - Delete duplicate test files from old package structure - Remove redundant parameter validation from extension methods Address Kiran's review comments: - Remove redundant 429 retry logic (SDK handles automatically) - Add explicit RequestEntityTooLarge error handling - Remove dead code in GetMessageCountAsync - Add defensive partition key validation comments --- .../CosmosChatMessageStore.cs | 40 +- .../CosmosDBChatExtensions.cs | 45 -- .../Microsoft.Agents.AI.Workflows.csproj | 3 - .../Microsoft.Agents.AI.csproj | 3 - .../CosmosChatMessageStoreTests.cs | 692 ------------------ .../CosmosCheckpointStoreTests.cs | 455 ------------ .../Microsoft.Agents.AI.UnitTests.csproj | 3 - 7 files changed, 25 insertions(+), 1216 deletions(-) delete mode 100644 dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs delete mode 100644 dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs index 2fb0377bc1..5b8a322417 100644 --- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs @@ -407,10 +407,18 @@ private async Task AddMessagesInBatchAsync(IReadOnlyCollection mess } /// - /// Executes a single batch operation with retry logic and enhanced error handling. + /// Executes a single batch operation with enhanced error handling. + /// Cosmos SDK handles throttling (429) retries automatically. /// private async Task ExecuteBatchOperationAsync(List messages, long timestamp, CancellationToken cancellationToken) { + // Validate all messages will use the same partition key (required for transactional batch) + // In practice, this is guaranteed by conversationId, but validate defensively + if (messages.Count > 0 && !messages.All(m => true)) // All messages in this store share the same partition key by design + { + throw new InvalidOperationException("All messages in a transactional batch must share the same partition key (conversationId)."); + } + var batch = this._container.CreateTransactionalBatch(this._partitionKey); foreach (var message in messages) @@ -445,12 +453,6 @@ private async Task ExecuteBatchOperationAsync(List messages, long t await this.ExecuteBatchOperationAsync(firstHalf, timestamp, cancellationToken).ConfigureAwait(false); await this.ExecuteBatchOperationAsync(secondHalf, timestamp, cancellationToken).ConfigureAwait(false); } - catch (CosmosException ex) when (ex.StatusCode == (System.Net.HttpStatusCode)429) // TooManyRequests - { - // Handle rate limiting with exponential backoff - await Task.Delay(TimeSpan.FromMilliseconds(ex.RetryAfter?.TotalMilliseconds ?? 1000), cancellationToken).ConfigureAwait(false); - await this.ExecuteBatchOperationAsync(messages, timestamp, cancellationToken).ConfigureAwait(false); - } } /// @@ -459,7 +461,19 @@ private async Task ExecuteBatchOperationAsync(List messages, long t private async Task AddSingleMessageAsync(ChatMessage message, CancellationToken cancellationToken) { var document = this.CreateMessageDocument(message, DateTimeOffset.UtcNow.ToUnixTimeSeconds()); - await this._container.CreateItemAsync(document, this._partitionKey, cancellationToken: cancellationToken).ConfigureAwait(false); + + try + { + await this._container.CreateItemAsync(document, this._partitionKey, cancellationToken: cancellationToken).ConfigureAwait(false); + } + catch (CosmosException ex) when (ex.StatusCode == System.Net.HttpStatusCode.RequestEntityTooLarge) + { + throw new InvalidOperationException( + "Message exceeds Cosmos DB's maximum item size limit of 2MB. " + + "Message ID: " + message.MessageId + ", Serialized size is too large. " + + "Consider reducing message content or splitting into smaller messages.", + ex); + } } /// @@ -531,13 +545,9 @@ public async Task GetMessageCountAsync(CancellationToken cancellationToken PartitionKey = this._partitionKey }); - if (iterator.HasMoreResults) - { - var response = await iterator.ReadNextAsync(cancellationToken).ConfigureAwait(false); - return response.FirstOrDefault(); - } - - return 0; + // COUNT queries always return a result + var response = await iterator.ReadNextAsync(cancellationToken).ConfigureAwait(false); + return response.FirstOrDefault(); } /// diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBChatExtensions.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBChatExtensions.cs index 5878aab5a6..4e3b66fd54 100644 --- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBChatExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBChatExtensions.cs @@ -35,21 +35,6 @@ public static ChatClientAgentOptions WithCosmosDBMessageStore( throw new ArgumentNullException(nameof(options)); } - if (string.IsNullOrWhiteSpace(connectionString)) - { - throw new ArgumentException("Cannot be null or whitespace", nameof(connectionString)); - } - - if (string.IsNullOrWhiteSpace(databaseId)) - { - throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); - } - - if (string.IsNullOrWhiteSpace(containerId)) - { - throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); - } - options.ChatMessageStoreFactory = context => new CosmosChatMessageStore(connectionString, databaseId, containerId); return options; } @@ -77,21 +62,6 @@ public static ChatClientAgentOptions WithCosmosDBMessageStoreUsingManagedIdentit throw new ArgumentNullException(nameof(options)); } - if (string.IsNullOrWhiteSpace(accountEndpoint)) - { - throw new ArgumentException("Cannot be null or whitespace", nameof(accountEndpoint)); - } - - if (string.IsNullOrWhiteSpace(databaseId)) - { - throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); - } - - if (string.IsNullOrWhiteSpace(containerId)) - { - throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); - } - options.ChatMessageStoreFactory = context => new CosmosChatMessageStore(accountEndpoint, new DefaultAzureCredential(), databaseId, containerId); return options; } @@ -119,21 +89,6 @@ public static ChatClientAgentOptions WithCosmosDBMessageStore( throw new ArgumentNullException(nameof(options)); } - if (cosmosClient is null) - { - throw new ArgumentNullException(nameof(cosmosClient)); - } - - if (string.IsNullOrWhiteSpace(databaseId)) - { - throw new ArgumentException("Cannot be null or whitespace", nameof(databaseId)); - } - - if (string.IsNullOrWhiteSpace(containerId)) - { - throw new ArgumentException("Cannot be null or whitespace", nameof(containerId)); - } - options.ChatMessageStoreFactory = context => new CosmosChatMessageStore(cosmosClient, databaseId, containerId); return options; } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Microsoft.Agents.AI.Workflows.csproj b/dotnet/src/Microsoft.Agents.AI.Workflows/Microsoft.Agents.AI.Workflows.csproj index 0ac25f0e1d..ff2e9dee64 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Microsoft.Agents.AI.Workflows.csproj +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Microsoft.Agents.AI.Workflows.csproj @@ -31,9 +31,6 @@ - - - diff --git a/dotnet/src/Microsoft.Agents.AI/Microsoft.Agents.AI.csproj b/dotnet/src/Microsoft.Agents.AI/Microsoft.Agents.AI.csproj index a5b52f9157..e3d7f00aa1 100644 --- a/dotnet/src/Microsoft.Agents.AI/Microsoft.Agents.AI.csproj +++ b/dotnet/src/Microsoft.Agents.AI/Microsoft.Agents.AI.csproj @@ -22,9 +22,6 @@ - - - diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs deleted file mode 100644 index 9f8daf4bdb..0000000000 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs +++ /dev/null @@ -1,692 +0,0 @@ -// Copyright (c) Microsoft. All rights reserved. - -using System; -using System.Collections.Generic; -using System.Linq; -using System.Text.Json; -using System.Text.Json.Serialization.Metadata; -using System.Threading.Tasks; -using Azure.Core; -using Azure.Identity; -using Microsoft.Azure.Cosmos; -using Microsoft.Extensions.AI; -using Microsoft.Agents.AI; -using Xunit; - -namespace Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests; - -/// -/// Contains tests for . -/// -/// Test Modes: -/// - Default Mode: Cleans up all test data after each test run (deletes database) -/// - Preserve Mode: Keeps containers and data for inspection in Cosmos DB Emulator Data Explorer -/// -/// To enable Preserve Mode, set environment variable: COSMOS_PRESERVE_CONTAINERS=true -/// Example: $env:COSMOS_PRESERVE_CONTAINERS="true"; dotnet test -/// -/// In Preserve Mode, you can view the data in Cosmos DB Emulator Data Explorer at: -/// https://localhost:8081/_explorer/index.html -/// Database: AgentFrameworkTests -/// Container: ChatMessages -/// -/// Environment Variable Reference: -/// | Variable | Values | Description | -/// |----------|--------|-------------| -/// | COSMOS_PRESERVE_CONTAINERS | true / false | Controls whether to preserve test data after completion | -/// -/// Usage Examples: -/// - Run all tests in preserve mode: $env:COSMOS_PRESERVE_CONTAINERS="true"; dotnet test tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/ -/// - Run specific test category in preserve mode: $env:COSMOS_PRESERVE_CONTAINERS="true"; dotnet test tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/ --filter "Category=CosmosDB" -/// - Reset to cleanup mode: $env:COSMOS_PRESERVE_CONTAINERS=""; dotnet test tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/ -/// -[Collection("CosmosDB")] -public sealed class CosmosChatMessageStoreTests : IAsyncLifetime, IDisposable -{ - // Cosmos DB Emulator connection settings - private const string EmulatorEndpoint = "https://localhost:8081"; - private const string EmulatorKey = "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw=="; - private const string TestContainerId = "ChatMessages"; - private const string HierarchicalTestContainerId = "HierarchicalChatMessages"; - // Use unique database ID per test class instance to avoid conflicts -#pragma warning disable CA1802 // Use literals where appropriate - private static readonly string TestDatabaseId = $"AgentFrameworkTests-ChatStore-{Guid.NewGuid():N}"; -#pragma warning restore CA1802 - - private string _connectionString = string.Empty; - private bool _emulatorAvailable; - private bool _preserveContainer; - private CosmosClient? _setupClient; // Only used for test setup/cleanup - - public async Task InitializeAsync() - { - // Check environment variable to determine if we should preserve containers - // Set COSMOS_PRESERVE_CONTAINERS=true to keep containers and data for inspection - this._preserveContainer = string.Equals(Environment.GetEnvironmentVariable("COSMOS_PRESERVE_CONTAINERS"), "true", StringComparison.OrdinalIgnoreCase); - - this._connectionString = $"AccountEndpoint={EmulatorEndpoint};AccountKey={EmulatorKey}"; - - try - { - // Only create CosmosClient for test setup - the actual tests will use connection string constructors - this._setupClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); - - // Test connection by attempting to create database - var databaseResponse = await this._setupClient.CreateDatabaseIfNotExistsAsync(TestDatabaseId); - - // Create container for simple partitioning tests - await databaseResponse.Database.CreateContainerIfNotExistsAsync( - TestContainerId, - "/conversationId", - throughput: 400); - - // Create container for hierarchical partitioning tests with hierarchical partition key - var hierarchicalContainerProperties = new ContainerProperties(HierarchicalTestContainerId, new List { "/tenantId", "/userId", "/sessionId" }); - await databaseResponse.Database.CreateContainerIfNotExistsAsync( - hierarchicalContainerProperties, - throughput: 400); - - this._emulatorAvailable = true; - } - catch (Exception) - { - // Emulator not available, tests will be skipped - this._emulatorAvailable = false; - this._setupClient?.Dispose(); - this._setupClient = null; - } - } - - public async Task DisposeAsync() - { - if (this._setupClient != null && this._emulatorAvailable) - { - try - { - if (this._preserveContainer) - { - // Preserve mode: Don't delete the database/container, keep data for inspection - // This allows viewing data in the Cosmos DB Emulator Data Explorer - // No cleanup needed - data persists for debugging - } - else - { - // Clean mode: Delete the test database and all data - var database = this._setupClient.GetDatabase(TestDatabaseId); - await database.DeleteAsync(); - } - } - catch (Exception ex) - { - // Ignore cleanup errors during test teardown - Console.WriteLine($"Warning: Cleanup failed: {ex.Message}"); - } - finally - { - this._setupClient.Dispose(); - } - } - } - - public void Dispose() - { - this._setupClient?.Dispose(); - GC.SuppressFinalize(this); - } - - private void SkipIfEmulatorNotAvailable() - { - if (!this._emulatorAvailable) - { - Assert.Fail("Cosmos DB Emulator is not available. Start the emulator to run these tests."); - } - } - - #region Constructor Tests - - [Fact] - [Trait("Category", "CosmosDB")] - public void Constructor_WithConnectionString_ShouldCreateInstance() - { - // Arrange & Act - this.SkipIfEmulatorNotAvailable(); - - // Act - using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, "test-conversation"); - - // Assert - Assert.NotNull(store); - Assert.Equal("test-conversation", store.ConversationId); - Assert.Equal(TestDatabaseId, store.DatabaseId); - Assert.Equal(TestContainerId, store.ContainerId); - } - - [Fact] - [Trait("Category", "CosmosDB")] - public void Constructor_WithConnectionStringNoConversationId_ShouldCreateInstance() - { - // Arrange - this.SkipIfEmulatorNotAvailable(); - - // Act - using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId); - - // Assert - Assert.NotNull(store); - Assert.NotNull(store.ConversationId); - Assert.Equal(TestDatabaseId, store.DatabaseId); - Assert.Equal(TestContainerId, store.ContainerId); - } - - [Fact] - [Trait("Category", "CosmosDB")] - public void Constructor_WithNullConnectionString_ShouldThrowArgumentException() - { - // Arrange & Act & Assert - Assert.Throws(() => - new CosmosChatMessageStore((string)null!, TestDatabaseId, TestContainerId, "test-conversation")); - } - - [Fact] - [Trait("Category", "CosmosDB")] - public void Constructor_WithEmptyConversationId_ShouldThrowArgumentException() - { - // Arrange & Act & Assert - this.SkipIfEmulatorNotAvailable(); - - Assert.Throws(() => - new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, "")); - } - - #endregion - - #region AddMessagesAsync Tests - - [Fact] - [Trait("Category", "CosmosDB")] - public async Task AddMessagesAsync_WithSingleMessage_ShouldAddMessageAsync() - { - // Arrange - this.SkipIfEmulatorNotAvailable(); - var conversationId = Guid.NewGuid().ToString(); - using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, conversationId); - var message = new ChatMessage(ChatRole.User, "Hello, world!"); - - // Act - await store.AddMessagesAsync([message]); - - // Wait a moment for eventual consistency - await Task.Delay(100); - - // Assert - var messages = await store.GetMessagesAsync(); - var messageList = messages.ToList(); - - // Simple assertion - if this fails, we know the deserialization is the issue - if (messageList.Count == 0) - { - // Let's check if we can find ANY items in the container for this conversation - var directQuery = new QueryDefinition("SELECT VALUE COUNT(1) FROM c WHERE c.conversationId = @conversationId") - .WithParameter("@conversationId", conversationId); - var countIterator = this._setupClient!.GetDatabase(TestDatabaseId).GetContainer(TestContainerId) - .GetItemQueryIterator(directQuery, requestOptions: new QueryRequestOptions - { - PartitionKey = new PartitionKey(conversationId) - }); - - var countResponse = await countIterator.ReadNextAsync(); - var count = countResponse.FirstOrDefault(); - - // Debug: Let's see what the raw query returns - var rawQuery = new QueryDefinition("SELECT * FROM c WHERE c.conversationId = @conversationId") - .WithParameter("@conversationId", conversationId); - var rawIterator = this._setupClient!.GetDatabase(TestDatabaseId).GetContainer(TestContainerId) - .GetItemQueryIterator(rawQuery, requestOptions: new QueryRequestOptions - { - PartitionKey = new PartitionKey(conversationId) - }); - - List rawResults = new List(); - while (rawIterator.HasMoreResults) - { - var rawResponse = await rawIterator.ReadNextAsync(); - rawResults.AddRange(rawResponse); - } - - string rawJson = rawResults.Count > 0 ? Newtonsoft.Json.JsonConvert.SerializeObject(rawResults[0], Newtonsoft.Json.Formatting.Indented) : "null"; - Assert.Fail($"GetMessagesAsync returned 0 messages, but direct count query found {count} items for conversation {conversationId}. Raw document: {rawJson}"); - } - - Assert.Single(messageList); - Assert.Equal("Hello, world!", messageList[0].Text); - Assert.Equal(ChatRole.User, messageList[0].Role); - } - - [Fact] - [Trait("Category", "CosmosDB")] - public async Task AddMessagesAsync_WithMultipleMessages_ShouldAddAllMessagesAsync() - { - // Arrange - SkipIfEmulatorNotAvailable(); - var conversationId = Guid.NewGuid().ToString(); - using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, conversationId); - var messages = new[] - { - new ChatMessage(ChatRole.User, "First message"), - new ChatMessage(ChatRole.Assistant, "Second message"), - new ChatMessage(ChatRole.User, "Third message") - }; - - // Act - await store.AddMessagesAsync(messages); - - // Assert - var retrievedMessages = await store.GetMessagesAsync(); - var messageList = retrievedMessages.ToList(); - Assert.Equal(3, messageList.Count); - Assert.Equal("First message", messageList[0].Text); - Assert.Equal("Second message", messageList[1].Text); - Assert.Equal("Third message", messageList[2].Text); - } - - #endregion - - #region GetMessagesAsync Tests - - [Fact] - [Trait("Category", "CosmosDB")] - public async Task GetMessagesAsync_WithNoMessages_ShouldReturnEmptyAsync() - { - // Arrange - SkipIfEmulatorNotAvailable(); - using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, Guid.NewGuid().ToString()); - - // Act - var messages = await store.GetMessagesAsync(); - - // Assert - Assert.Empty(messages); - } - - [Fact] - [Trait("Category", "CosmosDB")] - public async Task GetMessagesAsync_WithConversationIsolation_ShouldOnlyReturnMessagesForConversationAsync() - { - // Arrange - SkipIfEmulatorNotAvailable(); - var conversation1 = Guid.NewGuid().ToString(); - var conversation2 = Guid.NewGuid().ToString(); - - using var store1 = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, conversation1); - using var store2 = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, conversation2); - - await store1.AddMessagesAsync([new ChatMessage(ChatRole.User, "Message for conversation 1")]); - await store2.AddMessagesAsync([new ChatMessage(ChatRole.User, "Message for conversation 2")]); - - // Act - var messages1 = await store1.GetMessagesAsync(); - var messages2 = await store2.GetMessagesAsync(); - - // Assert - var messageList1 = messages1.ToList(); - var messageList2 = messages2.ToList(); - Assert.Single(messageList1); - Assert.Single(messageList2); - Assert.Equal("Message for conversation 1", messageList1[0].Text); - Assert.Equal("Message for conversation 2", messageList2[0].Text); - } - - #endregion - - #region Integration Tests - - [Fact] - [Trait("Category", "CosmosDB")] - public async Task FullWorkflow_AddAndGet_ShouldWorkCorrectlyAsync() - { - // Arrange - SkipIfEmulatorNotAvailable(); - var conversationId = $"test-conversation-{Guid.NewGuid():N}"; // Use unique conversation ID - using var originalStore = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, conversationId); - - var messages = new[] - { - new ChatMessage(ChatRole.System, "You are a helpful assistant."), - new ChatMessage(ChatRole.User, "Hello!"), - new ChatMessage(ChatRole.Assistant, "Hi there! How can I help you today?"), - new ChatMessage(ChatRole.User, "What's the weather like?"), - new ChatMessage(ChatRole.Assistant, "I'm sorry, I don't have access to current weather data.") - }; - - // Act 1: Add messages - await originalStore.AddMessagesAsync(messages); - - // Act 2: Verify messages were added - var retrievedMessages = await originalStore.GetMessagesAsync(); - var retrievedList = retrievedMessages.ToList(); - Assert.Equal(5, retrievedList.Count); - - // Act 3: Create new store instance for same conversation (test persistence) - using var newStore = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, conversationId); - var persistedMessages = await newStore.GetMessagesAsync(); - var persistedList = persistedMessages.ToList(); - - // Assert final state - Assert.Equal(5, persistedList.Count); - Assert.Equal("You are a helpful assistant.", persistedList[0].Text); - Assert.Equal("Hello!", persistedList[1].Text); - Assert.Equal("Hi there! How can I help you today?", persistedList[2].Text); - Assert.Equal("What's the weather like?", persistedList[3].Text); - Assert.Equal("I'm sorry, I don't have access to current weather data.", persistedList[4].Text); - } - - #endregion - - #region Disposal Tests - - [Fact] - [Trait("Category", "CosmosDB")] - public void Dispose_AfterUse_ShouldNotThrow() - { - // Arrange - SkipIfEmulatorNotAvailable(); - var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, Guid.NewGuid().ToString()); - - // Act & Assert - store.Dispose(); // Should not throw - } - - [Fact] - [Trait("Category", "CosmosDB")] - public void Dispose_MultipleCalls_ShouldNotThrow() - { - // Arrange - SkipIfEmulatorNotAvailable(); - var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, Guid.NewGuid().ToString()); - - // Act & Assert - store.Dispose(); // First call - store.Dispose(); // Second call - should not throw - } - - #endregion - - #region Hierarchical Partitioning Tests - - [Fact] - [Trait("Category", "CosmosDB")] - public void Constructor_WithHierarchicalConnectionString_ShouldCreateInstance() - { - // Arrange & Act - this.SkipIfEmulatorNotAvailable(); - - // Act - using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789"); - - // Assert - Assert.NotNull(store); - Assert.Equal("session-789", store.ConversationId); - Assert.Equal(TestDatabaseId, store.DatabaseId); - Assert.Equal(HierarchicalTestContainerId, store.ContainerId); - } - - [Fact] - [Trait("Category", "CosmosDB")] - public void Constructor_WithHierarchicalEndpoint_ShouldCreateInstance() - { - // Arrange & Act - this.SkipIfEmulatorNotAvailable(); - - // Act - TokenCredential credential = new DefaultAzureCredential(); - using var store = new CosmosChatMessageStore(EmulatorEndpoint, credential, TestDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789"); - - // Assert - Assert.NotNull(store); - Assert.Equal("session-789", store.ConversationId); - Assert.Equal(TestDatabaseId, store.DatabaseId); - Assert.Equal(HierarchicalTestContainerId, store.ContainerId); - } - - [Fact] - [Trait("Category", "CosmosDB")] - public void Constructor_WithHierarchicalCosmosClient_ShouldCreateInstance() - { - // Arrange & Act - this.SkipIfEmulatorNotAvailable(); - - using var cosmosClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); - using var store = new CosmosChatMessageStore(cosmosClient, TestDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789"); - - // Assert - Assert.NotNull(store); - Assert.Equal("session-789", store.ConversationId); - Assert.Equal(TestDatabaseId, store.DatabaseId); - Assert.Equal(HierarchicalTestContainerId, store.ContainerId); - } - - [Fact] - [Trait("Category", "CosmosDB")] - public void Constructor_WithHierarchicalNullTenantId_ShouldThrowArgumentException() - { - // Arrange & Act & Assert - this.SkipIfEmulatorNotAvailable(); - - Assert.Throws(() => - new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, null!, "user-456", "session-789")); - } - - [Fact] - [Trait("Category", "CosmosDB")] - public void Constructor_WithHierarchicalEmptyUserId_ShouldThrowArgumentException() - { - // Arrange & Act & Assert - this.SkipIfEmulatorNotAvailable(); - - Assert.Throws(() => - new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, "tenant-123", "", "session-789")); - } - - [Fact] - [Trait("Category", "CosmosDB")] - public void Constructor_WithHierarchicalWhitespaceSessionId_ShouldThrowArgumentException() - { - // Arrange & Act & Assert - this.SkipIfEmulatorNotAvailable(); - - Assert.Throws(() => - new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", " ")); - } - - [Fact] - [Trait("Category", "CosmosDB")] - public async Task AddMessagesAsync_WithHierarchicalPartitioning_ShouldAddMessageWithMetadataAsync() - { - // Arrange - this.SkipIfEmulatorNotAvailable(); - const string TenantId = "tenant-123"; - const string UserId = "user-456"; - const string SessionId = "session-789"; - // Test hierarchical partitioning constructor with connection string - using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, TenantId, UserId, SessionId); - var message = new ChatMessage(ChatRole.User, "Hello from hierarchical partitioning!"); - - // Act - await store.AddMessagesAsync([message]); - - // Wait a moment for eventual consistency - await Task.Delay(100); - - // Assert - var messages = await store.GetMessagesAsync(); - var messageList = messages.ToList(); - - Assert.Single(messageList); - Assert.Equal("Hello from hierarchical partitioning!", messageList[0].Text); - Assert.Equal(ChatRole.User, messageList[0].Role); - - // Verify that the document is stored with hierarchical partitioning metadata - var directQuery = new QueryDefinition("SELECT * FROM c WHERE c.conversationId = @conversationId AND c.type = @type") - .WithParameter("@conversationId", SessionId) - .WithParameter("@type", "ChatMessage"); - - var iterator = this._setupClient!.GetDatabase(TestDatabaseId).GetContainer(HierarchicalTestContainerId) - .GetItemQueryIterator(directQuery, requestOptions: new QueryRequestOptions - { - PartitionKey = new PartitionKeyBuilder().Add(TenantId).Add(UserId).Add(SessionId).Build() - }); - - var response = await iterator.ReadNextAsync(); - var document = response.FirstOrDefault(); - - Assert.NotNull(document); - // The document should have hierarchical metadata - Assert.Equal(SessionId, (string)document!.conversationId); - Assert.Equal(TenantId, (string)document!.tenantId); - Assert.Equal(UserId, (string)document!.userId); - Assert.Equal(SessionId, (string)document!.sessionId); - } - - [Fact] - [Trait("Category", "CosmosDB")] - public async Task AddMessagesAsync_WithHierarchicalMultipleMessages_ShouldAddAllMessagesAsync() - { - // Arrange - this.SkipIfEmulatorNotAvailable(); - const string TenantId = "tenant-batch"; - const string UserId = "user-batch"; - const string SessionId = "session-batch"; - // Test hierarchical partitioning constructor with connection string - using var store = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, TenantId, UserId, SessionId); - var messages = new[] - { - new ChatMessage(ChatRole.User, "First hierarchical message"), - new ChatMessage(ChatRole.Assistant, "Second hierarchical message"), - new ChatMessage(ChatRole.User, "Third hierarchical message") - }; - - // Act - await store.AddMessagesAsync(messages); - - // Wait a moment for eventual consistency - await Task.Delay(100); - - // Assert - var retrievedMessages = await store.GetMessagesAsync(); - var messageList = retrievedMessages.ToList(); - - Assert.Equal(3, messageList.Count); - Assert.Equal("First hierarchical message", messageList[0].Text); - Assert.Equal("Second hierarchical message", messageList[1].Text); - Assert.Equal("Third hierarchical message", messageList[2].Text); - } - - [Fact] - [Trait("Category", "CosmosDB")] - public async Task GetMessagesAsync_WithHierarchicalPartitionIsolation_ShouldIsolateMessagesByUserIdAsync() - { - // Arrange - this.SkipIfEmulatorNotAvailable(); - const string TenantId = "tenant-isolation"; - const string UserId1 = "user-1"; - const string UserId2 = "user-2"; - const string SessionId = "session-isolation"; - - // Different userIds create different hierarchical partitions, providing proper isolation - using var store1 = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, TenantId, UserId1, SessionId); - using var store2 = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, TenantId, UserId2, SessionId); - - // Add messages to both stores - await store1.AddMessagesAsync([new ChatMessage(ChatRole.User, "Message from user 1")]); - await store2.AddMessagesAsync([new ChatMessage(ChatRole.User, "Message from user 2")]); - - // Wait a moment for eventual consistency - await Task.Delay(100); - - // Act & Assert - var messages1 = await store1.GetMessagesAsync(); - var messageList1 = messages1.ToList(); - - var messages2 = await store2.GetMessagesAsync(); - var messageList2 = messages2.ToList(); - - // With true hierarchical partitioning, each user sees only their own messages - Assert.Single(messageList1); - Assert.Single(messageList2); - Assert.Equal("Message from user 1", messageList1[0].Text); - Assert.Equal("Message from user 2", messageList2[0].Text); - } - - [Fact] - [Trait("Category", "CosmosDB")] - public async Task SerializeDeserialize_WithHierarchicalPartitioning_ShouldPreserveStateAsync() - { - // Arrange - this.SkipIfEmulatorNotAvailable(); - const string TenantId = "tenant-serialize"; - const string UserId = "user-serialize"; - const string SessionId = "session-serialize"; - - using var originalStore = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, TenantId, UserId, SessionId); - await originalStore.AddMessagesAsync([new ChatMessage(ChatRole.User, "Test serialization message")]); - - // Act - Serialize the store state - var serializedState = originalStore.Serialize(); - - // Create a new store from the serialized state - using var cosmosClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); - var serializerOptions = new JsonSerializerOptions - { - TypeInfoResolver = new DefaultJsonTypeInfoResolver() - }; - using var deserializedStore = new CosmosChatMessageStore(serializedState, cosmosClient, serializerOptions); - - // Wait a moment for eventual consistency - await Task.Delay(100); - - // Assert - The deserialized store should have the same functionality - var messages = await deserializedStore.GetMessagesAsync(); - var messageList = messages.ToList(); - - Assert.Single(messageList); - Assert.Equal("Test serialization message", messageList[0].Text); - Assert.Equal(SessionId, deserializedStore.ConversationId); - Assert.Equal(TestDatabaseId, deserializedStore.DatabaseId); - Assert.Equal(HierarchicalTestContainerId, deserializedStore.ContainerId); - } - - [Fact] - [Trait("Category", "CosmosDB")] - public async Task HierarchicalAndSimplePartitioning_ShouldCoexistAsync() - { - // Arrange - this.SkipIfEmulatorNotAvailable(); - const string SessionId = "coexist-session"; - - // Create simple store using simple partitioning container and hierarchical store using hierarchical container - using var simpleStore = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, TestContainerId, SessionId); - using var hierarchicalStore = new CosmosChatMessageStore(this._connectionString, TestDatabaseId, HierarchicalTestContainerId, "tenant-coexist", "user-coexist", SessionId); - - // Add messages to both - await simpleStore.AddMessagesAsync([new ChatMessage(ChatRole.User, "Simple partitioning message")]); - await hierarchicalStore.AddMessagesAsync([new ChatMessage(ChatRole.User, "Hierarchical partitioning message")]); - - // Wait a moment for eventual consistency - await Task.Delay(100); - - // Act & Assert - var simpleMessages = await simpleStore.GetMessagesAsync(); - var simpleMessageList = simpleMessages.ToList(); - - var hierarchicalMessages = await hierarchicalStore.GetMessagesAsync(); - var hierarchicalMessageList = hierarchicalMessages.ToList(); - - // Each should only see its own messages since they use different containers - Assert.Single(simpleMessageList); - Assert.Single(hierarchicalMessageList); - Assert.Equal("Simple partitioning message", simpleMessageList[0].Text); - Assert.Equal("Hierarchical partitioning message", hierarchicalMessageList[0].Text); - } - - #endregion -} diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs deleted file mode 100644 index 0df3cecb18..0000000000 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests/CosmosCheckpointStoreTests.cs +++ /dev/null @@ -1,455 +0,0 @@ -// Copyright (c) Microsoft. All rights reserved. - -using System; -using System.Linq; -using System.Text.Json; -using System.Threading.Tasks; -using Microsoft.Azure.Cosmos; -using Microsoft.Agents.AI.Workflows; -using Microsoft.Agents.AI.Workflows.Checkpointing; -using Xunit; - -namespace Microsoft.Agents.AI.Abstractions.CosmosNoSql.UnitTests; - -/// -/// Contains tests for . -/// -/// Test Modes: -/// - Default Mode: Cleans up all test data after each test run (deletes database) -/// - Preserve Mode: Keeps containers and data for inspection in Cosmos DB Emulator Data Explorer -/// -/// To enable Preserve Mode, set environment variable: COSMOS_PRESERVE_CONTAINERS=true -/// Example: $env:COSMOS_PRESERVE_CONTAINERS="true"; dotnet test -/// -/// In Preserve Mode, you can view the data in Cosmos DB Emulator Data Explorer at: -/// https://localhost:8081/_explorer/index.html -/// Database: AgentFrameworkTests -/// Container: Checkpoints -/// -[Collection("CosmosDB")] -public class CosmosCheckpointStoreTests : IAsyncLifetime, IDisposable -{ - // Cosmos DB Emulator connection settings - private const string EmulatorEndpoint = "https://localhost:8081"; - private const string EmulatorKey = "C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw=="; - private const string TestContainerId = "Checkpoints"; - // Use unique database ID per test class instance to avoid conflicts -#pragma warning disable CA1802 // Use literals where appropriate - private static readonly string TestDatabaseId = $"AgentFrameworkTests-CheckpointStore-{Guid.NewGuid():N}"; -#pragma warning restore CA1802 - - private string _connectionString = string.Empty; - private CosmosClient? _cosmosClient; - private Database? _database; - private Container? _container; - private bool _emulatorAvailable; - private bool _preserveContainer; - - // JsonSerializerOptions configured for .NET 9+ compatibility - private static readonly JsonSerializerOptions s_jsonOptions = CreateJsonOptions(); - - private static JsonSerializerOptions CreateJsonOptions() - { - var options = new JsonSerializerOptions(); -#if NET9_0_OR_GREATER - options.TypeInfoResolver = new System.Text.Json.Serialization.Metadata.DefaultJsonTypeInfoResolver(); -#endif - return options; - } - - public async Task InitializeAsync() - { - // Check environment variable to determine if we should preserve containers - // Set COSMOS_PRESERVE_CONTAINERS=true to keep containers and data for inspection - _preserveContainer = string.Equals(Environment.GetEnvironmentVariable("COSMOS_PRESERVE_CONTAINERS"), "true", StringComparison.OrdinalIgnoreCase); - - _connectionString = $"AccountEndpoint={EmulatorEndpoint};AccountKey={EmulatorKey}"; - - try - { - _cosmosClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); - - // Test connection by attempting to create database - _database = await _cosmosClient.CreateDatabaseIfNotExistsAsync(TestDatabaseId); - _container = await _database.CreateContainerIfNotExistsAsync( - TestContainerId, - "/runId", - throughput: 400); - - _emulatorAvailable = true; - } - catch (Exception ex) when (!(ex is OutOfMemoryException || ex is StackOverflowException || ex is AccessViolationException)) - { - // Emulator not available, tests will be skipped - _emulatorAvailable = false; - _cosmosClient?.Dispose(); - _cosmosClient = null; - } - } - - public async Task DisposeAsync() - { - if (_cosmosClient != null && _emulatorAvailable) - { - try - { - if (_preserveContainer) - { - // Preserve mode: Don't delete the database/container, keep data for inspection - // This allows viewing data in the Cosmos DB Emulator Data Explorer - // No cleanup needed - data persists for debugging - } - else - { - // Clean mode: Delete the test database and all data - await _database!.DeleteAsync(); - } - } - catch (Exception ex) - { - // Ignore cleanup errors, but log for diagnostics - Console.WriteLine($"[DisposeAsync] Cleanup error: {ex.Message}\n{ex.StackTrace}"); - } - finally - { - _cosmosClient.Dispose(); - } - } - } - - private void SkipIfEmulatorNotAvailable() - { - if (!_emulatorAvailable) - { - // For now let's just fail to see if tests work - Assert.Fail("Cosmos DB Emulator is not available. Start the emulator to run these tests."); - } - } - - #region Constructor Tests - - [Fact] - public void Constructor_WithCosmosClient_SetsProperties() - { - // Arrange - SkipIfEmulatorNotAvailable(); - - // Act - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); - - // Assert - Assert.Equal(TestDatabaseId, store.DatabaseId); - Assert.Equal(TestContainerId, store.ContainerId); - } - - [Fact] - public void Constructor_WithConnectionString_SetsProperties() - { - // Arrange - SkipIfEmulatorNotAvailable(); - - // Act - using var store = new CosmosCheckpointStore(_connectionString, TestDatabaseId, TestContainerId); - - // Assert - Assert.Equal(TestDatabaseId, store.DatabaseId); - Assert.Equal(TestContainerId, store.ContainerId); - } - - [Fact] - public void Constructor_WithNullCosmosClient_ThrowsArgumentNullException() - { - // Act & Assert - Assert.Throws(() => - new CosmosCheckpointStore((CosmosClient)null!, TestDatabaseId, TestContainerId)); - } - - [Fact] - public void Constructor_WithNullConnectionString_ThrowsArgumentException() - { - // Act & Assert - Assert.Throws(() => - new CosmosCheckpointStore((string)null!, TestDatabaseId, TestContainerId)); - } - - #endregion - - #region Checkpoint Operations Tests - - [Fact] - public async Task CreateCheckpointAsync_NewCheckpoint_CreatesSuccessfully() - { - SkipIfEmulatorNotAvailable(); - - // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); - var runId = Guid.NewGuid().ToString(); - var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test checkpoint" }, s_jsonOptions); - - // Act - var checkpointInfo = await store.CreateCheckpointAsync(runId, checkpointValue); - - // Assert - Assert.NotNull(checkpointInfo); - Assert.Equal(runId, checkpointInfo.RunId); - Assert.NotNull(checkpointInfo.CheckpointId); - Assert.NotEmpty(checkpointInfo.CheckpointId); - } - - [Fact] - public async Task RetrieveCheckpointAsync_ExistingCheckpoint_ReturnsCorrectValue() - { - SkipIfEmulatorNotAvailable(); - - // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); - var runId = Guid.NewGuid().ToString(); - var originalData = new { message = "Hello, World!", timestamp = DateTimeOffset.UtcNow }; - var checkpointValue = JsonSerializer.SerializeToElement(originalData, s_jsonOptions); - - // Act - var checkpointInfo = await store.CreateCheckpointAsync(runId, checkpointValue); - var retrievedValue = await store.RetrieveCheckpointAsync(runId, checkpointInfo); - - // Assert - Assert.Equal(JsonValueKind.Object, retrievedValue.ValueKind); - Assert.True(retrievedValue.TryGetProperty("message", out var messageProp)); - Assert.Equal("Hello, World!", messageProp.GetString()); - } - - [Fact] - public async Task RetrieveCheckpointAsync_NonExistentCheckpoint_ThrowsInvalidOperationException() - { - SkipIfEmulatorNotAvailable(); - - // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); - var runId = Guid.NewGuid().ToString(); - var fakeCheckpointInfo = new CheckpointInfo(runId, "nonexistent-checkpoint"); - - // Act & Assert - await Assert.ThrowsAsync(() => - store.RetrieveCheckpointAsync(runId, fakeCheckpointInfo).AsTask()); - } - - [Fact] - public async Task RetrieveIndexAsync_EmptyStore_ReturnsEmptyCollection() - { - SkipIfEmulatorNotAvailable(); - - // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); - var runId = Guid.NewGuid().ToString(); - - // Act - var index = await store.RetrieveIndexAsync(runId); - - // Assert - Assert.NotNull(index); - Assert.Empty(index); - } - - [Fact] - public async Task RetrieveIndexAsync_WithCheckpoints_ReturnsAllCheckpoints() - { - SkipIfEmulatorNotAvailable(); - - // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); - var runId = Guid.NewGuid().ToString(); - var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); - - // Create multiple checkpoints - var checkpoint1 = await store.CreateCheckpointAsync(runId, checkpointValue); - var checkpoint2 = await store.CreateCheckpointAsync(runId, checkpointValue); - var checkpoint3 = await store.CreateCheckpointAsync(runId, checkpointValue); - - // Act - var index = (await store.RetrieveIndexAsync(runId)).ToList(); - - // Assert - Assert.Equal(3, index.Count); - Assert.Contains(index, c => c.CheckpointId == checkpoint1.CheckpointId); - Assert.Contains(index, c => c.CheckpointId == checkpoint2.CheckpointId); - Assert.Contains(index, c => c.CheckpointId == checkpoint3.CheckpointId); - } - - [Fact] - public async Task CreateCheckpointAsync_WithParent_CreatesHierarchy() - { - SkipIfEmulatorNotAvailable(); - - // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); - var runId = Guid.NewGuid().ToString(); - var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); - - // Act - var parentCheckpoint = await store.CreateCheckpointAsync(runId, checkpointValue); - var childCheckpoint = await store.CreateCheckpointAsync(runId, checkpointValue, parentCheckpoint); - - // Assert - Assert.NotEqual(parentCheckpoint.CheckpointId, childCheckpoint.CheckpointId); - Assert.Equal(runId, parentCheckpoint.RunId); - Assert.Equal(runId, childCheckpoint.RunId); - } - - [Fact] - public async Task RetrieveIndexAsync_WithParentFilter_ReturnsFilteredResults() - { - SkipIfEmulatorNotAvailable(); - - // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); - var runId = Guid.NewGuid().ToString(); - var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); - - // Create parent and child checkpoints - var parent = await store.CreateCheckpointAsync(runId, checkpointValue); - var child1 = await store.CreateCheckpointAsync(runId, checkpointValue, parent); - var child2 = await store.CreateCheckpointAsync(runId, checkpointValue, parent); - - // Create an orphan checkpoint - var orphan = await store.CreateCheckpointAsync(runId, checkpointValue); - - // Act - var allCheckpoints = (await store.RetrieveIndexAsync(runId)).ToList(); - var childrenOfParent = (await store.RetrieveIndexAsync(runId, parent)).ToList(); - - // Assert - Assert.Equal(4, allCheckpoints.Count); // parent + 2 children + orphan - Assert.Equal(2, childrenOfParent.Count); // only children - - Assert.Contains(childrenOfParent, c => c.CheckpointId == child1.CheckpointId); - Assert.Contains(childrenOfParent, c => c.CheckpointId == child2.CheckpointId); - Assert.DoesNotContain(childrenOfParent, c => c.CheckpointId == parent.CheckpointId); - Assert.DoesNotContain(childrenOfParent, c => c.CheckpointId == orphan.CheckpointId); - } - - #endregion - - #region Run Isolation Tests - - [Fact] - public async Task CheckpointOperations_DifferentRuns_IsolatesData() - { - SkipIfEmulatorNotAvailable(); - - // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); - var runId1 = Guid.NewGuid().ToString(); - var runId2 = Guid.NewGuid().ToString(); - var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); - - // Act - var checkpoint1 = await store.CreateCheckpointAsync(runId1, checkpointValue); - var checkpoint2 = await store.CreateCheckpointAsync(runId2, checkpointValue); - - var index1 = (await store.RetrieveIndexAsync(runId1)).ToList(); - var index2 = (await store.RetrieveIndexAsync(runId2)).ToList(); - - // Assert - Assert.Single(index1); - Assert.Single(index2); - Assert.Equal(checkpoint1.CheckpointId, index1[0].CheckpointId); - Assert.Equal(checkpoint2.CheckpointId, index2[0].CheckpointId); - Assert.NotEqual(checkpoint1.CheckpointId, checkpoint2.CheckpointId); - } - - #endregion - - #region Error Handling Tests - - [Fact] - public async Task CreateCheckpointAsync_WithNullRunId_ThrowsArgumentException() - { - SkipIfEmulatorNotAvailable(); - - // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); - var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); - - // Act & Assert - await Assert.ThrowsAsync(() => - store.CreateCheckpointAsync(null!, checkpointValue).AsTask()); - } - - [Fact] - public async Task CreateCheckpointAsync_WithEmptyRunId_ThrowsArgumentException() - { - SkipIfEmulatorNotAvailable(); - - // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); - var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); - - // Act & Assert - await Assert.ThrowsAsync(() => - store.CreateCheckpointAsync("", checkpointValue).AsTask()); - } - - [Fact] - public async Task RetrieveCheckpointAsync_WithNullCheckpointInfo_ThrowsArgumentNullException() - { - SkipIfEmulatorNotAvailable(); - - // Arrange - using var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); - var runId = Guid.NewGuid().ToString(); - - // Act & Assert - await Assert.ThrowsAsync(() => - store.RetrieveCheckpointAsync(runId, null!).AsTask()); - } - - #endregion - - #region Disposal Tests - - [Fact] - public async Task Dispose_AfterDisposal_ThrowsObjectDisposedException() - { - SkipIfEmulatorNotAvailable(); - - // Arrange - var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); - var checkpointValue = JsonSerializer.SerializeToElement(new { data = "test" }, s_jsonOptions); - - // Act - store.Dispose(); - - // Assert - await Assert.ThrowsAsync(() => - store.CreateCheckpointAsync("test-run", checkpointValue).AsTask()); - } - - [Fact] - public void Dispose_MultipleCalls_DoesNotThrow() - { - SkipIfEmulatorNotAvailable(); - - // Arrange - var store = new CosmosCheckpointStore(_cosmosClient!, TestDatabaseId, TestContainerId); - - // Act & Assert (should not throw) - store.Dispose(); - store.Dispose(); - store.Dispose(); - } - - #endregion - - public void Dispose() - { - Dispose(true); - GC.SuppressFinalize(this); - } - - protected virtual void Dispose(bool disposing) - { - if (disposing) - { - this._cosmosClient?.Dispose(); - } - } -} diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Microsoft.Agents.AI.UnitTests.csproj b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Microsoft.Agents.AI.UnitTests.csproj index a567b8be58..f871781d03 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Microsoft.Agents.AI.UnitTests.csproj +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Microsoft.Agents.AI.UnitTests.csproj @@ -18,9 +18,6 @@ - - - From 59b9cf8d037ba73bceaef3c58bb7c8c18c6590ef Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Tue, 11 Nov 2025 15:31:42 +0000 Subject: [PATCH 37/40] Fix IDE0001 formatting error in AgentProviderExtensions.cs. Use type alias to resolve namespace conflict between Azure.AI.Agents.Persistent.RunStatus and Microsoft.Agents.AI.Workflows.RunStatus. This eliminates the need for global:: qualifier which triggered the formatter warning. --- .../Extensions/AgentProviderExtensions.cs | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Extensions/AgentProviderExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Extensions/AgentProviderExtensions.cs index e594edd0f0..8d19978e0b 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Extensions/AgentProviderExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows.Declarative/Extensions/AgentProviderExtensions.cs @@ -6,6 +6,7 @@ using System.Threading.Tasks; #if NET9_0_OR_GREATER using Azure.AI.Agents.Persistent; +using AzureRunStatus = Azure.AI.Agents.Persistent.RunStatus; #endif using Microsoft.Extensions.AI; @@ -14,12 +15,12 @@ namespace Microsoft.Agents.AI.Workflows.Declarative.Extensions; internal static class AgentProviderExtensions { #if NET9_0_OR_GREATER - private static readonly HashSet s_failureStatus = + private static readonly HashSet s_failureStatus = [ - global::Azure.AI.Agents.Persistent.RunStatus.Failed, - global::Azure.AI.Agents.Persistent.RunStatus.Cancelled, - global::Azure.AI.Agents.Persistent.RunStatus.Cancelling, - global::Azure.AI.Agents.Persistent.RunStatus.Expired, + AzureRunStatus.Failed, + AzureRunStatus.Cancelled, + AzureRunStatus.Cancelling, + AzureRunStatus.Expired, ]; #endif From 108d7ab588646502dffa941fe0e7b976712ca53f Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Tue, 18 Nov 2025 07:29:13 -0800 Subject: [PATCH 38/40] Update package versions for Aspire 13.0.0 compatibility --- dotnet/Directory.Packages.props | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dotnet/Directory.Packages.props b/dotnet/Directory.Packages.props index 58e6c01dd8..5335e5582d 100644 --- a/dotnet/Directory.Packages.props +++ b/dotnet/Directory.Packages.props @@ -24,9 +24,9 @@ - + - + From d472ad4f9479f33bb3b45b84013da10d822e159d Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Mon, 24 Nov 2025 21:08:11 +0000 Subject: [PATCH 39/40] Fix TargetFrameworks in Cosmos DB projects - Replace with which is defined in Directory.Build.props - Fix package reference from System.Linq.Async to System.Linq.AsyncEnumerable to match Directory.Packages.props --- .../Microsoft.Agents.AI.CosmosNoSql.csproj | 3 +-- .../Microsoft.Agents.AI.CosmosNoSql.UnitTests.csproj | 4 ++-- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/Microsoft.Agents.AI.CosmosNoSql.csproj b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/Microsoft.Agents.AI.CosmosNoSql.csproj index b6f38b81be..7e13ec5998 100644 --- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/Microsoft.Agents.AI.CosmosNoSql.csproj +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/Microsoft.Agents.AI.CosmosNoSql.csproj @@ -1,8 +1,7 @@ - $(ProjectsTargetFrameworks) - $(ProjectsDebugTargetFrameworks) + $(TargetFrameworksCore) Microsoft.Agents.AI $(NoWarn);MEAI001 preview diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/Microsoft.Agents.AI.CosmosNoSql.UnitTests.csproj b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/Microsoft.Agents.AI.CosmosNoSql.UnitTests.csproj index d2050fb52f..d60418ee2c 100644 --- a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/Microsoft.Agents.AI.CosmosNoSql.UnitTests.csproj +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/Microsoft.Agents.AI.CosmosNoSql.UnitTests.csproj @@ -1,7 +1,7 @@ - $(ProjectsTargetFrameworks) + net10.0;net9.0 $(NoWarn);MEAI001 @@ -15,7 +15,7 @@ - + From 5f061776f6def086e018434b9b68c4290f7ab4d8 Mon Sep 17 00:00:00 2001 From: Theo van Kraay Date: Tue, 25 Nov 2025 19:41:22 +0000 Subject: [PATCH 40/40] Remove redundant counter, add partition key validation, use factory pattern for deserialization --- .../CosmosChatMessageStore.cs | 94 +++++++++++-------- .../CosmosChatMessageStoreTests.cs | 2 +- 2 files changed, 54 insertions(+), 42 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs index 5b8a322417..fff7f56fa5 100644 --- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatMessageStore.cs @@ -253,48 +253,37 @@ public CosmosChatMessageStore(CosmosClient cosmosClient, string databaseId, stri } /// - /// Initializes a new instance of the class from previously serialized state. + /// Creates a new instance of the class from previously serialized state. /// /// The instance to use for Cosmos DB operations. /// A representing the serialized state of the message store. /// The identifier of the Cosmos DB database. /// The identifier of the Cosmos DB container. /// Optional settings for customizing the JSON deserialization process. + /// A new instance of initialized from the serialized state. /// Thrown when is null. /// Thrown when the serialized state cannot be deserialized. - public CosmosChatMessageStore(CosmosClient cosmosClient, JsonElement serializedStoreState, string databaseId, string containerId, JsonSerializerOptions? jsonSerializerOptions = null) + public static CosmosChatMessageStore CreateFromSerializedState(CosmosClient cosmosClient, JsonElement serializedStoreState, string databaseId, string containerId, JsonSerializerOptions? jsonSerializerOptions = null) { - this._cosmosClient = cosmosClient ?? throw new ArgumentNullException(nameof(cosmosClient)); - this.DatabaseId = Throw.IfNullOrWhitespace(databaseId); - this.ContainerId = Throw.IfNullOrWhitespace(containerId); - this._container = this._cosmosClient.GetContainer(databaseId, containerId); - this._ownsClient = false; + Throw.IfNull(cosmosClient); + Throw.IfNullOrWhitespace(databaseId); + Throw.IfNullOrWhitespace(containerId); - if (serializedStoreState.ValueKind is JsonValueKind.Object) + if (serializedStoreState.ValueKind is not JsonValueKind.Object) { - var state = JsonSerializer.Deserialize(serializedStoreState, jsonSerializerOptions); - if (state?.ConversationIdentifier is { } conversationId) - { - this.ConversationId = conversationId; - - // Initialize hierarchical partitioning if available in state - this._tenantId = state.TenantId; - this._userId = state.UserId; - this._useHierarchicalPartitioning = state.UseHierarchicalPartitioning; - - this._partitionKey = (this._useHierarchicalPartitioning && this._tenantId != null && this._userId != null) - ? new PartitionKeyBuilder() - .Add(this._tenantId) - .Add(this._userId) - .Add(conversationId) - .Build() - : new PartitionKey(conversationId); + throw new ArgumentException("Invalid serialized state", nameof(serializedStoreState)); + } - return; - } + var state = JsonSerializer.Deserialize(serializedStoreState, jsonSerializerOptions); + if (state?.ConversationIdentifier is not { } conversationId) + { + throw new ArgumentException("Invalid serialized state", nameof(serializedStoreState)); } - throw new ArgumentException("Invalid serialized state", nameof(serializedStoreState)); + // Use the internal constructor with all parameters to ensure partition key logic is centralized + return state.UseHierarchicalPartitioning && state.TenantId != null && state.UserId != null + ? new CosmosChatMessageStore(cosmosClient, databaseId, containerId, conversationId, ownsClient: false, state.TenantId, state.UserId) + : new CosmosChatMessageStore(cosmosClient, databaseId, containerId, conversationId, ownsClient: false); } /// @@ -320,7 +309,6 @@ public override async Task> GetMessagesAsync(Cancellati }); var messages = new List(); - int messageCount = 0; while (iterator.HasMoreResults) { @@ -328,7 +316,7 @@ public override async Task> GetMessagesAsync(Cancellati foreach (var document in response) { - if (this.MaxMessagesToRetrieve.HasValue && messageCount >= this.MaxMessagesToRetrieve.Value) + if (this.MaxMessagesToRetrieve.HasValue && messages.Count >= this.MaxMessagesToRetrieve.Value) { break; } @@ -339,12 +327,11 @@ public override async Task> GetMessagesAsync(Cancellati if (message != null) { messages.Add(message); - messageCount++; } } } - if (this.MaxMessagesToRetrieve.HasValue && messageCount >= this.MaxMessagesToRetrieve.Value) + if (this.MaxMessagesToRetrieve.HasValue && messages.Count >= this.MaxMessagesToRetrieve.Value) { break; } @@ -412,18 +399,44 @@ private async Task AddMessagesInBatchAsync(IReadOnlyCollection mess /// private async Task ExecuteBatchOperationAsync(List messages, long timestamp, CancellationToken cancellationToken) { - // Validate all messages will use the same partition key (required for transactional batch) - // In practice, this is guaranteed by conversationId, but validate defensively - if (messages.Count > 0 && !messages.All(m => true)) // All messages in this store share the same partition key by design + // Create all documents upfront for validation and batch operation + var documents = new List(messages.Count); + foreach (var message in messages) { - throw new InvalidOperationException("All messages in a transactional batch must share the same partition key (conversationId)."); + documents.Add(this.CreateMessageDocument(message, timestamp)); } + // Defensive check: Verify all messages share the same partition key values + // In hierarchical partitioning, this means same tenantId, userId, and sessionId + // In simple partitioning, this means same conversationId + if (documents.Count > 0) + { + if (this._useHierarchicalPartitioning) + { + // Verify all documents have matching hierarchical partition key components + var firstDoc = documents[0]; + if (!documents.All(d => d.TenantId == firstDoc.TenantId && d.UserId == firstDoc.UserId && d.SessionId == firstDoc.SessionId)) + { + throw new InvalidOperationException("All messages in a batch must share the same partition key values (tenantId, userId, sessionId)."); + } + } + else + { + // Verify all documents have matching conversationId + var firstConversationId = documents[0].ConversationId; + if (!documents.All(d => d.ConversationId == firstConversationId)) + { + throw new InvalidOperationException("All messages in a batch must share the same partition key value (conversationId)."); + } + } + } + + // All messages in this store share the same partition key by design + // Transactional batches require all items to share the same partition key var batch = this._container.CreateTransactionalBatch(this._partitionKey); - foreach (var message in messages) + foreach (var document in documents) { - var document = this.CreateMessageDocument(message, timestamp); batch.CreateItem(document); } @@ -572,17 +585,16 @@ public async Task ClearMessagesAsync(CancellationToken cancellationToken = var iterator = this._container.GetItemQueryIterator(query, requestOptions: new QueryRequestOptions { - PartitionKey = new PartitionKey(this.ConversationId), + PartitionKey = this._partitionKey, MaxItemCount = this.MaxItemCount }); var deletedCount = 0; - var partitionKey = new PartitionKey(this.ConversationId); while (iterator.HasMoreResults) { var response = await iterator.ReadNextAsync(cancellationToken).ConfigureAwait(false); - var batch = this._container.CreateTransactionalBatch(partitionKey); + var batch = this._container.CreateTransactionalBatch(this._partitionKey); var batchItemCount = 0; foreach (var itemId in response) diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs index 6b62248a16..6f2a256206 100644 --- a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatMessageStoreTests.cs @@ -640,7 +640,7 @@ public async Task SerializeDeserialize_WithHierarchicalPartitioning_ShouldPreser { TypeInfoResolver = new DefaultJsonTypeInfoResolver() }; - using var deserializedStore = new CosmosChatMessageStore(cosmosClient, serializedState, s_testDatabaseId, HierarchicalTestContainerId, serializerOptions); + using var deserializedStore = CosmosChatMessageStore.CreateFromSerializedState(cosmosClient, serializedState, s_testDatabaseId, HierarchicalTestContainerId, serializerOptions); // Wait a moment for eventual consistency await Task.Delay(100);