From 880a7722b971cc3573beefad7c4755f3d3c1bbe1 Mon Sep 17 00:00:00 2001 From: SergeyMenshykh Date: Fri, 16 Jan 2026 13:33:00 +0000 Subject: [PATCH 1/4] allow passing token credentials to cosmosdb extensions --- .../CosmosDBChatExtensions.cs | 13 +++++++--- .../CosmosDBWorkflowExtensions.cs | 24 +++++++++++++++---- 2 files changed, 29 insertions(+), 8 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBChatExtensions.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBChatExtensions.cs index 45c0d09536..b8c921cddf 100644 --- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBChatExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBChatExtensions.cs @@ -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,6 +47,7 @@ public static ChatClientAgentOptions WithCosmosDBMessageStore( /// The Cosmos DB account endpoint URI. /// The identifier of the Cosmos DB database. /// The identifier of the Cosmos DB container. + /// The TokenCredential to use for authentication (e.g., DefaultAzureCredential, ManagedIdentityCredential). /// The configured . /// Thrown when is null. /// Thrown when any string parameter is null or whitespace. @@ -56,14 +57,20 @@ public static ChatClientAgentOptions WithCosmosDBMessageStoreUsingManagedIdentit 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(new CosmosChatMessageStore(accountEndpoint, new DefaultAzureCredential(), databaseId, containerId)); + if (tokenCredential is null) + { + throw new ArgumentNullException(nameof(tokenCredential)); + } + + options.ChatMessageStoreFactory = (context, ct) => new ValueTask(new CosmosChatMessageStore(accountEndpoint, tokenCredential, databaseId, containerId)); return options; } diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBWorkflowExtensions.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBWorkflowExtensions.cs index 9d8bc52e68..e00bffd080 100644 --- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBWorkflowExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBWorkflowExtensions.cs @@ -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,6 +52,7 @@ public static CosmosCheckpointStore CreateCheckpointStore( /// The Cosmos DB account endpoint URI. /// The identifier of the Cosmos DB database. /// The identifier of the Cosmos DB container. + /// The TokenCredential to use for authentication (e.g., DefaultAzureCredential, ManagedIdentityCredential). /// A new instance of . /// Thrown when any string parameter is null or whitespace. [RequiresUnreferencedCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with trimming.")] @@ -59,7 +60,8 @@ public static CosmosCheckpointStore CreateCheckpointStore( public static CosmosCheckpointStore CreateCheckpointStoreUsingManagedIdentity( string accountEndpoint, string databaseId, - string containerId) + string containerId, + TokenCredential tokenCredential) { if (string.IsNullOrWhiteSpace(accountEndpoint)) { @@ -76,7 +78,12 @@ public static CosmosCheckpointStore CreateCheckpointStoreUsingManagedIdentity( 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); } /// @@ -154,6 +161,7 @@ public static CosmosCheckpointStore CreateCheckpointStore( /// The Cosmos DB account endpoint URI. /// The identifier of the Cosmos DB database. /// The identifier of the Cosmos DB container. + /// The TokenCredential to use for authentication (e.g., DefaultAzureCredential, ManagedIdentityCredential). /// A new instance of . /// Thrown when any string parameter is null or whitespace. [RequiresUnreferencedCode("The CosmosCheckpointStore uses JSON serialization which is incompatible with trimming.")] @@ -161,7 +169,8 @@ public static CosmosCheckpointStore CreateCheckpointStore( public static CosmosCheckpointStore CreateCheckpointStoreUsingManagedIdentity( string accountEndpoint, string databaseId, - string containerId) + string containerId, + TokenCredential tokenCredential) { if (string.IsNullOrWhiteSpace(accountEndpoint)) { @@ -178,7 +187,12 @@ public static CosmosCheckpointStore CreateCheckpointStoreUsingManagedIdentity 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); } /// From daf32d707705ee4aa7ffdcdfbbaed95898465b0e Mon Sep 17 00:00:00 2001 From: SergeyMenshykh <68852919+SergeyMenshykh@users.noreply.github.com> Date: Fri, 16 Jan 2026 16:28:32 +0000 Subject: [PATCH 2/4] Update dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBWorkflowExtensions.cs Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- .../CosmosDBWorkflowExtensions.cs | 1 + 1 file changed, 1 insertion(+) diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBWorkflowExtensions.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBWorkflowExtensions.cs index e00bffd080..8cc0a33547 100644 --- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBWorkflowExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBWorkflowExtensions.cs @@ -55,6 +55,7 @@ public static CosmosCheckpointStore CreateCheckpointStore( /// The TokenCredential to use for authentication (e.g., DefaultAzureCredential, ManagedIdentityCredential). /// A new instance of . /// Thrown when any string parameter is null or whitespace. + /// Thrown when is null. [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( From 6dec2d49766d48b79fba440ea5099c53d1377392 Mon Sep 17 00:00:00 2001 From: SergeyMenshykh <68852919+SergeyMenshykh@users.noreply.github.com> Date: Fri, 16 Jan 2026 16:28:40 +0000 Subject: [PATCH 3/4] Update dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBWorkflowExtensions.cs Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- .../CosmosDBWorkflowExtensions.cs | 1 + 1 file changed, 1 insertion(+) diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBWorkflowExtensions.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBWorkflowExtensions.cs index 8cc0a33547..4005808dbe 100644 --- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBWorkflowExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBWorkflowExtensions.cs @@ -165,6 +165,7 @@ public static CosmosCheckpointStore CreateCheckpointStore( /// The TokenCredential to use for authentication (e.g., DefaultAzureCredential, ManagedIdentityCredential). /// A new instance of . /// Thrown when any string parameter is null or whitespace. + /// Thrown when is null. [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( From 402deeab848c103aca695f467c5e63f5f1cc2dcd Mon Sep 17 00:00:00 2001 From: SergeyMenshykh <68852919+SergeyMenshykh@users.noreply.github.com> Date: Fri, 16 Jan 2026 16:28:50 +0000 Subject: [PATCH 4/4] Update dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBChatExtensions.cs Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- .../Microsoft.Agents.AI.CosmosNoSql/CosmosDBChatExtensions.cs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBChatExtensions.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBChatExtensions.cs index b8c921cddf..061b64593c 100644 --- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBChatExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBChatExtensions.cs @@ -49,7 +49,7 @@ public static ChatClientAgentOptions WithCosmosDBMessageStore( /// The identifier of the Cosmos DB container. /// The TokenCredential to use for authentication (e.g., DefaultAzureCredential, ManagedIdentityCredential). /// The configured . - /// Thrown when is null. + /// Thrown when or is null. /// Thrown when any string parameter is null or whitespace. [RequiresUnreferencedCode("The CosmosChatMessageStore uses JSON serialization which is incompatible with trimming.")] [RequiresDynamicCode("The CosmosChatMessageStore uses JSON serialization which is incompatible with NativeAOT.")]