.NET: [Breaking] Allow passing auth token credential to cosmosdb extensions (#3250)

* allow passing token credentials to cosmosdb extensions

* Update dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBWorkflowExtensions.cs

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBWorkflowExtensions.cs

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBChatExtensions.cs

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
This commit is contained in:
SergeyMenshykh
2026-01-22 11:30:58 +00:00
committed by GitHub
Unverified
parent 2cf4980d77
commit f47645cdc8
2 changed files with 32 additions and 9 deletions
@@ -3,7 +3,7 @@
using System;
using System.Diagnostics.CodeAnalysis;
using System.Threading.Tasks;
using Azure.Identity;
using Azure.Core;
using Microsoft.Azure.Cosmos;
namespace Microsoft.Agents.AI;
@@ -47,8 +47,9 @@ public static class CosmosDBChatExtensions
/// <param name="accountEndpoint">The Cosmos DB account endpoint URI.</param>
/// <param name="databaseId">The identifier of the Cosmos DB database.</param>
/// <param name="containerId">The identifier of the Cosmos DB container.</param>
/// <param name="tokenCredential">The TokenCredential to use for authentication (e.g., DefaultAzureCredential, ManagedIdentityCredential).</param>
/// <returns>The configured <see cref="ChatClientAgentOptions"/>.</returns>
/// <exception cref="ArgumentNullException">Thrown when <paramref name="options"/> is null.</exception>
/// <exception cref="ArgumentNullException">Thrown when <paramref name="options"/> or <paramref name="tokenCredential"/> is null.</exception>
/// <exception cref="ArgumentException">Thrown when any string parameter is null or whitespace.</exception>
[RequiresUnreferencedCode("The CosmosChatMessageStore uses JSON serialization which is incompatible with trimming.")]
[RequiresDynamicCode("The CosmosChatMessageStore uses JSON serialization which is incompatible with NativeAOT.")]
@@ -56,14 +57,20 @@ public static class CosmosDBChatExtensions
this ChatClientAgentOptions options,
string accountEndpoint,
string databaseId,
string containerId)
string containerId,
TokenCredential tokenCredential)
{
if (options is null)
{
throw new ArgumentNullException(nameof(options));
}
options.ChatMessageStoreFactory = (context, ct) => new ValueTask<ChatMessageStore>(new CosmosChatMessageStore(accountEndpoint, new DefaultAzureCredential(), databaseId, containerId));
if (tokenCredential is null)
{
throw new ArgumentNullException(nameof(tokenCredential));
}
options.ChatMessageStoreFactory = (context, ct) => new ValueTask<ChatMessageStore>(new CosmosChatMessageStore(accountEndpoint, tokenCredential, databaseId, containerId));
return options;
}
@@ -2,7 +2,7 @@
using System;
using System.Diagnostics.CodeAnalysis;
using Azure.Identity;
using Azure.Core;
using Microsoft.Agents.AI.Workflows.Checkpointing;
using Microsoft.Azure.Cosmos;
@@ -52,14 +52,17 @@ public static class CosmosDBWorkflowExtensions
/// <param name="accountEndpoint">The Cosmos DB account endpoint URI.</param>
/// <param name="databaseId">The identifier of the Cosmos DB database.</param>
/// <param name="containerId">The identifier of the Cosmos DB container.</param>
/// <param name="tokenCredential">The TokenCredential to use for authentication (e.g., DefaultAzureCredential, ManagedIdentityCredential).</param>
/// <returns>A new instance of <see cref="CosmosCheckpointStore"/>.</returns>
/// <exception cref="ArgumentException">Thrown when any string parameter is null or whitespace.</exception>
/// <exception cref="ArgumentNullException">Thrown when <paramref name="tokenCredential"/> is null.</exception>
[RequiresUnreferencedCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with trimming.")]
[RequiresDynamicCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with NativeAOT.")]
public static CosmosCheckpointStore CreateCheckpointStoreUsingManagedIdentity(
string accountEndpoint,
string databaseId,
string containerId)
string containerId,
TokenCredential tokenCredential)
{
if (string.IsNullOrWhiteSpace(accountEndpoint))
{
@@ -76,7 +79,12 @@ public static class CosmosDBWorkflowExtensions
throw new ArgumentException("Cannot be null or whitespace", nameof(containerId));
}
return new CosmosCheckpointStore(accountEndpoint, new DefaultAzureCredential(), databaseId, containerId);
if (tokenCredential is null)
{
throw new ArgumentNullException(nameof(tokenCredential));
}
return new CosmosCheckpointStore(accountEndpoint, tokenCredential, databaseId, containerId);
}
/// <summary>
@@ -154,14 +162,17 @@ public static class CosmosDBWorkflowExtensions
/// <param name="accountEndpoint">The Cosmos DB account endpoint URI.</param>
/// <param name="databaseId">The identifier of the Cosmos DB database.</param>
/// <param name="containerId">The identifier of the Cosmos DB container.</param>
/// <param name="tokenCredential">The TokenCredential to use for authentication (e.g., DefaultAzureCredential, ManagedIdentityCredential).</param>
/// <returns>A new instance of <see cref="CosmosCheckpointStore{T}"/>.</returns>
/// <exception cref="ArgumentException">Thrown when any string parameter is null or whitespace.</exception>
/// <exception cref="ArgumentNullException">Thrown when <paramref name="tokenCredential"/> is null.</exception>
[RequiresUnreferencedCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with trimming.")]
[RequiresDynamicCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with NativeAOT.")]
public static CosmosCheckpointStore<T> CreateCheckpointStoreUsingManagedIdentity<T>(
string accountEndpoint,
string databaseId,
string containerId)
string containerId,
TokenCredential tokenCredential)
{
if (string.IsNullOrWhiteSpace(accountEndpoint))
{
@@ -178,7 +189,12 @@ public static class CosmosDBWorkflowExtensions
throw new ArgumentException("Cannot be null or whitespace", nameof(containerId));
}
return new CosmosCheckpointStore<T>(accountEndpoint, new DefaultAzureCredential(), databaseId, containerId);
if (tokenCredential is null)
{
throw new ArgumentNullException(nameof(tokenCredential));
}
return new CosmosCheckpointStore<T>(accountEndpoint, tokenCredential, databaseId, containerId);
}
/// <summary>