From f47645cdc8001e7bf14052d60615b49accff53ee Mon Sep 17 00:00:00 2001
From: SergeyMenshykh <68852919+SergeyMenshykh@users.noreply.github.com>
Date: Thu, 22 Jan 2026 11:30:58 +0000
Subject: [PATCH] .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>
---
 .../CosmosDBChatExtensions.cs                 | 15 ++++++++---
 .../CosmosDBWorkflowExtensions.cs             | 26 +++++++++++++++----
 2 files changed, 32 insertions(+), 9 deletions(-)

diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBChatExtensions.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBChatExtensions.cs
index 45c0d09536..061b64593c 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,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;
     }
 
diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBWorkflowExtensions.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBWorkflowExtensions.cs
index 9d8bc52e68..4005808dbe 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,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>