diff --git a/dotnet/Directory.Packages.props b/dotnet/Directory.Packages.props index cd140a52e7..d1f31f3b5f 100644 --- a/dotnet/Directory.Packages.props +++ b/dotnet/Directory.Packages.props @@ -10,6 +10,9 @@ + + + @@ -21,6 +24,9 @@ + + + diff --git a/dotnet/agent-framework-dotnet.slnx b/dotnet/agent-framework-dotnet.slnx index ea7898a16e..7e976a0493 100644 --- a/dotnet/agent-framework-dotnet.slnx +++ b/dotnet/agent-framework-dotnet.slnx @@ -145,6 +145,9 @@ + + + @@ -163,5 +166,14 @@ + + + + + + + + + diff --git a/dotnet/samples/HelloHttpApi/HelloHttpApi.ApiService/HelloHttpApi.ApiService.csproj b/dotnet/samples/HelloHttpApi/HelloHttpApi.ApiService/HelloHttpApi.ApiService.csproj index e3f02a5f4c..6efdd1f81e 100644 --- a/dotnet/samples/HelloHttpApi/HelloHttpApi.ApiService/HelloHttpApi.ApiService.csproj +++ b/dotnet/samples/HelloHttpApi/HelloHttpApi.ApiService/HelloHttpApi.ApiService.csproj @@ -12,12 +12,14 @@ + + diff --git a/dotnet/samples/HelloHttpApi/HelloHttpApi.ApiService/HostApplicationBuilderAgentExtensions.cs b/dotnet/samples/HelloHttpApi/HelloHttpApi.ApiService/HostApplicationBuilderAgentExtensions.cs index 3773a2ca56..284e3f18fa 100644 --- a/dotnet/samples/HelloHttpApi/HelloHttpApi.ApiService/HostApplicationBuilderAgentExtensions.cs +++ b/dotnet/samples/HelloHttpApi/HelloHttpApi.ApiService/HostApplicationBuilderAgentExtensions.cs @@ -4,6 +4,7 @@ using Microsoft.Agents.Orchestration; using Microsoft.Extensions.AI; using Microsoft.Extensions.AI.Agents; using Microsoft.Extensions.AI.Agents.Runtime; +using Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB; namespace HelloHttpApi.ApiService; @@ -26,8 +27,12 @@ public static class HostApplicationBuilderAgentExtensions .Add(triage, customerService, "Hand off to the customer service agent for handling rude customer inquiries.") .Build("PirateWorkflow"); }); + var actorBuilder = builder.AddActorRuntime(); + // Add CosmosDB state storage to override default storage + builder.Services.AddCosmosActorStateStorage("actor-state-db", "ActorState"); + actorBuilder.AddActorType( new ActorType(agentKey), (sp, ctx) => new ChatClientAgentActor( diff --git a/dotnet/samples/HelloHttpApi/HelloHttpApi.ApiService/Program.cs b/dotnet/samples/HelloHttpApi/HelloHttpApi.ApiService/Program.cs index 4b963646b6..81a8534708 100644 --- a/dotnet/samples/HelloHttpApi/HelloHttpApi.ApiService/Program.cs +++ b/dotnet/samples/HelloHttpApi/HelloHttpApi.ApiService/Program.cs @@ -8,6 +8,9 @@ var builder = WebApplication.CreateBuilder(args); // Add service defaults & Aspire client integrations. builder.AddServiceDefaults(); +// Add CosmosDB client integration +builder.AddAzureCosmosClient("hello-http-api-cosmosdb"); + // Add services to the container. builder.Services.AddProblemDetails(); diff --git a/dotnet/samples/HelloHttpApi/HelloHttpApi.AppHost/HelloHttpApi.AppHost.csproj b/dotnet/samples/HelloHttpApi/HelloHttpApi.AppHost/HelloHttpApi.AppHost.csproj index c0b12bead9..95c5809b39 100644 --- a/dotnet/samples/HelloHttpApi/HelloHttpApi.AppHost/HelloHttpApi.AppHost.csproj +++ b/dotnet/samples/HelloHttpApi/HelloHttpApi.AppHost/HelloHttpApi.AppHost.csproj @@ -14,6 +14,7 @@ + diff --git a/dotnet/samples/HelloHttpApi/HelloHttpApi.AppHost/Program.cs b/dotnet/samples/HelloHttpApi/HelloHttpApi.AppHost/Program.cs index 97701fcf10..ddf06aa01e 100644 --- a/dotnet/samples/HelloHttpApi/HelloHttpApi.AppHost/Program.cs +++ b/dotnet/samples/HelloHttpApi/HelloHttpApi.AppHost/Program.cs @@ -8,8 +8,15 @@ var azOpenAiResource = builder.AddParameterFromConfiguration("AzureOpenAIName", var azOpenAiResourceGroup = builder.AddParameterFromConfiguration("AzureOpenAIResourceGroup", "AzureOpenAI:ResourceGroup"); var chatModel = builder.AddAIModel("chat-model").AsAzureOpenAI("gpt-4o", o => o.AsExisting(azOpenAiResource, azOpenAiResourceGroup)); +var cosmosDbResource = builder.AddParameterFromConfiguration("CosmosDbName", "CosmosDb:Name"); +var cosmosDbResourceGroup = builder.AddParameterFromConfiguration("CosmosDbResourceGroup", "CosmosDb:ResourceGroup"); +var cosmos = builder.AddAzureCosmosDB("hello-http-api-cosmosdb").RunAsExisting(cosmosDbResource, cosmosDbResourceGroup); + +var stateDb = cosmos.AddCosmosDatabase("actor-state-db"); + var apiService = builder.AddProject("apiservice") - .WithReference(chatModel); + .WithReference(chatModel) + .WithReference(cosmos).WaitFor(cosmos); builder.AddProject("webfrontend") .WithExternalHttpEndpoints() diff --git a/dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB/ActorDocuments.cs b/dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB/ActorDocuments.cs new file mode 100644 index 0000000000..4b8b3cb296 --- /dev/null +++ b/dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB/ActorDocuments.cs @@ -0,0 +1,82 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Text.Json; + +namespace Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB; + +/// +/// Root document for each actor that provides actor-level ETag semantics. +/// Every write operation updates this document to ensure a single ETag represents +/// the entire actor's state for optimistic concurrency control. +/// This document contains no actor state data. It only serves to track last modified +/// time and provide a single ETag for the actor's state. +/// +/// Example structure: +/// { +/// "id": "rootdoc", // Root document ID (constant per actor partition) +/// "actorId": "actor-123", // Partition key (actor ID) +/// "lastModified": "2024-...", // Timestamp +/// } +/// +public sealed class ActorRootDocument +{ + /// + /// The document ID. + /// + public string Id { get; set; } = default!; + + /// + /// The actor ID. + /// + public string ActorId { get; set; } = default!; + + /// + /// The last modified timestamp. + /// + public DateTimeOffset LastModified { get; set; } +} + +/// +/// Actor state document that represents a single key-value pair in the actor's state. +/// Document Structure (one per actor key): +/// { +/// "id": "state_sanitizedkey", // Unique document ID for the state entry +/// "actorId": "actor-123", // Partition key (actor ID) +/// "key": "foo", // Logical key for the state entry +/// "value": { "bar": 42, "baz": "hello" } // Arbitrary JsonElement payload +/// } +/// +public sealed class ActorStateDocument +{ + /// + /// The document ID. + /// + public string Id { get; set; } = default!; + + /// + /// The actor ID. + /// + public string ActorId { get; set; } = default!; + + /// + /// The logical key for the state entry. + /// + public string Key { get; set; } = default!; + + /// + /// The value payload. + /// + public JsonElement Value { get; set; } = default!; +} + +/// +/// Projection class for Cosmos DB queries to retrieve keys. +/// +public sealed class KeyProjection +{ + /// + /// The key value. + /// + public string Key { get; set; } = default!; +} diff --git a/dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB/CosmosActorStateJsonContext.cs b/dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB/CosmosActorStateJsonContext.cs new file mode 100644 index 0000000000..e88dbe310e --- /dev/null +++ b/dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB/CosmosActorStateJsonContext.cs @@ -0,0 +1,19 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System.Text.Json; +using System.Text.Json.Serialization; + +namespace Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB; + +/// +/// Source-generated JSON type information for Cosmos DB actor state documents. +/// +[JsonSourceGenerationOptions( + JsonSerializerDefaults.Web, + UseStringEnumConverter = true, + DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull, + WriteIndented = false)] +[JsonSerializable(typeof(ActorStateDocument))] +[JsonSerializable(typeof(ActorRootDocument))] +[JsonSerializable(typeof(KeyProjection))] +internal sealed partial class CosmosActorStateJsonContext : JsonSerializerContext; diff --git a/dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB/CosmosActorStateStorage.cs b/dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB/CosmosActorStateStorage.cs new file mode 100644 index 0000000000..547485e91b --- /dev/null +++ b/dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB/CosmosActorStateStorage.cs @@ -0,0 +1,237 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Azure.Cosmos; + +namespace Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB; + +/// +/// Cosmos DB implementation of actor state storage. +/// +public class CosmosActorStateStorage : IActorStateStorage +{ + private readonly LazyCosmosContainer _lazyContainer; + private const string InitialEtag = "0"; // Initial ETag value when no state exists + + /// + /// Constructs a new instance of with the specified Cosmos DB container. + /// + /// The Cosmos DB container to use for storage. + /// Thrown when is null. + public CosmosActorStateStorage(Container container) => this._lazyContainer = new LazyCosmosContainer(container); + + /// + /// This constructor is used by dependency injection to create an instance of + /// with a lazy-loaded Cosmos container whose initialization is deferred until first access. + /// + /// The lazy-loaded Cosmos container. + /// Thrown when is null. + internal CosmosActorStateStorage(LazyCosmosContainer lazyContainer) => + this._lazyContainer = lazyContainer ?? throw new ArgumentNullException(nameof(lazyContainer)); + + /// + /// Writes state changes to the actor's persistent storage. + /// + public async ValueTask WriteStateAsync( + ActorId actorId, + IReadOnlyCollection operations, + string etag, + CancellationToken cancellationToken = default) + { + if (operations.Count == 0) + { + throw new InvalidOperationException("No operations provided for write. At least one operation is required."); + } + + var container = await this._lazyContainer.GetContainerAsync().ConfigureAwait(false); + var batch = container.CreateTransactionalBatch(GetPartitionKey(actorId)); + var actorIdStr = actorId.ToString(); + + // Add data operations to batch + foreach (var op in operations) + { + switch (op) + { + case SetValueOperation set: + var docId = GetDocumentId(set.Key); + + var item = new ActorStateDocument + { + Id = docId, + ActorId = actorIdStr, + Key = set.Key, + Value = set.Value + }; + + batch.UpsertItem(item); + break; + + case RemoveKeyOperation remove: + var docToRemove = GetDocumentId(remove.Key); + batch.DeleteItem(docToRemove); + break; + + default: + throw new ArgumentException($"Unsupported write operation: {op.GetType().Name}"); + } + } + + // Add root document update to batch + var newRoot = new ActorRootDocument + { + Id = RootDocumentId, + ActorId = actorId.ToString(), + LastModified = DateTimeOffset.UtcNow, + }; + + if (string.IsNullOrEmpty(etag) || etag == InitialEtag) + { + // No eTag provided or initial eTag - create new root document (will fail if it already exists) + batch.CreateItem(newRoot); + } + else + { + // eTag provided - replace existing root document with eTag check + batch.ReplaceItem(RootDocumentId, newRoot, new TransactionalBatchItemRequestOptions { IfMatchEtag = etag }); + } + + try + { + var result = await batch.ExecuteAsync(cancellationToken).ConfigureAwait(false); + if (!result.IsSuccessStatusCode) + { + return new WriteResponse(eTag: string.Empty, success: false); + } + + // Get the ETag from the root document operation (last operation in batch) + var rootResult = result[result.Count - 1]; + return new WriteResponse(eTag: rootResult.ETag, success: true); + } + catch (CosmosException) + { + // If any operation in the batch fails, we return failure + return new WriteResponse(eTag: string.Empty, success: false); + } + } + + /// + /// Reads state data from the actor's persistent storage. + /// + public async ValueTask ReadStateAsync( + ActorId actorId, + IReadOnlyCollection operations, + CancellationToken cancellationToken = default) + { + if (operations.Count == 0) + { + throw new InvalidOperationException("No operations provided for read. At least one operation is required."); + } + + var container = await this._lazyContainer.GetContainerAsync().ConfigureAwait(false); + var results = new List(); + + // Read root document first to get actor-level ETag + string actorETag = await this.GetActorETagAsync(container, actorId, cancellationToken).ConfigureAwait(false); + + foreach (var op in operations) + { + switch (op) + { + case GetValueOperation get: + var id = GetDocumentId(get.Key); + try + { + var response = await container.ReadItemAsync( + id, + GetPartitionKey(actorId), + cancellationToken: cancellationToken) + .ConfigureAwait(false); + + results.Add(new GetValueResult(response.Resource.Value)); + } + catch (CosmosException ex) when (ex.StatusCode == System.Net.HttpStatusCode.NotFound) + { + results.Add(new GetValueResult(null)); + } + break; + + case ListKeysOperation list: + QueryDefinition query; + if (!string.IsNullOrEmpty(list.KeyPrefix)) + { + query = new QueryDefinition("SELECT c.key FROM c WHERE c.actorId = @actorId AND c.key != null AND STARTSWITH(c.key, @keyPrefix)") + .WithParameter("@actorId", actorId.ToString()) + .WithParameter("@keyPrefix", list.KeyPrefix); + } + else + { + query = new QueryDefinition("SELECT c.key FROM c WHERE c.actorId = @actorId AND c.key != null") + .WithParameter("@actorId", actorId.ToString()); + } + + var requestOptions = new QueryRequestOptions + { + PartitionKey = GetPartitionKey(actorId), + MaxItemCount = -1 // Use dynamic page size + }; + + var iterator = container.GetItemQueryIterator( + query, + list.ContinuationToken, + requestOptions); + + var keys = new List(); + string? continuationToken = null; + + while (iterator.HasMoreResults) + { + var page = await iterator.ReadNextAsync(cancellationToken).ConfigureAwait(false); + foreach (var projection in page) + { + keys.Add(projection.Key); + } + + continuationToken = page.ContinuationToken; + } + + results.Add(new ListKeysResult(keys, continuationToken)); + break; + + default: + throw new NotSupportedException($"Unsupported read operation: {op.GetType().Name}"); + } + } + + return new ReadResponse(actorETag, results); + } + + private static string GetDocumentId(string key) => $"state_{CosmosIdSanitizer.Sanitize(key)}"; + private const string RootDocumentId = "rootdoc"; + + private static PartitionKey GetPartitionKey(ActorId actorId) + => new(actorId.ToString()); + + /// + /// Gets the current ETag for the actor's root document. + /// Returns a generated ETag if no root document exists. + /// + private async ValueTask GetActorETagAsync(Container container, ActorId actorId, CancellationToken cancellationToken) + { + try + { + var rootResponse = await container.ReadItemAsync( + RootDocumentId, + GetPartitionKey(actorId), + cancellationToken: cancellationToken).ConfigureAwait(false); + return rootResponse.ETag; + } + catch (CosmosException ex) when (ex.StatusCode == System.Net.HttpStatusCode.NotFound) + { + // No root document means no actor state exists + return Guid.NewGuid().ToString("N"); + } + } +} diff --git a/dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB/CosmosIdSanitizer.cs b/dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB/CosmosIdSanitizer.cs new file mode 100644 index 0000000000..38b6362157 --- /dev/null +++ b/dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB/CosmosIdSanitizer.cs @@ -0,0 +1,133 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; + +namespace Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB; + +// CosmosIdSanitizer is used to sanitize Cosmos DB IDs by replacing characters that are +// not allowed in Cosmos DB IDs with a safe escape sequence. This implementation was +// inspired heavily by the one in Orleans, with some modifications made to enable +// targeting NET472. +internal static class CosmosIdSanitizer +{ + private const char EscapeChar = '~'; + public const char SeparatorChar = '_'; + + private static ReadOnlySpan SanitizedCharacters => ['/', '\\', '?', '#', SeparatorChar, EscapeChar]; + private static ReadOnlySpan ReplacementCharacters => ['0', '1', '2', '3', '4', '5']; + + public static string Sanitize(string input) + { + int extraChars = CountSanitizedCharacters(input.AsSpan()); + + if (extraChars == 0) + { + return input; + } + +#if NET8_0_OR_GREATER + return string.Create(input.Length + extraChars, input, (output, state) => + { + Encode(state.AsSpan(), output); + }); +#else + var result = new char[input.Length + extraChars]; + Encode(input.AsSpan(), result); + return new string(result, 0, input.Length + extraChars); +#endif + } + + public static string Unsanitize(string input) + { + int escapeCount = CountEscapeCharacters(input.AsSpan()); + + if (escapeCount == 0) + { + return input; + } + +#if NET8_0_OR_GREATER + return string.Create(input.Length - escapeCount, input, (output, state) => + { + Decode(state.AsSpan(), output); + }); +#else + var result = new char[input.Length - escapeCount]; + Decode(input.AsSpan(), result); + return new string(result, 0, input.Length - escapeCount); +#endif + } + + private static int CountSanitizedCharacters(ReadOnlySpan input) + { + int count = 0; + foreach (var c in input) + { + if (SanitizedCharacters.IndexOf(c) >= 0) + { + count++; + } + } + return count; + } + + private static int CountEscapeCharacters(ReadOnlySpan input) + { + int count = 0; + foreach (var c in input) + { + if (c == EscapeChar) + { + count++; + } + } + return count; + } + + private static void Encode(ReadOnlySpan input, Span output) + { + int j = 0; + foreach (var c in input) + { + int idx = SanitizedCharacters.IndexOf(c); + if (idx < 0) + { + output[j++] = c; + } + else + { + output[j++] = EscapeChar; + output[j++] = ReplacementCharacters[idx]; + } + } + } + + private static void Decode(ReadOnlySpan input, Span output) + { + int j = 0; + bool isEscaped = false; + + foreach (var c in input) + { + if (isEscaped) + { + int idx = ReplacementCharacters.IndexOf(c); + if (idx < 0) + { + throw new ArgumentException("Input is not in a valid format: Encountered unsupported escape sequence"); + } + + output[j++] = SanitizedCharacters[idx]; + isEscaped = false; + } + else if (c == EscapeChar) + { + isEscaped = true; + } + else + { + output[j++] = c; + } + } + } +} diff --git a/dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB/LazyCosmosContainer.cs b/dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB/LazyCosmosContainer.cs new file mode 100644 index 0000000000..1695786636 --- /dev/null +++ b/dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB/LazyCosmosContainer.cs @@ -0,0 +1,80 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.ObjectModel; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Azure.Cosmos; + +namespace Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB; + +#pragma warning disable VSTHRD011 // Use AsyncLazy + +/// +/// A lazy wrapper around a Cosmos DB Container. +/// This avoids performing async I/O-bound operations (i.e. Cosmos DB setup) during +/// DI registration, deferring them until first access. +/// +internal sealed class LazyCosmosContainer +{ + private readonly CosmosClient? _cosmosClient; + private readonly string? _databaseName; + private readonly string? _containerName; + private readonly Lazy> _lazyContainer; + + /// + /// LazyCosmosContainer constructor that initializes the container lazily. + /// + public LazyCosmosContainer(CosmosClient cosmosClient, string databaseName, string containerName) + { + this._cosmosClient = cosmosClient ?? throw new ArgumentNullException(nameof(cosmosClient)); + this._databaseName = databaseName ?? throw new ArgumentNullException(nameof(databaseName)); + this._containerName = containerName ?? throw new ArgumentNullException(nameof(containerName)); + this._lazyContainer = new Lazy>(this.InitializeContainerAsync, LazyThreadSafetyMode.ExecutionAndPublication); + } + + /// + /// LazyCosmosContainer constructor that accepts an existing Container instance. + /// + public LazyCosmosContainer(Container container) + { + if (container is null) + { + throw new ArgumentNullException(nameof(container)); + } + + this._lazyContainer = new Lazy>(() => Task.FromResult(container), LazyThreadSafetyMode.ExecutionAndPublication); + } + + /// + /// Gets the Container, initializing it if necessary. + /// + public Task GetContainerAsync() => this._lazyContainer.Value; + + private async Task InitializeContainerAsync() + { + // Create database if it doesn't exist + var database = await this._cosmosClient!.CreateDatabaseIfNotExistsAsync(this._databaseName!).ConfigureAwait(false); + + var containerProperties = new ContainerProperties(this._containerName!, "/actorId") + { + Id = this._containerName!, + IndexingPolicy = new IndexingPolicy + { + IndexingMode = IndexingMode.Consistent, + Automatic = true + }, + PartitionKeyPaths = ["/actorId"] + }; + + // Add composite index for efficient queries + containerProperties.IndexingPolicy.CompositeIndexes.Add(new Collection + { + new() { Path = "/actorId", Order = CompositePathSortOrder.Ascending }, + new() { Path = "/key", Order = CompositePathSortOrder.Ascending } + }); + + var container = await database.Database.CreateContainerIfNotExistsAsync(containerProperties).ConfigureAwait(false); + return container.Container; + } +} diff --git a/dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.csproj b/dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.csproj new file mode 100644 index 0000000000..4b451d2d35 --- /dev/null +++ b/dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.csproj @@ -0,0 +1,23 @@ + + + + $(ProjectsTargetFrameworks) + $(ProjectsDebugTargetFrameworks) + $(NoWarn);IDE1006;IDE0130 + alpha + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB/ServiceCollectionExtensions.cs b/dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB/ServiceCollectionExtensions.cs new file mode 100644 index 0000000000..1ca7e6c995 --- /dev/null +++ b/dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB/ServiceCollectionExtensions.cs @@ -0,0 +1,93 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System.Text.Json; +using Microsoft.Azure.Cosmos; +using Microsoft.Extensions.DependencyInjection; + +namespace Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB; + +#pragma warning disable VSTHRD002 + +/// +/// Extension methods for configuring Cosmos DB actor state storage in dependency injection. +/// +public static class ServiceCollectionExtensions +{ + /// + /// Adds Cosmos DB actor state storage to the service collection. + /// + /// The service collection to add services to. + /// The Cosmos DB connection string. + /// The database name to use for actor state storage. + /// The container name to use for actor state storage. Defaults to "ActorState". + /// The service collection for chaining. + public static IServiceCollection AddCosmosActorStateStorage( + this IServiceCollection services, + string connectionString, + string databaseName, + string containerName = "ActorState") + { + // Register CosmosClient as singleton + services.AddSingleton(serviceProvider => + { + var cosmosClientOptions = new CosmosClientOptions + { + ApplicationName = "AgentFramework", + ConnectionMode = ConnectionMode.Direct, + ConsistencyLevel = ConsistencyLevel.Session, + UseSystemTextJsonSerializerWithOptions = new JsonSerializerOptions + { + PropertyNamingPolicy = JsonNamingPolicy.CamelCase, + TypeInfoResolver = CosmosActorStateJsonContext.Default + } + }; + + return new CosmosClient(connectionString, cosmosClientOptions); + }); + + // Register LazyCosmosContainer as singleton + services.AddSingleton(serviceProvider => + { + var cosmosClient = serviceProvider.GetRequiredService(); + return new LazyCosmosContainer(cosmosClient, databaseName, containerName); + }); + + // Register the storage implementation + services.AddSingleton(serviceProvider => + { + var lazyContainer = serviceProvider.GetRequiredService(); + return new CosmosActorStateStorage(lazyContainer); + }); + + return services; + } + + /// + /// Adds Cosmos DB actor state storage to the service collection using an existing CosmosClient from DI. + /// + /// The service collection to add services to. + /// The database name to use for actor state storage. + /// The container name to use for actor state storage. Defaults to "ActorState". + /// The service collection for chaining. + public static IServiceCollection AddCosmosActorStateStorage( + this IServiceCollection services, + string databaseName, + string containerName = "ActorState") + { + // Register LazyCosmosContainer as singleton using existing CosmosClient + services.AddSingleton(serviceProvider => + { + var cosmosClient = serviceProvider.GetRequiredService(); + return new LazyCosmosContainer(cosmosClient, databaseName, containerName); + }); + + // Register the storage implementation + services.AddSingleton(serviceProvider => + { + var lazyContainer = serviceProvider.GetRequiredService(); + return new CosmosActorStateStorage(lazyContainer); + }); + + return services; + } +} diff --git a/dotnet/src/Microsoft.Extensions.AI.Agents.Runtime/InProcessActorContext.cs b/dotnet/src/Microsoft.Extensions.AI.Agents.Runtime/InProcessActorContext.cs index e28a61dcb1..ad10dd07e6 100644 --- a/dotnet/src/Microsoft.Extensions.AI.Agents.Runtime/InProcessActorContext.cs +++ b/dotnet/src/Microsoft.Extensions.AI.Agents.Runtime/InProcessActorContext.cs @@ -239,11 +239,23 @@ internal sealed class InProcessActorContext : IActorRuntimeContext, IAsyncDispos // TODO: Turn send & update message operations into storage writes to outbox - var result = await this.Storage.WriteStateAsync( - this.ActorId, - [.. operations.Operations.OfType()], - operations.ETag, - cancellationToken).ConfigureAwait(false); + IReadOnlyCollection writeOps = + [.. operations.Operations.OfType()]; + + WriteResponse result; + if (writeOps.Count == 0) + { + // Nothing to write + result = new WriteResponse(operations.ETag, success: true); + } + else + { + result = await this.Storage.WriteStateAsync( + this.ActorId, + writeOps, + operations.ETag, + cancellationToken).ConfigureAwait(false); + } Log.WriteOperationCompleted(this._logger, this.ActorId.ToString(), result.Success); diff --git a/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests.AppHost/AppHost.cs b/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests.AppHost/AppHost.cs new file mode 100644 index 0000000000..61090b7f5e --- /dev/null +++ b/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests.AppHost/AppHost.cs @@ -0,0 +1,20 @@ +// Copyright (c) Microsoft. All rights reserved. +using Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests; + +var builder = DistributedApplication.CreateBuilder(args); +var cosmosDb = builder.AddAzureCosmosDB(CosmosDBTestConstants.TestCosmosDbName); + +if (CosmosDBTestConstants.UseEmulatorForTesting) +{ + cosmosDb.RunAsEmulator(emulator => emulator.WithLifetime(ContainerLifetime.Persistent)); +} +else +{ + var cosmosDbResource = builder.AddParameterFromConfiguration("CosmosDbName", "CosmosDb:Name"); + var cosmosDbResourceGroup = builder.AddParameterFromConfiguration("CosmosDbResourceGroup", "CosmosDb:ResourceGroup"); + cosmosDb.RunAsExisting(cosmosDbResource, cosmosDbResourceGroup); +} + +cosmosDb.AddCosmosDatabase(CosmosDBTestConstants.TestCosmosDbDatabaseName); + +builder.Build().Run(); diff --git a/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests.AppHost/CosmosDBTestConstants.cs b/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests.AppHost/CosmosDBTestConstants.cs new file mode 100644 index 0000000000..43707f208b --- /dev/null +++ b/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests.AppHost/CosmosDBTestConstants.cs @@ -0,0 +1,18 @@ +// Copyright (c) Microsoft. All rights reserved. + +namespace Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests; + +public static class CosmosDBTestConstants +{ + public const string TestCosmosDbName = "ActorStateStorageTests"; + public const string TestCosmosDbDatabaseName = "state-database"; + + // Set to use the CosmosDB emulator for testing via environment variable. + // Example: set COSMOSDB_TESTS_USE_EMULATOR=true in your environment. + // Warning: Using the emulator may cause test flakiness. + public static bool UseEmulatorForTesting => + string.Equals( + Environment.GetEnvironmentVariable("COSMOSDB_TESTS_USE_EMULATOR"), + "true", + StringComparison.OrdinalIgnoreCase); +} diff --git a/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests.AppHost/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests.AppHost.csproj b/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests.AppHost/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests.AppHost.csproj new file mode 100644 index 0000000000..8f7ef08ad8 --- /dev/null +++ b/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests.AppHost/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests.AppHost.csproj @@ -0,0 +1,20 @@ + + + + + + + net9.0 + enable + enable + 80048040-aaf1-4f44-9970-8aef39651edf + true + false + + + + + + + + diff --git a/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests.AppHost/Properties/launchSettings.json b/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests.AppHost/Properties/launchSettings.json new file mode 100644 index 0000000000..596bd53611 --- /dev/null +++ b/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests.AppHost/Properties/launchSettings.json @@ -0,0 +1,29 @@ +{ + "$schema": "https://json.schemastore.org/launchsettings.json", + "profiles": { + "https": { + "commandName": "Project", + "dotnetRunMessages": true, + "launchBrowser": true, + "applicationUrl": "https://localhost:17163;http://localhost:15113", + "environmentVariables": { + "ASPNETCORE_ENVIRONMENT": "Development", + "DOTNET_ENVIRONMENT": "Development", + "ASPIRE_DASHBOARD_OTLP_ENDPOINT_URL": "https://localhost:21207", + "ASPIRE_RESOURCE_SERVICE_ENDPOINT_URL": "https://localhost:22258" + } + }, + "http": { + "commandName": "Project", + "dotnetRunMessages": true, + "launchBrowser": true, + "applicationUrl": "http://localhost:15113", + "environmentVariables": { + "ASPNETCORE_ENVIRONMENT": "Development", + "DOTNET_ENVIRONMENT": "Development", + "ASPIRE_DASHBOARD_OTLP_ENDPOINT_URL": "http://localhost:19080", + "ASPIRE_RESOURCE_SERVICE_ENDPOINT_URL": "http://localhost:20201" + } + } + } +} diff --git a/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests/CosmosActorStateStorageConcurrencyTests.cs b/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests/CosmosActorStateStorageConcurrencyTests.cs new file mode 100644 index 0000000000..fd3fafaa23 --- /dev/null +++ b/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests/CosmosActorStateStorageConcurrencyTests.cs @@ -0,0 +1,250 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System.Text.Json; + +namespace Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests; + +#pragma warning disable CA5394 // Insecure randomness is okay for test purposes + +/// +/// Integration tests for CosmosActorStateStorage focusing on concurrency control and ETag progression. +/// +[Collection("Cosmos Test Collection")] +public class CosmosActorStateStorageConcurrencyTests +{ + private readonly CosmosTestFixture _fixture; + + public CosmosActorStateStorageConcurrencyTests(CosmosTestFixture fixture) + { + this._fixture = fixture; + } + + private static readonly TimeSpan s_defaultTimeout = TimeSpan.FromSeconds(300); + + [Fact] + public async Task ETagProgression_ShouldChangeWithEachWriteAsync() + { + // CosmosDB ETags are not guaranteed to be numeric or monotonically increasing + // They are opaque strings that change with each update, which is sufficient for optimistic concurrency + + // Arrange + using var cts = new CancellationTokenSource(s_defaultTimeout); + var cancellationToken = cts.Token; + + var storage = new CosmosActorStateStorage(this._fixture.Container); + var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString()); + + var key = "testKey"; + var value1 = JsonSerializer.SerializeToElement("value1"); + var value2 = JsonSerializer.SerializeToElement("value2"); + + // Act - First write + var operations1 = new List { new SetValueOperation(key, value1) }; + var result1 = await storage.WriteStateAsync(testActorId, operations1, "0", cancellationToken); + + // Act - Second write + var operations2 = new List { new SetValueOperation(key, value2) }; + var result2 = await storage.WriteStateAsync(testActorId, operations2, result1.ETag, cancellationToken); + + // Act - Third write + var operations3 = new List { new RemoveKeyOperation(key) }; + var result3 = await storage.WriteStateAsync(testActorId, operations3, result2.ETag, cancellationToken); + + // Assert + Assert.True(result1.Success); + Assert.True(result2.Success); + Assert.True(result3.Success); + Assert.NotEqual("0", result1.ETag); + Assert.NotEqual(result1.ETag, result2.ETag); + Assert.NotEqual(result2.ETag, result3.ETag); + + // Verify ETags are all different and represent progression + string[] etags = [result1.ETag, result2.ETag, result3.ETag]; + Assert.Equal(3, etags.Distinct().Count()); + } + + [Fact] + public async Task ConcurrentWrites_ShouldHandleOptimisticConcurrencyCorrectlyAsync() + { + // Arrange + using var cts = new CancellationTokenSource(s_defaultTimeout); + var cancellationToken = cts.Token; + + var storage = new CosmosActorStateStorage(this._fixture.Container); + var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString()); + + // Setup initial state + var initialOperations = new List + { + new SetValueOperation("counter", JsonSerializer.SerializeToElement(0)) + }; + var initialResult = await storage.WriteStateAsync(testActorId, initialOperations, "0", cancellationToken); + Assert.True(initialResult.Success); + + const int ConcurrentOperations = 10; + var tasks = new List>(); + + // Act - Simulate concurrent writes with retry logic + for (int i = 0; i < ConcurrentOperations; i++) + { + var operationNumber = i; + tasks.Add(Task.Run(async () => + { + var success = false; + var retryCount = 0; + const int MaxRetries = 20; + string? finalETag = null; + + while (!success && retryCount < MaxRetries) + { + try + { + // Read current state to get latest ETag + var readOps = new List + { + new GetValueOperation("counter") + }; + var readResult = await storage.ReadStateAsync(testActorId, readOps, cancellationToken); + var currentETag = readResult.ETag; + + var currentValue = readResult.Results[0] as GetValueResult; + var currentCounter = currentValue?.Value?.GetInt32() ?? 0; + + // Try to increment the counter + var writeOps = new List + { + new SetValueOperation("counter", JsonSerializer.SerializeToElement(currentCounter + 1)), + new SetValueOperation($"operation_{operationNumber}", JsonSerializer.SerializeToElement($"completed_attempt_{retryCount}")) + }; + + var writeResult = await storage.WriteStateAsync(testActorId, writeOps, currentETag, cancellationToken); + + if (writeResult.Success) + { + success = true; + finalETag = writeResult.ETag; + } + else + { + retryCount++; + // Small delay to reduce contention + await Task.Delay(Random.Shared.Next(1, 10), cancellationToken); + } + } + catch (Exception) + { + retryCount++; + await Task.Delay(Random.Shared.Next(1, 10), cancellationToken); + } + } + + return (success, finalETag, retryCount); + })); + } + + // Wait for all operations to complete + var results = await Task.WhenAll(tasks); + + // Assert - All operations should eventually succeed + Assert.All(results, result => Assert.True(result.Success, $"Operation failed after {result.AttemptNumber} attempts")); + + // Act - Verify final state + var finalReadOps = new List + { + new GetValueOperation("counter"), + new ListKeysOperation(continuationToken: null) + }; + var finalResult = await storage.ReadStateAsync(testActorId, finalReadOps, cancellationToken); + + var finalCounter = finalResult.Results[0] as GetValueResult; + var finalKeys = finalResult.Results[1] as ListKeysResult; + + // Assert final state is consistent + Assert.NotNull(finalCounter); + Assert.NotNull(finalKeys); + Assert.Equal(ConcurrentOperations, finalCounter.Value?.GetInt32()); // Counter should equal number of operations + Assert.Equal(ConcurrentOperations + 1, finalKeys.Keys.Count); // counter + operation_N keys + + // Verify all operation keys are present + Assert.Contains("counter", finalKeys.Keys); + for (int i = 0; i < ConcurrentOperations; i++) + { + Assert.Contains($"operation_{i}", finalKeys.Keys); + } + + // Log retry statistics for debugging + var totalRetries = results.Sum(r => r.AttemptNumber); + var maxRetries = results.Max(r => r.AttemptNumber); + Console.WriteLine($"Concurrent operations completed. Total retries: {totalRetries}, Max retries for single operation: {maxRetries}"); + } + + [Fact] + public async Task WriteStateAsync_InitialETagHandling_ShouldWorkCorrectlyAsync() + { + // Arrange + using var cts = new CancellationTokenSource(s_defaultTimeout); + var cancellationToken = cts.Token; + + var storage = new CosmosActorStateStorage(this._fixture.Container); + var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString()); + + var key = "testKey"; + var value = JsonSerializer.SerializeToElement("testValue"); + var operations = new List + { + new SetValueOperation(key, value) + }; + + // Act & Assert - Test null eTag (should create new document) +#pragma warning disable CS8625 // Cannot convert null literal to non-nullable reference type. + var resultWithNullETag = await storage.WriteStateAsync(testActorId, operations, null, cancellationToken); +#pragma warning restore CS8625 // Cannot convert null literal to non-nullable reference type. + Assert.True(resultWithNullETag.Success); + Assert.NotNull(resultWithNullETag.ETag); + Assert.NotEmpty(resultWithNullETag.ETag); + + // Clean up for next test + var uniqueActorId1 = new ActorId("TestActor", Guid.NewGuid().ToString()); + + // Act & Assert - Test empty eTag (should create new document) + var resultWithEmptyETag = await storage.WriteStateAsync(uniqueActorId1, operations, string.Empty, cancellationToken); + Assert.True(resultWithEmptyETag.Success); + Assert.NotNull(resultWithEmptyETag.ETag); + Assert.NotEmpty(resultWithEmptyETag.ETag); + + // Clean up for next test + var uniqueActorId2 = new ActorId("TestActor", Guid.NewGuid().ToString()); + + // Act & Assert - Test "0" initial eTag (should create new document) + var resultWithInitialETag = await storage.WriteStateAsync(uniqueActorId2, operations, "0", cancellationToken); + Assert.True(resultWithInitialETag.Success); + Assert.NotNull(resultWithInitialETag.ETag); + Assert.NotEmpty(resultWithInitialETag.ETag); + Assert.NotEqual("0", resultWithInitialETag.ETag); + + // Act & Assert - Test writing again with "0" should fail (document already exists) + var secondWriteWithInitialETag = await storage.WriteStateAsync(uniqueActorId2, operations, "0", cancellationToken); + Assert.False(secondWriteWithInitialETag.Success); + Assert.Empty(secondWriteWithInitialETag.ETag); + + // Act & Assert - Test writing with correct eTag should succeed + var updateOperations = new List + { + new SetValueOperation(key, JsonSerializer.SerializeToElement("updatedValue")) + }; + var resultWithCorrectETag = await storage.WriteStateAsync(uniqueActorId2, updateOperations, resultWithInitialETag.ETag, cancellationToken); + Assert.True(resultWithCorrectETag.Success); + Assert.NotNull(resultWithCorrectETag.ETag); + Assert.NotEqual(resultWithInitialETag.ETag, resultWithCorrectETag.ETag); + + // Verify the value was actually updated + var readOperations = new List + { + new GetValueOperation(key) + }; + var readResult = await storage.ReadStateAsync(uniqueActorId2, readOperations, cancellationToken); + var getValue = readResult.Results[0] as GetValueResult; + Assert.NotNull(getValue); + Assert.Equal("updatedValue", getValue.Value?.GetString()); + } +} diff --git a/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests/CosmosActorStateStorageListKeysTests.cs b/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests/CosmosActorStateStorageListKeysTests.cs new file mode 100644 index 0000000000..87d3c2fbb9 --- /dev/null +++ b/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests/CosmosActorStateStorageListKeysTests.cs @@ -0,0 +1,290 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System.Text.Json; + +namespace Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests; + +/// +/// Integration tests for CosmosActorStateStorage focusing on ListKeys functionality. +/// +[Collection("Cosmos Test Collection")] +public class CosmosActorStateStorageListKeysTests +{ + private readonly CosmosTestFixture _fixture; + + public CosmosActorStateStorageListKeysTests(CosmosTestFixture fixture) + { + this._fixture = fixture; + } + + private static readonly TimeSpan s_defaultTimeout = TimeSpan.FromSeconds(300); + + [Fact] + public async Task ReadStateAsync_WithListKeysAndKeyPrefix_ShouldReturnFilteredKeysAsync() + { + // Arrange + using var cts = new CancellationTokenSource(s_defaultTimeout); + var cancellationToken = cts.Token; + + var storage = new CosmosActorStateStorage(this._fixture.Container); + var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString()); + + var prefixKey1 = "prefix_key1"; + var prefixKey2 = "prefix_key2"; + var otherKey = "other_key"; + var value1 = JsonSerializer.SerializeToElement("value1"); + var value2 = JsonSerializer.SerializeToElement("value2"); + var value3 = JsonSerializer.SerializeToElement("value3"); + + var writeOperations = new List + { + new SetValueOperation(prefixKey1, value1), + new SetValueOperation(prefixKey2, value2), + new SetValueOperation(otherKey, value3) + }; + + await storage.WriteStateAsync(testActorId, writeOperations, "0", cancellationToken); + + // Act - List keys with prefix filter + var readOperations = new List + { + new ListKeysOperation(continuationToken: null, keyPrefix: "prefix_") + }; + var result = await storage.ReadStateAsync(testActorId, readOperations, cancellationToken); + + // Assert + Assert.Single(result.Results); + var listKeys = result.Results[0] as ListKeysResult; + Assert.NotNull(listKeys); + Assert.Equal(2, listKeys.Keys.Count); + Assert.Contains(prefixKey1, listKeys.Keys); + Assert.Contains(prefixKey2, listKeys.Keys); + Assert.DoesNotContain(otherKey, listKeys.Keys); + } + + [Fact] + public async Task ReadStateAsync_WithListKeysAndNonMatchingPrefix_ShouldReturnEmptyListAsync() + { + // Arrange + using var cts = new CancellationTokenSource(s_defaultTimeout); + var cancellationToken = cts.Token; + + var storage = new CosmosActorStateStorage(this._fixture.Container); + var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString()); + + var key1 = "key1"; + var key2 = "key2"; + var value1 = JsonSerializer.SerializeToElement("value1"); + var value2 = JsonSerializer.SerializeToElement("value2"); + + var writeOperations = new List + { + new SetValueOperation(key1, value1), + new SetValueOperation(key2, value2) + }; + + await storage.WriteStateAsync(testActorId, writeOperations, "0", cancellationToken); + + // Act - List keys with non-matching prefix + var readOperations = new List + { + new ListKeysOperation(continuationToken: null, keyPrefix: "nonexistent_") + }; + var result = await storage.ReadStateAsync(testActorId, readOperations, cancellationToken); + + // Assert + Assert.Single(result.Results); + var listKeys = result.Results[0] as ListKeysResult; + Assert.NotNull(listKeys); + Assert.Empty(listKeys.Keys); + Assert.Null(listKeys.ContinuationToken); + } + + [Fact] + public async Task ReadStateAsync_WithListKeysForEmptyActor_ShouldReturnEmptyListAsync() + { + // Arrange + using var cts = new CancellationTokenSource(s_defaultTimeout); + var cancellationToken = cts.Token; + + var storage = new CosmosActorStateStorage(this._fixture.Container); + var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString()); + + // Act - List keys for actor with no state + var readOperations = new List + { + new ListKeysOperation(continuationToken: null) + }; + var result = await storage.ReadStateAsync(testActorId, readOperations, cancellationToken); + + // Assert + Assert.Single(result.Results); + var listKeys = result.Results[0] as ListKeysResult; + Assert.NotNull(listKeys); + Assert.Empty(listKeys.Keys); + Assert.Null(listKeys.ContinuationToken); + } + + [Fact] + public async Task ReadStateAsync_WithListKeysOperation_ShouldReturnAllKeysAsync() + { + // Arrange + using var cts = new CancellationTokenSource(s_defaultTimeout); + var cancellationToken = cts.Token; + + var storage = new CosmosActorStateStorage(this._fixture.Container); + var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString()); + + var key1 = "key1"; + var key2 = "key2"; + var value1 = JsonSerializer.SerializeToElement("value1"); + var value2 = JsonSerializer.SerializeToElement("value2"); + + var writeOperations = new List + { + new SetValueOperation(key1, value1), + new SetValueOperation(key2, value2) + }; + + // First write some data + var writeResult = await storage.WriteStateAsync(testActorId, writeOperations, "0", cancellationToken); + Assert.True(writeResult.Success); + + // Act - List keys + var readOperations = new List + { + new ListKeysOperation(continuationToken: null) + }; + var readResult = await storage.ReadStateAsync(testActorId, readOperations, cancellationToken); + + // Assert + Assert.Single(readResult.Results); + var listKeys = readResult.Results[0] as ListKeysResult; + Assert.NotNull(listKeys); + Assert.Equal(2, listKeys.Keys.Count); + Assert.Contains(key1, listKeys.Keys); + Assert.Contains(key2, listKeys.Keys); + } + + [Fact] + public async Task ReadStateAsync_WithListKeysAfterKeyRemoval_ShouldNotIncludeRemovedKeysAsync() + { + // Arrange + using var cts = new CancellationTokenSource(s_defaultTimeout); + var cancellationToken = cts.Token; + + var storage = new CosmosActorStateStorage(this._fixture.Container); + var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString()); + + var key1 = "key1"; + var key2 = "key2"; + var key3 = "key3"; + var value1 = JsonSerializer.SerializeToElement("value1"); + var value2 = JsonSerializer.SerializeToElement("value2"); + var value3 = JsonSerializer.SerializeToElement("value3"); + + // Setup initial state with 3 keys + var writeOperations = new List + { + new SetValueOperation(key1, value1), + new SetValueOperation(key2, value2), + new SetValueOperation(key3, value3) + }; + var writeResult = await storage.WriteStateAsync(testActorId, writeOperations, "0", cancellationToken); + Assert.True(writeResult.Success); + + // Remove one key + var removeOperations = new List + { + new RemoveKeyOperation(key2) + }; + var removeResult = await storage.WriteStateAsync(testActorId, removeOperations, writeResult.ETag, cancellationToken); + Assert.True(removeResult.Success); + + // Act - List keys after removal + var readOperations = new List + { + new ListKeysOperation(continuationToken: null) + }; + var readResult = await storage.ReadStateAsync(testActorId, readOperations, cancellationToken); + + // Assert - Only remaining keys should be listed + Assert.Single(readResult.Results); + var listKeys = readResult.Results[0] as ListKeysResult; + Assert.NotNull(listKeys); + Assert.Equal(2, listKeys.Keys.Count); + Assert.Contains(key1, listKeys.Keys); + Assert.Contains(key3, listKeys.Keys); + Assert.DoesNotContain(key2, listKeys.Keys); + } + + [Fact] + public async Task ReadStateAsync_WithListKeysAndMultiplePrefixes_ShouldFilterCorrectlyAsync() + { + // Arrange + using var cts = new CancellationTokenSource(s_defaultTimeout); + var cancellationToken = cts.Token; + + var storage = new CosmosActorStateStorage(this._fixture.Container); + var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString()); + + // Create keys with different prefixes + string[] userKeys = ["user_profile", "user_settings", "user_preferences"]; + string[] sessionKeys = ["session_token", "session_data"]; + string[] cacheKeys = ["cache_item1", "cache_item2", "cache_item3"]; + string[] miscKeys = ["config", "metadata"]; + + var writeOperations = new List(); + var allKeys = userKeys.Concat(sessionKeys).Concat(cacheKeys).Concat(miscKeys); + + foreach (var key in allKeys) + { + writeOperations.Add(new SetValueOperation(key, JsonSerializer.SerializeToElement($"value_for_{key}"))); + } + + await storage.WriteStateAsync(testActorId, writeOperations, "0", cancellationToken); + + // Test user_ prefix + var userReadOps = new List + { + new ListKeysOperation(continuationToken: null, keyPrefix: "user_") + }; + var userResult = await storage.ReadStateAsync(testActorId, userReadOps, cancellationToken); + var userListKeys = userResult.Results[0] as ListKeysResult; + Assert.NotNull(userListKeys); + Assert.Equal(3, userListKeys.Keys.Count); + Assert.All(userKeys, key => Assert.Contains(key, userListKeys.Keys)); + + // Test session_ prefix + var sessionReadOps = new List + { + new ListKeysOperation(continuationToken: null, keyPrefix: "session_") + }; + var sessionResult = await storage.ReadStateAsync(testActorId, sessionReadOps, cancellationToken); + var sessionListKeys = sessionResult.Results[0] as ListKeysResult; + Assert.NotNull(sessionListKeys); + Assert.Equal(2, sessionListKeys.Keys.Count); + Assert.All(sessionKeys, key => Assert.Contains(key, sessionListKeys.Keys)); + + // Test cache_ prefix + var cacheReadOps = new List + { + new ListKeysOperation(continuationToken: null, keyPrefix: "cache_") + }; + var cacheResult = await storage.ReadStateAsync(testActorId, cacheReadOps, cancellationToken); + var cacheListKeys = cacheResult.Results[0] as ListKeysResult; + Assert.NotNull(cacheListKeys); + Assert.Equal(3, cacheListKeys.Keys.Count); + Assert.All(cacheKeys, key => Assert.Contains(key, cacheListKeys.Keys)); + + // Test no prefix (should return all keys) + var allReadOps = new List + { + new ListKeysOperation(continuationToken: null) + }; + var allResult = await storage.ReadStateAsync(testActorId, allReadOps, cancellationToken); + var allListKeys = allResult.Results[0] as ListKeysResult; + Assert.NotNull(allListKeys); + Assert.Equal(10, allListKeys.Keys.Count); // 3 + 2 + 3 + 2 = 10 total keys + } +} diff --git a/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests/CosmosActorStateStorageTests.cs b/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests/CosmosActorStateStorageTests.cs new file mode 100644 index 0000000000..14d6529b13 --- /dev/null +++ b/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests/CosmosActorStateStorageTests.cs @@ -0,0 +1,492 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System.Text.Json; + +namespace Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests; + +/// +/// Integration tests for CosmosActorStateStorage covering basic CRUD operations and advanced scenarios. +/// +[Collection("Cosmos Test Collection")] +public class CosmosActorStateStorageTests +{ + private readonly CosmosTestFixture _fixture; + + public CosmosActorStateStorageTests(CosmosTestFixture fixture) + { + this._fixture = fixture; + } + + private static readonly TimeSpan s_defaultTimeout = TimeSpan.FromSeconds(300); + + [Fact] + public async Task WriteStateAsync_WithSetValueOperation_ShouldStoreValueAsync() + { + // Arrange + using var cts = new CancellationTokenSource(s_defaultTimeout); + var cancellationToken = cts.Token; + + var storage = new CosmosActorStateStorage(this._fixture.Container); + var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString()); + + var key = "testKey"; + var value = JsonSerializer.SerializeToElement("testValue"); + + var operations = new List + { + new SetValueOperation(key, value) + }; + + // Act + var result = await storage.WriteStateAsync(testActorId, operations, "0", cancellationToken); + + // Assert + Assert.True(result.Success); + Assert.NotEqual("0", result.ETag); + } + + [Fact] + public async Task WriteAndReadState_WithMultipleOperations_ShouldMaintainConsistencyAsync() + { + // Arrange + using var cts = new CancellationTokenSource(s_defaultTimeout); + var cancellationToken = cts.Token; + + var storage = new CosmosActorStateStorage(this._fixture.Container); + var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString()); + + var key1 = "key1"; + var key2 = "key2"; + var value1 = JsonSerializer.SerializeToElement("value1"); + var value2 = JsonSerializer.SerializeToElement(42); + + var writeOperations = new List + { + new SetValueOperation(key1, value1), + new SetValueOperation(key2, value2) + }; + + // Act - Write state + var writeResult = await storage.WriteStateAsync(testActorId, writeOperations, "0", cancellationToken); + + // Assert write succeeded + Assert.True(writeResult.Success); + Assert.NotNull(writeResult.ETag); + Assert.NotEqual("0", writeResult.ETag); + + // Act - Read individual values + var readOperations = new List + { + new GetValueOperation(key1), + new GetValueOperation(key2) + }; + var readResult = await storage.ReadStateAsync(testActorId, readOperations, cancellationToken); + + // Assert read succeeded and values match + Assert.Equal(2, readResult.Results.Count); + + var getValue1 = readResult.Results[0] as GetValueResult; + var getValue2 = readResult.Results[1] as GetValueResult; + + Assert.NotNull(getValue1); + Assert.NotNull(getValue2); + Assert.Equal("value1", getValue1.Value?.GetString()); + Assert.Equal(42, getValue2.Value?.GetInt32()); + + // Act - List keys + var listKeysOperations = new List + { + new ListKeysOperation(continuationToken: null) + }; + var listResult = await storage.ReadStateAsync(testActorId, listKeysOperations, cancellationToken); + + // Assert keys are listed correctly + Assert.Single(listResult.Results); + var listKeys = listResult.Results[0] as ListKeysResult; + Assert.NotNull(listKeys); + Assert.Equal(2, listKeys.Keys.Count); + Assert.Contains(key1, listKeys.Keys); + Assert.Contains(key2, listKeys.Keys); + + // Act - Update with correct ETag + var updateOperations = new List + { + new SetValueOperation(key1, JsonSerializer.SerializeToElement("updated_value1")), + new RemoveKeyOperation(key2) + }; + var updateResult = await storage.WriteStateAsync(testActorId, updateOperations, writeResult.ETag, cancellationToken); + + // Assert update succeeded + Assert.True(updateResult.Success); + Assert.NotEqual(writeResult.ETag, updateResult.ETag); + + // Act - Verify final state + var finalReadOperations = new List + { + new GetValueOperation(key1), + new GetValueOperation(key2), + new ListKeysOperation(continuationToken: null) + }; + var finalResult = await storage.ReadStateAsync(testActorId, finalReadOperations, cancellationToken); + + // Assert final state is correct + Assert.Equal(3, finalResult.Results.Count); + + var finalValue1 = finalResult.Results[0] as GetValueResult; + var finalValue2 = finalResult.Results[1] as GetValueResult; + var finalKeys = finalResult.Results[2] as ListKeysResult; + + Assert.NotNull(finalValue1); + Assert.NotNull(finalValue2); + Assert.NotNull(finalKeys); + + Assert.Equal("updated_value1", finalValue1.Value?.GetString()); + Assert.Null(finalValue2.Value); // key2 was removed + Assert.Single(finalKeys.Keys); + Assert.Contains(key1, finalKeys.Keys); + Assert.DoesNotContain(key2, finalKeys.Keys); + } + + [Fact] + public async Task WriteStateAsync_WithIncorrectETag_ShouldReturnFailureAsync() + { + // Arrange + using var cts = new CancellationTokenSource(s_defaultTimeout); + var cancellationToken = cts.Token; + + var storage = new CosmosActorStateStorage(this._fixture.Container); + var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString()); + + var key = "testKey"; + var value = JsonSerializer.SerializeToElement("testValue"); + var operations = new List + { + new SetValueOperation(key, value) + }; + + // First write to establish state + var firstResult = await storage.WriteStateAsync(testActorId, operations, "0", cancellationToken); + Assert.True(firstResult.Success); + + // Act - Try to write with incorrect ETag + var incorrectOperations = new List + { + new SetValueOperation(key, JsonSerializer.SerializeToElement("newValue")) + }; + var result = await storage.WriteStateAsync(testActorId, incorrectOperations, "incorrect-etag", cancellationToken); + + // Assert + Assert.False(result.Success); + Assert.Empty(result.ETag); + + // Verify original value is unchanged + var readOperations = new List + { + new GetValueOperation(key) + }; + var readResult = await storage.ReadStateAsync(testActorId, readOperations, cancellationToken); + var getValue = readResult.Results[0] as GetValueResult; + Assert.Equal("testValue", getValue?.Value?.GetString()); + } + + [Fact] + public async Task DifferentActors_ShouldHaveIsolatedStateAsync() + { + // Arrange + using var cts = new CancellationTokenSource(s_defaultTimeout); + var cancellationToken = cts.Token; + + var storage = new CosmosActorStateStorage(this._fixture.Container); + var testActorId1 = new ActorId("TestActor1", Guid.NewGuid().ToString()); + var testActorId2 = new ActorId("TestActor2", Guid.NewGuid().ToString()); + + var key = "sharedKey"; + var value1 = JsonSerializer.SerializeToElement("value1"); + var value2 = JsonSerializer.SerializeToElement("value2"); + + var operations1 = new List + { + new SetValueOperation(key, value1) + }; + var operations2 = new List + { + new SetValueOperation(key, value2) + }; + + // Act - Write to both actors + await storage.WriteStateAsync(testActorId1, operations1, "0", cancellationToken); + await storage.WriteStateAsync(testActorId2, operations2, "0", cancellationToken); + + // Assert - Verify values are different + var readOperations = new List + { + new GetValueOperation(key) + }; + + var result1 = await storage.ReadStateAsync(testActorId1, readOperations, cancellationToken); + var result2 = await storage.ReadStateAsync(testActorId2, readOperations, cancellationToken); + + var getValue1 = result1.Results[0] as GetValueResult; + var getValue2 = result2.Results[0] as GetValueResult; + + Assert.NotNull(getValue1); + Assert.NotNull(getValue2); + Assert.Equal("value1", getValue1.Value?.GetString()); + Assert.Equal("value2", getValue2.Value?.GetString()); + Assert.NotEqual(result1.ETag, result2.ETag); + } + + [Fact] + public async Task WriteStateAsync_WithEmptyOperations_ShouldThrowExceptionAsync() + { + // Arrange + using var cts = new CancellationTokenSource(s_defaultTimeout); + var cancellationToken = cts.Token; + var storage = new CosmosActorStateStorage(this._fixture.Container); + var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString()); + var emptyOperations = new List(); + // Act & Assert + await Assert.ThrowsAsync(async () => + { + await storage.WriteStateAsync(testActorId, emptyOperations, "0", cancellationToken); + }); + } + + [Fact] + public async Task ReadStateAsync_WithGetValueForNonExistentKey_ShouldReturnNullAsync() + { + // Arrange + using var cts = new CancellationTokenSource(s_defaultTimeout); + var cancellationToken = cts.Token; + + var storage = new CosmosActorStateStorage(this._fixture.Container); + var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString()); + + var readOperations = new List + { + new GetValueOperation("nonExistentKey") + }; + + // Act + var result = await storage.ReadStateAsync(testActorId, readOperations, cancellationToken); + + // Assert + Assert.Single(result.Results); + var getValue = result.Results[0] as GetValueResult; + Assert.NotNull(getValue); + Assert.Null(getValue.Value); + } + + [Fact] + public async Task WriteStateAsync_WithComplexJsonValue_ShouldSerializeCorrectlyAsync() + { + // Arrange + using var cts = new CancellationTokenSource(s_defaultTimeout); + var cancellationToken = cts.Token; + + var storage = new CosmosActorStateStorage(this._fixture.Container); + var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString()); + + // Create a complex object with various types +#pragma warning disable CA1861 // Avoid constant arrays as arguments + var complexObject = new + { + Id = 123, + Name = "Test Object", + Properties = new Dictionary + { + { "StringProp", "value" }, + { "NumberProp", 42.5 }, + { "BoolProp", true }, + { "ArrayProp", new[] { 1, 2, 3 } }, + { "NestedProp", new { Inner = "nested value" } } + }, + Tags = new[] { "tag1", "tag2", "tag3" }, + Metadata = new Dictionary + { + { "version", "1.0" }, + { "author", "test" } + } + }; +#pragma warning restore CA1861 // Avoid constant arrays as arguments + + var key = "complexObject"; + var value = JsonSerializer.SerializeToElement(complexObject); + + var operations = new List + { + new SetValueOperation(key, value) + }; + + // Act - Write complex object + var writeResult = await storage.WriteStateAsync(testActorId, operations, "0", cancellationToken); + Assert.True(writeResult.Success); + + // Act - Read back complex object + var readOperations = new List + { + new GetValueOperation(key) + }; + var readResult = await storage.ReadStateAsync(testActorId, readOperations, cancellationToken); + + // Assert - Verify complex object was stored and retrieved correctly + Assert.Single(readResult.Results); + var getValue = readResult.Results[0] as GetValueResult; + Assert.NotNull(getValue); + Assert.NotNull(getValue.Value); + + // Deserialize and verify structure + var retrievedObject = JsonSerializer.Deserialize(getValue.Value!.Value.GetRawText()); + Assert.Equal(123, retrievedObject.GetProperty("Id").GetInt32()); + Assert.Equal("Test Object", retrievedObject.GetProperty("Name").GetString()); + + var properties = retrievedObject.GetProperty("Properties"); + Assert.Equal("value", properties.GetProperty("StringProp").GetString()); + Assert.Equal(42.5, properties.GetProperty("NumberProp").GetDouble()); + Assert.True(properties.GetProperty("BoolProp").GetBoolean()); + + var tags = retrievedObject.GetProperty("Tags"); + Assert.Equal(3, tags.GetArrayLength()); + Assert.Equal("tag1", tags[0].GetString()); + } + + [Fact] + public async Task MultipleOperationsInSequence_ShouldBeProcessedInOrderAsync() + { + // Arrange + using var cts = new CancellationTokenSource(s_defaultTimeout); + var cancellationToken = cts.Token; + + var storage = new CosmosActorStateStorage(this._fixture.Container); + var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString()); + + var key1 = "key1"; + var key2 = "key2"; + var key3 = "key3"; + var value1 = JsonSerializer.SerializeToElement("value1"); + var value2 = JsonSerializer.SerializeToElement("value2"); + var value3 = JsonSerializer.SerializeToElement("value3"); + + // Act - Perform multiple operations in a single batch + var operations = new List + { + new SetValueOperation(key1, value1), // Set key1 + new SetValueOperation(key2, value2), // Set key2 + new SetValueOperation(key3, value3), // Set key3 + new RemoveKeyOperation(key1), // Remove key1 + new SetValueOperation(key1, JsonSerializer.SerializeToElement("new_value1")) // Re-add key1 with new value + }; + + var result = await storage.WriteStateAsync(testActorId, operations, "0", cancellationToken); + + // Assert write succeeded + Assert.True(result.Success); + + // Act - Verify final state + var readOperations = new List + { + new GetValueOperation(key1), + new GetValueOperation(key2), + new GetValueOperation(key3), + new ListKeysOperation(continuationToken: null) + }; + var readResult = await storage.ReadStateAsync(testActorId, readOperations, cancellationToken); + + // Assert final state is correct + Assert.Equal(4, readResult.Results.Count); + + var getValue1 = readResult.Results[0] as GetValueResult; + var getValue2 = readResult.Results[1] as GetValueResult; + var getValue3 = readResult.Results[2] as GetValueResult; + var listKeys = readResult.Results[3] as ListKeysResult; + + Assert.NotNull(getValue1); + Assert.NotNull(getValue2); + Assert.NotNull(getValue3); + Assert.NotNull(listKeys); + + // key1 should have the final value from the last operation + Assert.Equal("new_value1", getValue1.Value?.GetString()); + Assert.Equal("value2", getValue2.Value?.GetString()); + Assert.Equal("value3", getValue3.Value?.GetString()); + + // All three keys should be present + Assert.Equal(3, listKeys.Keys.Count); + Assert.Contains(key1, listKeys.Keys); + Assert.Contains(key2, listKeys.Keys); + Assert.Contains(key3, listKeys.Keys); + } + + [SkipOnEmulatorFact] + public async Task WriteAndReadState_WithSpecialCharactersInKeys_ShouldHandleSanitizationAsync() + { + // Arrange + using var cts = new CancellationTokenSource(s_defaultTimeout); + var cancellationToken = cts.Token; + + var storage = new CosmosActorStateStorage(this._fixture.Container); + var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString()); + + // Test keys with special characters that need sanitization + var specialKeys = new[] + { + "key/with/slashes", + "key with spaces", + "key:with:colons", + "key@with@symbols", + "key%with%percent", + "key#with#hash", + "key?with?query", + "key&with&ersand" + }; + + var writeOperations = new List(); + for (int i = 0; i < specialKeys.Length; i++) + { + var value = JsonSerializer.SerializeToElement($"value{i}"); + writeOperations.Add(new SetValueOperation(specialKeys[i], value)); + } + + // Act - Write keys with special characters + var writeResult = await storage.WriteStateAsync(testActorId, writeOperations, "0", cancellationToken); + + // Assert write succeeded + Assert.True(writeResult.Success); + Assert.NotNull(writeResult.ETag); + + // Act - Read back each key individually + for (int i = 0; i < specialKeys.Length; i++) + { + var readOperations = new List + { + new GetValueOperation(specialKeys[i]) + }; + var readResult = await storage.ReadStateAsync(testActorId, readOperations, cancellationToken); + + // Assert each key can be read back correctly + Assert.Single(readResult.Results); + var getValue = readResult.Results[0] as GetValueResult; + Assert.NotNull(getValue); + Assert.NotNull(getValue.Value); + Assert.Equal($"value{i}", getValue.Value?.GetString()); + } + + // Act - List all keys + var listOperations = new List + { + new ListKeysOperation(continuationToken: null) + }; + var listResult = await storage.ReadStateAsync(testActorId, listOperations, cancellationToken); + + // Assert all keys are present in the list + Assert.Single(listResult.Results); + var listKeys = listResult.Results[0] as ListKeysResult; + Assert.NotNull(listKeys); + Assert.Equal(specialKeys.Length, listKeys.Keys.Count); + + foreach (var specialKey in specialKeys) + { + Assert.Contains(specialKey, listKeys.Keys); + } + } +} diff --git a/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests/CosmosIdSanitizerTests.cs b/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests/CosmosIdSanitizerTests.cs new file mode 100644 index 0000000000..18966a9d0a --- /dev/null +++ b/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests/CosmosIdSanitizerTests.cs @@ -0,0 +1,268 @@ +// Copyright (c) Microsoft. All rights reserved. + +namespace Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests; + +public class CosmosIdSanitizerTests +{ + [Fact] + public void Sanitize_WithValidInputNoSpecialChars_ReturnsOriginalString() + { + // Arrange + const string Input = "ValidId123"; + + // Act + string result = CosmosIdSanitizer.Sanitize(Input); + + // Assert + Assert.Equal(Input, result); + } + + [Fact] + public void Sanitize_WithEmptyString_ReturnsEmptyString() + { + // Arrange + const string Input = ""; + + // Act + string result = CosmosIdSanitizer.Sanitize(Input); + + // Assert + Assert.Equal(Input, result); + } + + [Theory] + [InlineData("/", "~0")] + [InlineData("\\", "~1")] + [InlineData("?", "~2")] + [InlineData("#", "~3")] + [InlineData("_", "~4")] + [InlineData("~", "~5")] + public void Sanitize_WithSingleSpecialChar_ReturnsCorrectEscapeSequence(string input, string expected) + { + // Act + string result = CosmosIdSanitizer.Sanitize(input); + + // Assert + Assert.Equal(expected, result); + } + + [Fact] + public void Sanitize_WithMultipleSpecialChars_ReturnsCorrectEscapeSequences() + { + // Arrange + const string Input = "test/path\\file?query#fragment_underscore~tilde"; + const string Expected = "test~0path~1file~2query~3fragment~4underscore~5tilde"; + + // Act + string result = CosmosIdSanitizer.Sanitize(Input); + + // Assert + Assert.Equal(Expected, result); + } + + [Fact] + public void Sanitize_WithMixedValidAndSpecialChars_ReturnsCorrectResult() + { + // Arrange + const string Input = "user/123"; + const string Expected = "user~0123"; + + // Act + string result = CosmosIdSanitizer.Sanitize(Input); + + // Assert + Assert.Equal(Expected, result); + } + + [Fact] + public void Unsanitize_WithValidInputNoEscapeChars_ReturnsOriginalString() + { + // Arrange + const string Input = "ValidId123"; + + // Act + string result = CosmosIdSanitizer.Unsanitize(Input); + + // Assert + Assert.Equal(Input, result); + } + + [Fact] + public void Unsanitize_WithEmptyString_ReturnsEmptyString() + { + // Arrange + const string Input = ""; + + // Act + string result = CosmosIdSanitizer.Unsanitize(Input); + + // Assert + Assert.Equal(Input, result); + } + + [Theory] + [InlineData("~0", "/")] + [InlineData("~1", "\\")] + [InlineData("~2", "?")] + [InlineData("~3", "#")] + [InlineData("~4", "_")] + [InlineData("~5", "~")] + public void Unsanitize_WithSingleEscapeSequence_ReturnsCorrectChar(string input, string expected) + { + // Act + string result = CosmosIdSanitizer.Unsanitize(input); + + // Assert + Assert.Equal(expected, result); + } + + [Fact] + public void Unsanitize_WithMultipleEscapeSequences_ReturnsCorrectResult() + { + // Arrange + const string Input = "test~0path~1file~2query~3fragment~4underscore~5tilde"; + const string Expected = "test/path\\file?query#fragment_underscore~tilde"; + + // Act + string result = CosmosIdSanitizer.Unsanitize(Input); + + // Assert + Assert.Equal(Expected, result); + } + + [Fact] + public void Unsanitize_WithMixedValidAndEscapeChars_ReturnsCorrectResult() + { + // Arrange + const string Input = "user~0123"; + const string Expected = "user/123"; + + // Act + string result = CosmosIdSanitizer.Unsanitize(Input); + + // Assert + Assert.Equal(Expected, result); + } + + [Theory] + [InlineData("~6")] + [InlineData("~A")] + [InlineData("~z")] + public void Unsanitize_WithInvalidEscapeSequence_ThrowsArgumentException(string input) + { + // Act & Assert + var exception = Assert.Throws(() => CosmosIdSanitizer.Unsanitize(input)); + Assert.Contains("Input is not in a valid format: Encountered unsupported escape sequence", exception.Message); + } + + [Fact] + public void SanitizeUnsanitize_RoundTrip_ReturnsOriginalString() + { + // Arrange + const string Original = "user/path\\to?file#with_underscore~and~tildes"; + + // Act + string sanitized = CosmosIdSanitizer.Sanitize(Original); + string unsanitized = CosmosIdSanitizer.Unsanitize(sanitized); + + // Assert + Assert.Equal(Original, unsanitized); + } + + [Theory] + [InlineData("")] + [InlineData("a")] + [InlineData("abc")] + [InlineData("simple-text")] + [InlineData("user123")] + [InlineData("user/path")] + [InlineData("/\\?#_~")] + [InlineData("complex/path\\with?query#fragment_underscore~tilde")] + public void SanitizeUnsanitize_RoundTripProperty_AlwaysReturnsOriginal(string original) + { + // Act + string sanitized = CosmosIdSanitizer.Sanitize(original); + string unsanitized = CosmosIdSanitizer.Unsanitize(sanitized); + + // Assert + Assert.Equal(original, unsanitized); + } + + [Fact] + public void Sanitize_WithLongString_HandlesCorrectly() + { + // Arrange + var input = new string('a', 1000) + "/" + new string('b', 1000) + "\\" + new string('c', 1000); + + // Act + string result = CosmosIdSanitizer.Sanitize(input); + + // Assert + Assert.Contains("~0", result); + Assert.Contains("~1", result); + Assert.Equal(3004, result.Length); // Original 3002 chars + 2 escape chars + } + + [Fact] + public void Unsanitize_WithLongString_HandlesCorrectly() + { + // Arrange + var input = new string('a', 1000) + "~0" + new string('b', 1000) + "~1" + new string('c', 1000); + + // Act + string result = CosmosIdSanitizer.Unsanitize(input); + + // Assert + Assert.Contains("/", result); + Assert.Contains("\\", result); + Assert.Equal(3002, result.Length); // 3004 chars - 2 escape chars + } + + [Fact] + public void SeparatorChar_HasCorrectValue() + { + // Assert + Assert.Equal('_', CosmosIdSanitizer.SeparatorChar); + } + + [Fact] + public void Sanitize_WithOnlySeparatorChar_EscapesCorrectly() + { + // Arrange + const string Input = "_"; + + // Act + string result = CosmosIdSanitizer.Sanitize(Input); + + // Assert + Assert.Equal("~4", result); + } + + [Fact] + public void Sanitize_WithConsecutiveSpecialChars_HandlesCorrectly() + { + // Arrange + const string Input = "//\\\\??##__~~"; + const string Expected = "~0~0~1~1~2~2~3~3~4~4~5~5"; + + // Act + string result = CosmosIdSanitizer.Sanitize(Input); + + // Assert + Assert.Equal(Expected, result); + } + + [Fact] + public void Unsanitize_WithConsecutiveEscapeSequences_HandlesCorrectly() + { + // Arrange + const string Input = "~0~0~1~1~2~2~3~3~4~4~5~5"; + const string Expected = "//\\\\??##__~~"; + + // Act + string result = CosmosIdSanitizer.Unsanitize(Input); + + // Assert + Assert.Equal(Expected, result); + } +} diff --git a/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests/CosmosTestFixture.cs b/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests/CosmosTestFixture.cs new file mode 100644 index 0000000000..09c0b5b6ea --- /dev/null +++ b/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests/CosmosTestFixture.cs @@ -0,0 +1,86 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System.Text.Json; +using Aspire.Hosting; +using Azure.Identity; +using Microsoft.Azure.Cosmos; +using Microsoft.Extensions.Logging; + +#pragma warning disable CA2007, VSTHRD111, CS1591 + +namespace Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests; + +[CollectionDefinition("Cosmos Test Collection")] +public class CosmosTests : ICollectionFixture { } + +/// +/// Shared test fixture for CosmosDB integration tests. +/// Sets up and manages the CosmosDB container for all tests. +/// +public class CosmosTestFixture : IAsyncLifetime +{ + public DistributedApplication App { get; private set; } = default!; + public CosmosClient CosmosClient { get; private set; } = default!; + public Container Container { get; private set; } = default!; + + /// + public async Task InitializeAsync() + { + using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(300)); + var cancellationToken = cts.Token; + + var appHost = await DistributedApplicationTestingBuilder + .CreateAsync(cancellationToken); + + appHost.Services.AddLogging(logging => + { + logging.SetMinimumLevel(LogLevel.Debug); + logging.AddFilter(appHost.Environment.ApplicationName, LogLevel.Debug); + logging.AddFilter("Aspire.", LogLevel.Debug); + }); + + appHost.Services.ConfigureHttpClientDefaults(clientBuilder => + { + clientBuilder.AddStandardResilienceHandler(); + }); + + this.App = await appHost.BuildAsync(cancellationToken).WaitAsync(cancellationToken); + await this.App.StartAsync(cancellationToken).WaitAsync(cancellationToken); + + var cs = await this.App.GetConnectionStringAsync(CosmosDBTestConstants.TestCosmosDbName, cancellationToken); + CosmosClientOptions ccoptions = new() + { + UseSystemTextJsonSerializerWithOptions = new JsonSerializerOptions() + { + PropertyNamingPolicy = JsonNamingPolicy.CamelCase, + } + }; + + if (CosmosDBTestConstants.UseEmulatorForTesting) + { + ccoptions.ConnectionMode = ConnectionMode.Gateway; + ccoptions.LimitToEndpoint = true; + this.CosmosClient = new CosmosClient(cs, ccoptions); + } + else + { + this.CosmosClient = new CosmosClient(cs, new DefaultAzureCredential(), ccoptions); + } + + var database = this.CosmosClient.GetDatabase(CosmosDBTestConstants.TestCosmosDbDatabaseName); + + var containerProperties = new ContainerProperties() + { + Id = "CosmosActorStateStorageTests", + PartitionKeyPath = "/actorId" + }; + + this.Container = await database.CreateContainerIfNotExistsAsync(containerProperties); + } + + public async Task DisposeAsync() + { + await this.App.DisposeAsync(); + this.CosmosClient.Dispose(); + } +} diff --git a/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests/LazyCosmosContainerTests.cs b/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests/LazyCosmosContainerTests.cs new file mode 100644 index 0000000000..07925d2a60 --- /dev/null +++ b/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests/LazyCosmosContainerTests.cs @@ -0,0 +1,285 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System.Text.Json; +using Microsoft.Azure.Cosmos; + +namespace Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests; + +/// +/// Integration tests for LazyCosmosContainer to verify lazy initialization behavior. +/// +[Collection("Cosmos Test Collection")] +public class LazyCosmosContainerTests +{ + private readonly CosmosTestFixture _fixture; + + public LazyCosmosContainerTests(CosmosTestFixture fixture) + { + this._fixture = fixture; + } + + private static readonly TimeSpan s_defaultTimeout = TimeSpan.FromSeconds(300); + + [Fact] + public async Task GetContainerAsync_WithExistingContainer_ShouldReturnImmediatelyAsync() + { + // Arrange + using var cts = new CancellationTokenSource(s_defaultTimeout); + var lazyContainer = new LazyCosmosContainer(this._fixture.Container); + + // Act + var result = await lazyContainer.GetContainerAsync(); + + // Assert + Assert.Same(this._fixture.Container, result); + } + + [Fact] + public async Task GetContainerAsync_WithExistingContainer_MultipleCalls_ShouldReturnSameInstanceAsync() + { + // Arrange + using var cts = new CancellationTokenSource(s_defaultTimeout); + var lazyContainer = new LazyCosmosContainer(this._fixture.Container); + + // Act + var result1 = await lazyContainer.GetContainerAsync(); + var result2 = await lazyContainer.GetContainerAsync(); + var result3 = await lazyContainer.GetContainerAsync(); + + // Assert + Assert.Same(result1, result2); + Assert.Same(result2, result3); + Assert.Same(this._fixture.Container, result1); + } + + [Fact] + public async Task GetContainerAsync_WithCosmosClient_ShouldInitializeAndWorkCorrectlyAsync() + { + // Arrange + using var cts = new CancellationTokenSource(s_defaultTimeout); + var cancellationToken = cts.Token; + + // Create a unique container name for this test + var testContainerName = $"LazyContainerTest_{Guid.NewGuid():N}"; + var lazyContainer = new LazyCosmosContainer(this._fixture.CosmosClient, CosmosDBTestConstants.TestCosmosDbDatabaseName, testContainerName); + + try + { + // Act + var container = await lazyContainer.GetContainerAsync(); + + // Assert - Container should be usable for actual operations + Assert.NotNull(container); + Assert.Equal(testContainerName, container.Id); + + // Verify the container can perform basic operations + var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString()); + var storage = new CosmosActorStateStorage(lazyContainer); + + var key = "testKey"; + var value = JsonSerializer.SerializeToElement("testValue"); + var operations = new List + { + new SetValueOperation(key, value) + }; + + // This should work if the container was properly initialized + var writeResult = await storage.WriteStateAsync(testActorId, operations, "0", cancellationToken); + Assert.True(writeResult.Success); + Assert.NotEqual("0", writeResult.ETag); + } + finally + { + // Cleanup - delete the test container + try + { + var container = await lazyContainer.GetContainerAsync(); + await container.DeleteContainerAsync(); + } + catch + { + // Ignore cleanup errors + } + } + } + + [Fact] + public async Task GetContainerAsync_WithCosmosClient_MultipleCalls_ShouldReturnSameInstanceAsync() + { + // Arrange + using var cts = new CancellationTokenSource(s_defaultTimeout); + + // Create a unique container name for this test + var testContainerName = $"LazyContainerTest_{Guid.NewGuid():N}"; + var lazyContainer = new LazyCosmosContainer(this._fixture.CosmosClient, CosmosDBTestConstants.TestCosmosDbDatabaseName, testContainerName); + + try + { + // Act + var result1 = await lazyContainer.GetContainerAsync(); + var result2 = await lazyContainer.GetContainerAsync(); + var result3 = await lazyContainer.GetContainerAsync(); + + // Assert + Assert.Same(result1, result2); + Assert.Same(result2, result3); + Assert.Equal(testContainerName, result1.Id); + } + finally + { + // Cleanup + try + { + var container = await lazyContainer.GetContainerAsync(); + await container.DeleteContainerAsync(); + } + catch + { + // Ignore cleanup errors + } + } + } + + [Fact] + public async Task GetContainerAsync_WithCosmosClient_ConcurrentAccess_ShouldInitializeOnlyOnceAsync() + { + // Arrange + using var cts = new CancellationTokenSource(s_defaultTimeout); + + // Create a unique container name for this test + var testContainerName = $"LazyContainerTest_{Guid.NewGuid():N}"; + var lazyContainer = new LazyCosmosContainer(this._fixture.CosmosClient, CosmosDBTestConstants.TestCosmosDbDatabaseName, testContainerName); + + try + { + // Act - Execute multiple concurrent calls + var tasks = new List>(); + for (int i = 0; i < 10; i++) + { + tasks.Add(lazyContainer.GetContainerAsync()); + } + var results = await Task.WhenAll(tasks); + + // Assert + // All results should be the same instance + for (int i = 1; i < results.Length; i++) + { + Assert.Same(results[0], results[i]); + } + Assert.Equal(testContainerName, results[0].Id); + } + finally + { + // Cleanup + try + { + var container = await lazyContainer.GetContainerAsync(); + await container.DeleteContainerAsync(); + } + catch + { + // Ignore cleanup errors + } + } + } + + [Fact] + public void Constructor_WithNullContainer_ShouldThrowArgumentNullException() + { + // Act & Assert + Assert.Throws(() => new LazyCosmosContainer((Container)null!)); + } + + [Fact] + public void Constructor_WithNullCosmosClient_ShouldThrowArgumentNullException() + { + // Act & Assert + Assert.Throws(() => new LazyCosmosContainer(null!, "test-db", "test-container")); + } + + [Fact] + public void Constructor_WithNullDatabaseName_ShouldThrowArgumentNullException() + { + // Act & Assert + Assert.Throws(() => new LazyCosmosContainer(this._fixture.CosmosClient, null!, "test-container")); + } + + [Fact] + public void Constructor_WithNullContainerName_ShouldThrowArgumentNullException() + { + // Act & Assert + Assert.Throws(() => new LazyCosmosContainer(this._fixture.CosmosClient, "test-db", null!)); + } + + [Fact] + public async Task LazyCosmosContainer_WithInternalConstructor_ShouldWorkWithCosmosActorStateStorageAsync() + { + // Arrange + using var cts = new CancellationTokenSource(s_defaultTimeout); + var cancellationToken = cts.Token; + + // Create a unique container name for this test + var testContainerName = $"LazyContainerTest_{Guid.NewGuid():N}"; + var lazyContainer = new LazyCosmosContainer(this._fixture.CosmosClient, CosmosDBTestConstants.TestCosmosDbDatabaseName, testContainerName); + + try + { + // Act - Create storage using the internal constructor (like DI would) + var storage = new CosmosActorStateStorage(lazyContainer); + var testActorId = new ActorId("TestActor", Guid.NewGuid().ToString()); + + var key = "testKey"; + var value = JsonSerializer.SerializeToElement("testValue"); + var operations = new List + { + new SetValueOperation(key, value) + }; + + // This should work - container should be initialized on first storage operation + var writeResult = await storage.WriteStateAsync(testActorId, operations, "0", cancellationToken); + + // Assert + Assert.True(writeResult.Success); + Assert.NotEqual("0", writeResult.ETag); + + // Verify we can read back the value + var readOperations = new List + { + new GetValueOperation(key) + }; + var readResult = await storage.ReadStateAsync(testActorId, readOperations, cancellationToken); + + Assert.Single(readResult.Results); + var getValue = readResult.Results[0] as GetValueResult; + Assert.NotNull(getValue); + Assert.Equal("testValue", getValue.Value?.GetString()); + } + finally + { + // Cleanup + try + { + var container = await lazyContainer.GetContainerAsync(); + await container.DeleteContainerAsync(); + } + catch + { + // Ignore cleanup errors + } + } + } + + [Fact] + public async Task GetContainerAsync_WithInvalidDatabaseName_ShouldThrowCosmosExceptionAsync() + { + // Arrange + using var cts = new CancellationTokenSource(s_defaultTimeout); + + // Use an invalid database name that should cause Cosmos to reject it + var invalidDatabaseName = new string('a', 256); // Database names have limits + var lazyContainer = new LazyCosmosContainer(this._fixture.CosmosClient, invalidDatabaseName, "test-container"); + + // Act & Assert + await Assert.ThrowsAsync(async () => await lazyContainer.GetContainerAsync()); + } +} diff --git a/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests.csproj b/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests.csproj new file mode 100644 index 0000000000..9f3271c68a --- /dev/null +++ b/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests.csproj @@ -0,0 +1,28 @@ + + + + net9.0 + enable + enable + false + true + + + + + + + + + + + + + + + + + + + + diff --git a/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests/SkipOnEmulatorFactAttribute.cs b/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests/SkipOnEmulatorFactAttribute.cs new file mode 100644 index 0000000000..b1508873f6 --- /dev/null +++ b/dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests/SkipOnEmulatorFactAttribute.cs @@ -0,0 +1,21 @@ +// Copyright (c) Microsoft. All rights reserved. + +namespace Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests; + +/// +/// Skip test if running on CosmosDB emulator. +/// +[AttributeUsage(AttributeTargets.Method | AttributeTargets.Class, AllowMultiple = false)] +public sealed class SkipOnEmulatorFactAttribute : FactAttribute +{ + /// + /// Initializes a new instance of the class. + /// + public SkipOnEmulatorFactAttribute() + { + if (CosmosDBTestConstants.UseEmulatorForTesting) + { + this.Skip = "Skipping test on CosmosDB emulator."; + } + } +}