agent/agent-framework
2
0
Fork 0
mirror of https://github.com/microsoft/agent-framework.git synced 2026-06-16 21:04:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

.NET: chore: support retries on Cosmos storage creation (#402)

Browse Source 
* support retries

* tests + registration options

* fix ordering ..

* HK + update packages

* fix paths

* Update dotnet/tests/CosmosDB.IntegrationTests/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.Tests/CosmosTestFixture.cs

* re create project and fix some pk usage

* fix all tests

* try workflow?

* wip 1

* fix definition

* try with cosmos_use_emulator env?

* try ignore SSL errors?

* other cert verifications

* hardcode to 8081?

* proper valuation of ENV

* logging

* ensure db exsists for CI

* bump

* cleanup

* fix usage

* nit comment

* try only release for stability?

* try skip some flaky tests

* merge fixes + rollback container

* reimplement with iasyncdisposable pattern

* remove example doc struct
This commit is contained in:
Korolev Dmitry
2025-08-21 18:37:12 +00:00
committed by GitHub
parent 1116916737
commit 25291de8cb
22 changed files with 436 additions and 110 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB/ActorDocuments.cs
+12 -11
View File
@@ -11,13 +11,6 @@ namespace Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB;
/// 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
{
@@ -27,9 +20,13 @@ public sealed class ActorRootDocument
public string Id { get; set; } = default!;
/// <summary>
/// The actor ID.
/// The actor type.
/// </summary>
public string ActorId { get; set; } = default!;
public string ActorType { get; set; } = default!;
/// <summary>
/// The actor key.
/// </summary>
public string ActorKey { get; set; } = default!;
/// <summary>
/// The last modified timestamp.
@@ -55,9 +52,13 @@ public sealed class ActorStateDocument
public string Id { get; set; } = default!;
/// <summary>
/// The actor ID.
/// The actor type.
/// </summary>
public string ActorId { get; set; } = default!;
public string ActorType { get; set; } = default!;
/// <summary>
/// The actor key.
/// </summary>
public string ActorKey { get; set; } = default!;
/// <summary>
/// The logical key for the state entry.
dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB/CosmosActorStateStorage.cs
+37 -10
View File
@@ -11,7 +11,7 @@ namespace Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB;
/// <summary>
/// Cosmos DB implementation of actor state storage.
/// </summary>
public class CosmosActorStateStorage : IActorStateStorage
public class CosmosActorStateStorage : IActorStateStorage, IAsyncDisposable
{
private readonly LazyCosmosContainer _lazyContainer;
private const string InitialEtag = "0"; // Initial ETag value when no state exists
@@ -47,8 +47,8 @@ public class CosmosActorStateStorage : IActorStateStorage
}
var container = await this._lazyContainer.GetContainerAsync().ConfigureAwait(false);
var batch = container.CreateTransactionalBatch(GetPartitionKey(actorId));
var actorIdStr = actorId.ToString();
var (partitionKey, actorType, actorKey) = BuildPartitionKey(actorId);
var batch = container.CreateTransactionalBatch(partitionKey);
// Add data operations to batch
foreach (var op in operations)
@@ -61,7 +61,8 @@ public class CosmosActorStateStorage : IActorStateStorage
var item = new ActorStateDocument
{
Id = docId,
ActorId = actorIdStr,
ActorType = actorType,
ActorKey = actorKey,
Key = set.Key,
Value = set.Value
};
@@ -83,7 +84,8 @@ public class CosmosActorStateStorage : IActorStateStorage
var newRoot = new ActorRootDocument
{
Id = RootDocumentId,
ActorId = actorId.ToString(),
ActorType = actorType,
ActorKey = actorKey,
LastModified = DateTimeOffset.UtcNow,
};
@@ -103,6 +105,7 @@ public class CosmosActorStateStorage : IActorStateStorage
var result = await batch.ExecuteAsync(cancellationToken).ConfigureAwait(false);
if (!result.IsSuccessStatusCode)
{
_ = result.ErrorMessage;
return new WriteResponse(eTag: string.Empty, success: false);
}
@@ -135,6 +138,8 @@ public class CosmosActorStateStorage : IActorStateStorage
// Read root document first to get actor-level ETag
string actorETag = await this.GetActorETagAsync(container, actorId, cancellationToken).ConfigureAwait(false);
var actorType = actorId.Type.ToString();
var actorKey = actorId.Key;
foreach (var op in operations)
{
@@ -162,14 +167,16 @@ public class CosmosActorStateStorage : IActorStateStorage
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())
query = new QueryDefinition("SELECT c.key FROM c WHERE c.actorType = @actorType AND c.actorKey = @actorKey AND c.key != null AND STARTSWITH(c.key, @keyPrefix)")
.WithParameter("@actorType", actorType)
.WithParameter("@actorKey", actorKey)
.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());
query = new QueryDefinition("SELECT c.key FROM c WHERE c.actorType = @actorType AND c.actorKey = @actorKey AND c.key != null")
.WithParameter("@actorType", actorType)
.WithParameter("@actorKey", actorKey);
}
var requestOptions = new QueryRequestOptions
@@ -212,7 +219,18 @@ public class CosmosActorStateStorage : IActorStateStorage
private const string RootDocumentId = "rootdoc";
private static PartitionKey GetPartitionKey(ActorId actorId)
=> new(actorId.ToString());
{
var (partitionKey, _, _) = BuildPartitionKey(actorId);
return partitionKey;
}
private static (PartitionKey partitionKey, string actorType, string actorKey) BuildPartitionKey(ActorId actorId)
{
var actorType = actorId.Type.ToString();
var actorKey = actorId.Key;
var partitionKey = new PartitionKeyBuilder().Add(actorType).Add(actorKey).Build();
return (partitionKey, actorType, actorKey);
}
/// <summary>
/// Gets the current ETag for the actor's root document.
@@ -234,4 +252,13 @@ public class CosmosActorStateStorage : IActorStateStorage
return InitialEtag;
}
}
/// <summary>
/// Disposes the Cosmos DB container asynchronously.
/// </summary>
public async ValueTask DisposeAsync()
{
await this._lazyContainer.DisposeAsync().ConfigureAwait(false);
GC.SuppressFinalize(this);
}
}
dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB/LazyCosmosContainer.cs
+90 -16
View File
@@ -1,26 +1,30 @@
// Copyright (c) Microsoft. All rights reserved.
using System;
using System.Collections.ObjectModel;
using System.Net;
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
internal sealed class LazyCosmosContainer : IAsyncDisposable
{
private readonly static Random s_random = new();
private readonly CosmosClient? _cosmosClient;
private readonly string? _databaseName;
private readonly string? _containerName;
private readonly Lazy<Task<Container>> _lazyContainer;
private readonly CancellationTokenSource _cts = new();
private Task<Container>? _initTask;
// internal for testing
internal readonly static string[] CosmosPartitionKeyPaths = ["/actorType", "/actorKey"];
/// <summary>
/// LazyCosmosContainer constructor that initializes the container lazily.
@@ -30,9 +34,7 @@ internal sealed class LazyCosmosContainer
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>
@@ -43,20 +45,59 @@ internal sealed class LazyCosmosContainer
throw new ArgumentNullException(nameof(container));
}
this._lazyContainer = new Lazy<Task<Container>>(() => Task.FromResult(container), LazyThreadSafetyMode.ExecutionAndPublication);
this._initTask = Task.FromResult(container);
}
/// <summary>
/// Gets the Container, initializing it if necessary.
/// </summary>
public Task<Container> GetContainerAsync() => this._lazyContainer.Value;
public Task<Container> GetContainerAsync()
=> this._initTask ??= this.InitializeWithRetryAsync(this._cts.Token);
private async Task<Container> InitializeContainerAsync()
private async Task<Container> InitializeWithRetryAsync(CancellationToken cancellationToken)
{
var baseDelay = TimeSpan.FromSeconds(1);
var maxDelay = TimeSpan.FromSeconds(30);
var previousDelay = baseDelay;
while (true)
{
cancellationToken.ThrowIfCancellationRequested();
try
{
return await this.InitializeContainerAsync(cancellationToken).ConfigureAwait(false);
}
catch (CosmosException ex) when (IsTransient(ex))
{
// If server provided RetryAfter, respect it but add a small jitter so clients don't retry in perfect sync.
if (ex.RetryAfter is not null && ex.RetryAfter > TimeSpan.Zero)
{
var retry = ex.RetryAfter.Value;
var jitterMs = this.RandomNextDouble() * retry.TotalMilliseconds; // 0..retry
var delay = retry + TimeSpan.FromMilliseconds(jitterMs);
await Task.Delay(delay, cancellationToken).ConfigureAwait(false);
previousDelay = delay;
continue;
}
// sleep = min(maxDelay, random(baseDelay, previousDelay * 3))
var minMs = baseDelay.TotalMilliseconds;
var maxMs = Math.Min(maxDelay.TotalMilliseconds, Math.Max(minMs, previousDelay.TotalMilliseconds * 3));
var sleepMs = this.RandomNextDouble() * (maxMs - minMs) + minMs;
var jitterDelay = TimeSpan.FromMilliseconds(sleepMs);
await Task.Delay(jitterDelay, cancellationToken).ConfigureAwait(false);
previousDelay = jitterDelay;
}
}
}
private async Task<Container> InitializeContainerAsync(CancellationToken cancellationToken)
{
// Create database if it doesn't exist
var database = await this._cosmosClient!.CreateDatabaseIfNotExistsAsync(this._databaseName!).ConfigureAwait(false);
var database = await this._cosmosClient!.CreateDatabaseIfNotExistsAsync(this._databaseName!, cancellationToken: cancellationToken).ConfigureAwait(false);
var containerProperties = new ContainerProperties(this._containerName!, "/actorId")
var containerProperties = new ContainerProperties(this._containerName!, CosmosPartitionKeyPaths)
{
Id = this._containerName!,
IndexingPolicy = new IndexingPolicy
@@ -64,17 +105,50 @@ internal sealed class LazyCosmosContainer
IndexingMode = IndexingMode.Consistent,
Automatic = true
},
PartitionKeyPaths = ["/actorId"]
PartitionKeyPaths = CosmosPartitionKeyPaths
};
// Add composite index for efficient queries
containerProperties.IndexingPolicy.CompositeIndexes.Add(new Collection<CompositePath>
{
new() { Path = "/actorId", Order = CompositePathSortOrder.Ascending },
new() { Path = "/actorType", Order = CompositePathSortOrder.Ascending },
new() { Path = "/actorKey", Order = CompositePathSortOrder.Ascending },
new() { Path = "/key", Order = CompositePathSortOrder.Ascending }
});
var container = await database.Database.CreateContainerIfNotExistsAsync(containerProperties).ConfigureAwait(false);
var container = await database.Database.CreateContainerIfNotExistsAsync(containerProperties, cancellationToken: cancellationToken).ConfigureAwait(false);
return container.Container;
}
private static bool IsTransient(Exception exception)
{
return exception switch
{
CosmosException cosmosEx => cosmosEx.StatusCode switch
{
#if NET9_0_OR_GREATER
HttpStatusCode.TooManyRequests => true, // 429 - Rate limited
#endif
HttpStatusCode.InternalServerError => true, // 500 - Server error
HttpStatusCode.BadGateway => true, // 502 - Bad gateway
HttpStatusCode.ServiceUnavailable => true, // 503 - Service unavailable
HttpStatusCode.GatewayTimeout => true, // 504 - Gateway timeout
HttpStatusCode.RequestTimeout => true, // 408 - Request timeout
_ => false
},
TaskCanceledException or OperationCanceledException or ArgumentException => false,
_ => true // Retry other exceptions (network issues, etc.)
};
}
#pragma warning disable CA5394 // Do not use insecure randomness
private double RandomNextDouble() => s_random.NextDouble();
#pragma warning restore CA5394 // Do not use insecure randomness
public ValueTask DisposeAsync()
{
this._cts?.Cancel();
this._cts?.Dispose();
return default;
}
}
dotnet/src/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB/Microsoft.Extensions.AI.Agents.Runtime.Storage.CosmosDB.csproj
+1
View File
@@ -9,6 +9,7 @@
<ItemGroup>
<PackageReference Include="Microsoft.Azure.Cosmos" />
<PackageReference Include="Microsoft.Extensions.Options" />
<PackageReference Include="Newtonsoft.Json" />
</ItemGroup>