.NET: CosmosDB Actor State Storage (#262)

* Implement CosmosDB actor state storage.

* Fix.

* Minor fixes.

* Fixes.

* Make CosmosDB initialization be lazy.

* Remove unnecessary read from write path.

* Throw on empty writes.

* Add arg validation for read.

* Add CosmosIdSanitizer.

* Fix.

* Fix.

* Simplify doc IDs.

* Update comment.

* fb

* Make LazyCosmosContainer internal and add tests.

* Make test constants public and remove IVT.

* Use source generated JSON context for future nativeAOT support.

* Re-add dropped comments.

dotnet/Directory.Packages.props

@@ -10,6 +10,9 @@
     <PackageVersion Include="Aspire.Azure.AI.OpenAI" Version="9.3.1-preview.1.25305.6" />
     <PackageVersion Include="Aspire.Hosting.AppHost" Version="9.3.1" />
     <PackageVersion Include="Aspire.Hosting.Azure.CognitiveServices" Version="9.3.1" />
+    <PackageVersion Include="Aspire.Hosting.Azure.CosmosDB" Version="9.3.1" />
+    <PackageVersion Include="Aspire.Microsoft.Azure.Cosmos" Version="9.3.1" />
+    <PackageVersion Include="Aspire.Hosting.Testing" Version="9.3.1" />
     <PackageVersion Include="Azure.AI.Agents.Persistent" Version="1.2.0-beta.1" />
     <PackageVersion Include="Azure.AI.OpenAI" Version="2.2.0-beta.5" />
     <PackageVersion Include="Azure.Identity" Version="1.14.2" />
@@ -21,6 +24,9 @@
     <PackageVersion Include="OpenTelemetry.Instrumentation.AspNetCore" Version="1.12.0" />
     <PackageVersion Include="OpenTelemetry.Instrumentation.Http" Version="1.12.0" />
     <PackageVersion Include="OpenTelemetry.Instrumentation.Runtime" Version="1.12.0" />
+    <PackageVersion Include="Microsoft.Azure.Cosmos" Version="3.52.0" />
+    <!-- Newtonsoft (Required by CosmosClient) -->
+    <PackageVersion Include="Newtonsoft.Json" Version="13.0.3" />
     <!-- System.* -->
     <PackageVersion Include="System.Linq.Async" Version="6.0.3" />
     <PackageVersion Include="System.Text.Json" Version="9.0.7" />

dotnet/agent-framework-dotnet.slnx

@@ -145,6 +145,9 @@
     <Project Path="src/Microsoft.Extensions.AI.Agents.CopilotStudio/Microsoft.Extensions.AI.Agents.CopilotStudio.csproj" />
     <Project Path="src/Microsoft.Extensions.AI.Agents.OpenAI/Microsoft.Extensions.AI.Agents.OpenAI.csproj" />
     <Project Path="src/Microsoft.Extensions.AI.Agents.Runtime.Abstractions/Microsoft.Extensions.AI.Agents.Runtime.Abstractions.csproj" />
+    <Project Path="src/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.csproj">
+      <BuildType Solution="Publish|*" Project="Release" />
+    </Project>
     <Project Path="src/Microsoft.Extensions.AI.Agents.Runtime/Microsoft.Extensions.AI.Agents.Runtime.csproj" Id="35d72ad5-61e1-45cc-a9ad-fa8490dbd146">
       <BuildType Solution="Publish|*" Project="Release" />
     </Project>
@@ -163,5 +166,14 @@
     <Project Path="tests/Microsoft.Extensions.AI.Agents.UnitTests/Microsoft.Extensions.AI.Agents.UnitTests.csproj">
       <BuildType Solution="Publish|*" Project="Debug" />
     </Project>
+      <Project Path="tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests.AppHost/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests.AppHost.csproj">
+        <BuildType Solution="Publish|*" Project="Release" />
+      </Project>
+      <Project Path="tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests.csproj">
+        <BuildType Solution="Publish|*" Project="Release" />
+      </Project>
+    <Folder Name="/CosmosDB/">
+
+    </Folder>
   </Folder>
 </Solution>

dotnet/samples/HelloHttpApi/HelloHttpApi.ApiService/HelloHttpApi.ApiService.csproj

@@ -12,12 +12,14 @@
     <ProjectReference Include="..\..\..\src\Microsoft.Extensions.AI.Agents.Runtime.Abstractions\Microsoft.Extensions.AI.Agents.Runtime.Abstractions.csproj" />
     <ProjectReference Include="..\..\..\src\Microsoft.Extensions.AI.Agents.Runtime\Microsoft.Extensions.AI.Agents.Runtime.csproj" />
     <ProjectReference Include="..\..\..\src\Microsoft.Extensions.AI.Agents\Microsoft.Extensions.AI.Agents.csproj" />
+    <ProjectReference Include="..\..\..\src\Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB\Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.csproj" />
     <ProjectReference Include="..\HelloHttpApi.ServiceDefaults\HelloHttpApi.ServiceDefaults.csproj" />
   </ItemGroup>
 
   <ItemGroup>
     <PackageReference Include="Aspire.Azure.AI.OpenAI" />
     <PackageReference Include="Aspire.Hosting.Azure.CognitiveServices" />
+    <PackageReference Include="Aspire.Microsoft.Azure.Cosmos" />
     <PackageReference Include="CommunityToolkit.Aspire.OllamaSharp" />
     <PackageReference Include="Microsoft.Extensions.AI" />
     <PackageReference Include="Microsoft.Extensions.AI.Abstractions" />

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(

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();
 

dotnet/samples/HelloHttpApi/HelloHttpApi.AppHost/HelloHttpApi.AppHost.csproj

@@ -14,6 +14,7 @@
   <ItemGroup>
     <PackageReference Include="Aspire.Hosting.AppHost" />
     <PackageReference Include="Aspire.Hosting.Azure.CognitiveServices" />
+    <PackageReference Include="Aspire.Hosting.Azure.CosmosDB" />
   </ItemGroup>
 
   <ItemGroup>

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<Projects.HelloHttpApi_ApiService>("apiservice")
-        .WithReference(chatModel);
+        .WithReference(chatModel)
+        .WithReference(cosmos).WaitFor(cosmos);
 
 builder.AddProject<Projects.HelloHttpApi_Web>("webfrontend")
     .WithExternalHttpEndpoints()

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;
+
+/// <summary>
+/// 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
+/// }
+/// </summary>
+public sealed class ActorRootDocument
+{
+    /// <summary>
+    /// The document ID.
+    /// </summary>
+    public string Id { get; set; } = default!;
+
+    /// <summary>
+    /// The actor ID.
+    /// </summary>
+    public string ActorId { get; set; } = default!;
+
+    /// <summary>
+    /// The last modified timestamp.
+    /// </summary>
+    public DateTimeOffset LastModified { get; set; }
+}
+
+/// <summary>
+/// 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
+/// }
+/// </summary>
+public sealed class ActorStateDocument
+{
+    /// <summary>
+    /// The document ID.
+    /// </summary>
+    public string Id { get; set; } = default!;
+
+    /// <summary>
+    /// The actor ID.
+    /// </summary>
+    public string ActorId { get; set; } = default!;
+
+    /// <summary>
+    /// The logical key for the state entry.
+    /// </summary>
+    public string Key { get; set; } = default!;
+
+    /// <summary>
+    /// The value payload.
+    /// </summary>
+    public JsonElement Value { get; set; } = default!;
+}
+
+/// <summary>
+/// Projection class for Cosmos DB queries to retrieve keys.
+/// </summary>
+public sealed class KeyProjection
+{
+    /// <summary>
+    /// The key value.
+    /// </summary>
+    public string Key { get; set; } = default!;
+}

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;
+
+/// <summary>
+/// Source-generated JSON type information for Cosmos DB actor state documents.
+/// </summary>
+[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;

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;
+
+/// <summary>
+/// Cosmos DB implementation of actor state storage.
+/// </summary>
+public class CosmosActorStateStorage : IActorStateStorage
+{
+    private readonly LazyCosmosContainer _lazyContainer;
+    private const string InitialEtag = "0"; // Initial ETag value when no state exists
+
+    /// <summary>
+    /// Constructs a new instance of <see cref="CosmosActorStateStorage"/> with the specified Cosmos DB container.
+    /// </summary>
+    /// <param name="container">The Cosmos DB container to use for storage.</param>
+    /// <exception cref="ArgumentNullException">Thrown when <paramref name="container"/> is null.</exception>
+    public CosmosActorStateStorage(Container container) => this._lazyContainer = new LazyCosmosContainer(container);
+
+    /// <summary>
+    /// This constructor is used by dependency injection to create an instance of <see cref="CosmosActorStateStorage"/>
+    /// with a lazy-loaded Cosmos container whose initialization is deferred until first access.
+    /// </summary>
+    /// <param name="lazyContainer">The lazy-loaded Cosmos container.</param>
+    /// <throws cref="ArgumentNullException">Thrown when <paramref name="lazyContainer"/> is null.</throws>
+    internal CosmosActorStateStorage(LazyCosmosContainer lazyContainer) =>
+        this._lazyContainer = lazyContainer ?? throw new ArgumentNullException(nameof(lazyContainer));
+
+    /// <summary>
+    /// Writes state changes to the actor's persistent storage.
+    /// </summary>
+    public async ValueTask<WriteResponse> WriteStateAsync(
+        ActorId actorId,
+        IReadOnlyCollection<ActorStateWriteOperation> 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);
+        }
+    }
+
+    /// <summary>
+    /// Reads state data from the actor's persistent storage.
+    /// </summary>
+    public async ValueTask<ReadResponse> ReadStateAsync(
+        ActorId actorId,
+        IReadOnlyCollection<ActorStateReadOperation> 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<ActorReadResult>();
+
+        // 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<ActorStateDocument>(
+                            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<KeyProjection>(
+                        query,
+                        list.ContinuationToken,
+                        requestOptions);
+
+                    var keys = new List<string>();
+                    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());
+
+    /// <summary>
+    /// Gets the current ETag for the actor's root document.
+    /// Returns a generated ETag if no root document exists.
+    /// </summary>
+    private async ValueTask<string> GetActorETagAsync(Container container, ActorId actorId, CancellationToken cancellationToken)
+    {
+        try
+        {
+            var rootResponse = await container.ReadItemAsync<ActorRootDocument>(
+                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");
+        }
+    }
+}

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<char> SanitizedCharacters => ['/', '\\', '?', '#', SeparatorChar, EscapeChar];
+    private static ReadOnlySpan<char> 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<char> input)
+    {
+        int count = 0;
+        foreach (var c in input)
+        {
+            if (SanitizedCharacters.IndexOf(c) >= 0)
+            {
+                count++;
+            }
+        }
+        return count;
+    }
+
+    private static int CountEscapeCharacters(ReadOnlySpan<char> input)
+    {
+        int count = 0;
+        foreach (var c in input)
+        {
+            if (c == EscapeChar)
+            {
+                count++;
+            }
+        }
+        return count;
+    }
+
+    private static void Encode(ReadOnlySpan<char> input, Span<char> 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<char> input, Span<char> 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;
+            }
+        }
+    }
+}

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<T>
+
+/// <summary>
+/// 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.
+/// </summary>
+internal sealed class LazyCosmosContainer
+{
+    private readonly CosmosClient? _cosmosClient;
+    private readonly string? _databaseName;
+    private readonly string? _containerName;
+    private readonly Lazy<Task<Container>> _lazyContainer;
+
+    /// <summary>
+    /// LazyCosmosContainer constructor that initializes the container lazily.
+    /// </summary>
+    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<Task<Container>>(this.InitializeContainerAsync, LazyThreadSafetyMode.ExecutionAndPublication);
+    }
+
+    /// <summary>
+    /// LazyCosmosContainer constructor that accepts an existing Container instance.
+    /// </summary>
+    public LazyCosmosContainer(Container container)
+    {
+        if (container is null)
+        {
+            throw new ArgumentNullException(nameof(container));
+        }
+
+        this._lazyContainer = new Lazy<Task<Container>>(() => Task.FromResult(container), LazyThreadSafetyMode.ExecutionAndPublication);
+    }
+
+    /// <summary>
+    /// Gets the Container, initializing it if necessary.
+    /// </summary>
+    public Task<Container> GetContainerAsync() => this._lazyContainer.Value;
+
+    private async Task<Container> 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<CompositePath>
+        {
+            new() { Path = "/actorId", Order = CompositePathSortOrder.Ascending },
+            new() { Path = "/key", Order = CompositePathSortOrder.Ascending }
+        });
+
+        var container = await database.Database.CreateContainerIfNotExistsAsync(containerProperties).ConfigureAwait(false);
+        return container.Container;
+    }
+}

dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.csproj

@@ -0,0 +1,23 @@
+<Project Sdk="Microsoft.NET.Sdk">
+
+  <PropertyGroup>
+    <TargetFrameworks>$(ProjectsTargetFrameworks)</TargetFrameworks>
+    <TargetFrameworks Condition="'$(Configuration)' == 'Debug'">$(ProjectsDebugTargetFrameworks)</TargetFrameworks>
+    <NoWarn>$(NoWarn);IDE1006;IDE0130</NoWarn>
+    <VersionSuffix>alpha</VersionSuffix>
+  </PropertyGroup>
+
+  <ItemGroup>
+    <PackageReference Include="Microsoft.Azure.Cosmos" />
+    <PackageReference Include="Newtonsoft.Json" />
+  </ItemGroup>
+
+  <ItemGroup>
+    <ProjectReference Include="../Microsoft.Extensions.AI.Agents.Runtime.Abstractions/Microsoft.Extensions.AI.Agents.Runtime.Abstractions.csproj" />
+  </ItemGroup>
+
+  <ItemGroup>
+    <InternalsVisibleTo Include="Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests" />
+  </ItemGroup>
+
+</Project>

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
+
+/// <summary>
+/// Extension methods for configuring Cosmos DB actor state storage in dependency injection.
+/// </summary>
+public static class ServiceCollectionExtensions
+{
+    /// <summary>
+    /// Adds Cosmos DB actor state storage to the service collection.
+    /// </summary>
+    /// <param name="services">The service collection to add services to.</param>
+    /// <param name="connectionString">The Cosmos DB connection string.</param>
+    /// <param name="databaseName">The database name to use for actor state storage.</param>
+    /// <param name="containerName">The container name to use for actor state storage. Defaults to "ActorState".</param>
+    /// <returns>The service collection for chaining.</returns>
+    public static IServiceCollection AddCosmosActorStateStorage(
+        this IServiceCollection services,
+        string connectionString,
+        string databaseName,
+        string containerName = "ActorState")
+    {
+        // Register CosmosClient as singleton
+        services.AddSingleton<CosmosClient>(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<LazyCosmosContainer>(serviceProvider =>
+        {
+            var cosmosClient = serviceProvider.GetRequiredService<CosmosClient>();
+            return new LazyCosmosContainer(cosmosClient, databaseName, containerName);
+        });
+
+        // Register the storage implementation
+        services.AddSingleton<IActorStateStorage>(serviceProvider =>
+        {
+            var lazyContainer = serviceProvider.GetRequiredService<LazyCosmosContainer>();
+            return new CosmosActorStateStorage(lazyContainer);
+        });
+
+        return services;
+    }
+
+    /// <summary>
+    /// Adds Cosmos DB actor state storage to the service collection using an existing CosmosClient from DI.
+    /// </summary>
+    /// <param name="services">The service collection to add services to.</param>
+    /// <param name="databaseName">The database name to use for actor state storage.</param>
+    /// <param name="containerName">The container name to use for actor state storage. Defaults to "ActorState".</param>
+    /// <returns>The service collection for chaining.</returns>
+    public static IServiceCollection AddCosmosActorStateStorage(
+        this IServiceCollection services,
+        string databaseName,
+        string containerName = "ActorState")
+    {
+        // Register LazyCosmosContainer as singleton using existing CosmosClient
+        services.AddSingleton<LazyCosmosContainer>(serviceProvider =>
+        {
+            var cosmosClient = serviceProvider.GetRequiredService<CosmosClient>();
+            return new LazyCosmosContainer(cosmosClient, databaseName, containerName);
+        });
+
+        // Register the storage implementation
+        services.AddSingleton<IActorStateStorage>(serviceProvider =>
+        {
+            var lazyContainer = serviceProvider.GetRequiredService<LazyCosmosContainer>();
+            return new CosmosActorStateStorage(lazyContainer);
+        });
+
+        return services;
+    }
+}

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<ActorStateWriteOperation>()],
-            operations.ETag,
-            cancellationToken).ConfigureAwait(false);
+        IReadOnlyCollection<ActorStateWriteOperation> writeOps =
+            [.. operations.Operations.OfType<ActorStateWriteOperation>()];
+
+        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);
 

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();

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);
+}

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 @@
+<Project Sdk="Microsoft.NET.Sdk">
+
+  <Sdk Name="Aspire.AppHost.Sdk" Version="9.3.1" />
+
+  <PropertyGroup>
+    <!--<OutputType>Exe</OutputType>-->
+    <TargetFramework>net9.0</TargetFramework>
+    <ImplicitUsings>enable</ImplicitUsings>
+    <Nullable>enable</Nullable>
+    <UserSecretsId>80048040-aaf1-4f44-9970-8aef39651edf</UserSecretsId>
+    <IsAspireHost>true</IsAspireHost>
+    <GenerateProgramFile>false</GenerateProgramFile>
+  </PropertyGroup>
+
+  <ItemGroup>
+    <PackageReference Include="Aspire.Hosting.AppHost" />
+    <PackageReference Include="Aspire.Hosting.Azure.CosmosDB" />
+  </ItemGroup>
+
+</Project>

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"
+      }
+    }
+  }
+}

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
+
+/// <summary>
+/// Integration tests for CosmosActorStateStorage focusing on concurrency control and ETag progression.
+/// </summary>
+[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<ActorStateWriteOperation> { new SetValueOperation(key, value1) };
+        var result1 = await storage.WriteStateAsync(testActorId, operations1, "0", cancellationToken);
+
+        // Act - Second write
+        var operations2 = new List<ActorStateWriteOperation> { new SetValueOperation(key, value2) };
+        var result2 = await storage.WriteStateAsync(testActorId, operations2, result1.ETag, cancellationToken);
+
+        // Act - Third write
+        var operations3 = new List<ActorStateWriteOperation> { 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<ActorStateWriteOperation>
+        {
+            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<Task<(bool Success, string? ETag, int AttemptNumber)>>();
+
+        // 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<ActorStateReadOperation>
+                        {
+                            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<ActorStateWriteOperation>
+                        {
+                            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<ActorStateReadOperation>
+        {
+            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<ActorStateWriteOperation>
+        {
+            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<ActorStateWriteOperation>
+        {
+            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<ActorStateReadOperation>
+        {
+            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());
+    }
+}

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;
+
+/// <summary>
+/// Integration tests for CosmosActorStateStorage focusing on ListKeys functionality.
+/// </summary>
+[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<ActorStateWriteOperation>
+        {
+            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<ActorStateReadOperation>
+        {
+            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<ActorStateWriteOperation>
+        {
+            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<ActorStateReadOperation>
+        {
+            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<ActorStateReadOperation>
+        {
+            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<ActorStateWriteOperation>
+        {
+            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<ActorStateReadOperation>
+        {
+            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<ActorStateWriteOperation>
+        {
+            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<ActorStateWriteOperation>
+        {
+            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<ActorStateReadOperation>
+        {
+            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<ActorStateWriteOperation>();
+        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<ActorStateReadOperation>
+        {
+            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<ActorStateReadOperation>
+        {
+            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<ActorStateReadOperation>
+        {
+            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<ActorStateReadOperation>
+        {
+            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
+    }
+}

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;
+
+/// <summary>
+/// Integration tests for CosmosActorStateStorage covering basic CRUD operations and advanced scenarios.
+/// </summary>
+[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<ActorStateWriteOperation>
+        {
+            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<ActorStateWriteOperation>
+        {
+            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<ActorStateReadOperation>
+        {
+            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<ActorStateReadOperation>
+        {
+            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<ActorStateWriteOperation>
+        {
+            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<ActorStateReadOperation>
+        {
+            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<ActorStateWriteOperation>
+        {
+            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<ActorStateWriteOperation>
+        {
+            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<ActorStateReadOperation>
+        {
+            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<ActorStateWriteOperation>
+        {
+            new SetValueOperation(key, value1)
+        };
+        var operations2 = new List<ActorStateWriteOperation>
+        {
+            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<ActorStateReadOperation>
+        {
+            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<ActorStateWriteOperation>();
+        // Act & Assert
+        await Assert.ThrowsAsync<InvalidOperationException>(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<ActorStateReadOperation>
+        {
+            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<string, object>
+            {
+                { "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<string, string>
+            {
+                { "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<ActorStateWriteOperation>
+        {
+            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<ActorStateReadOperation>
+        {
+            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<JsonElement>(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<ActorStateWriteOperation>
+        {
+            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<ActorStateReadOperation>
+        {
+            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&ampersand"
+        };
+
+        var writeOperations = new List<ActorStateWriteOperation>();
+        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<ActorStateReadOperation>
+            {
+                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<ActorStateReadOperation>
+        {
+            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);
+        }
+    }
+}

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<ArgumentException>(() => 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);
+    }
+}

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<CosmosTestFixture> { }
+
+/// <summary>
+/// Shared test fixture for CosmosDB integration tests.
+/// Sets up and manages the CosmosDB container for all tests.
+/// </summary>
+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!;
+
+    /// <inheritdoc/>
+    public async Task InitializeAsync()
+    {
+        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(300));
+        var cancellationToken = cts.Token;
+
+        var appHost = await DistributedApplicationTestingBuilder
+            .CreateAsync<Projects.Microsoft_Extensions_AI_Agents_Runtime_Storage_CosmosDB_Tests_AppHost>(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();
+    }
+}

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;
+
+/// <summary>
+/// Integration tests for LazyCosmosContainer to verify lazy initialization behavior.
+/// </summary>
+[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<ActorStateWriteOperation>
+            {
+                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<Task<Container>>();
+            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<ArgumentNullException>(() => new LazyCosmosContainer((Container)null!));
+    }
+
+    [Fact]
+    public void Constructor_WithNullCosmosClient_ShouldThrowArgumentNullException()
+    {
+        // Act & Assert
+        Assert.Throws<ArgumentNullException>(() => new LazyCosmosContainer(null!, "test-db", "test-container"));
+    }
+
+    [Fact]
+    public void Constructor_WithNullDatabaseName_ShouldThrowArgumentNullException()
+    {
+        // Act & Assert
+        Assert.Throws<ArgumentNullException>(() => new LazyCosmosContainer(this._fixture.CosmosClient, null!, "test-container"));
+    }
+
+    [Fact]
+    public void Constructor_WithNullContainerName_ShouldThrowArgumentNullException()
+    {
+        // Act & Assert
+        Assert.Throws<ArgumentNullException>(() => 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<ActorStateWriteOperation>
+            {
+                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<ActorStateReadOperation>
+            {
+                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<CosmosException>(async () => await lazyContainer.GetContainerAsync());
+    }
+}

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 @@
+<Project Sdk="Microsoft.NET.Sdk">
+
+  <PropertyGroup>
+    <TargetFramework>net9.0</TargetFramework>
+    <ImplicitUsings>enable</ImplicitUsings>
+    <Nullable>enable</Nullable>
+    <IsPackable>false</IsPackable>
+    <IsTestProject>true</IsTestProject>
+  </PropertyGroup>
+
+  <ItemGroup>
+    <PackageReference Include="Aspire.Hosting.Testing" />
+  </ItemGroup>
+
+  <ItemGroup>
+    <ProjectReference Include="..\Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests.AppHost\Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests.AppHost.csproj" />
+    <ProjectReference Include="..\..\..\src\Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB\Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.csproj" />
+  </ItemGroup>
+
+  <ItemGroup>
+    <Using Include="System.Net" />
+    <Using Include="Microsoft.Extensions.DependencyInjection" />
+    <Using Include="Aspire.Hosting.ApplicationModel" />
+    <Using Include="Aspire.Hosting.Testing" />
+    <Using Include="Xunit" />
+  </ItemGroup>
+
+</Project>

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;
+
+/// <summary>
+/// Skip test if running on CosmosDB emulator.
+/// </summary>
+[AttributeUsage(AttributeTargets.Method | AttributeTargets.Class, AllowMultiple = false)]
+public sealed class SkipOnEmulatorFactAttribute : FactAttribute
+{
+    /// <summary>
+    /// Initializes a new instance of the <see cref="SkipOnEmulatorFactAttribute"/> class.
+    /// </summary>
+    public SkipOnEmulatorFactAttribute()
+    {
+        if (CosmosDBTestConstants.UseEmulatorForTesting)
+        {
+            this.Skip = "Skipping test on CosmosDB emulator.";
+        }
+    }
+}