// Copyright (c) Microsoft. All rights reserved.
using System;
using System.Collections.Generic;
using System.Diagnostics.CodeAnalysis;
using System.Linq;
using System.Text.Json;
using System.Threading;
using System.Threading.Tasks;
using Azure.Core;
using Microsoft.Azure.Cosmos;
using Microsoft.Extensions.AI;
using Microsoft.Shared.Diagnostics;
namespace Microsoft.Agents.AI;
///
/// Provides a Cosmos DB implementation of the abstract class.
///
[RequiresUnreferencedCode("The CosmosChatMessageStore uses JSON serialization which is incompatible with trimming.")]
[RequiresDynamicCode("The CosmosChatMessageStore uses JSON serialization which is incompatible with NativeAOT.")]
public sealed class CosmosChatMessageStore : ChatMessageStore, IDisposable
{
private readonly CosmosClient _cosmosClient;
private readonly Container _container;
private readonly bool _ownsClient;
private bool _disposed;
// Hierarchical partition key support
private readonly string? _tenantId;
private readonly string? _userId;
private readonly PartitionKey _partitionKey;
private readonly bool _useHierarchicalPartitioning;
///
/// Cached JSON serializer options for .NET 9.0 compatibility.
///
private static readonly JsonSerializerOptions s_defaultJsonOptions = CreateDefaultJsonOptions();
private static JsonSerializerOptions CreateDefaultJsonOptions()
{
var options = new JsonSerializerOptions();
#if NET9_0_OR_GREATER
// Configure TypeInfoResolver for .NET 9.0 to enable JSON serialization
options.TypeInfoResolver = new System.Text.Json.Serialization.Metadata.DefaultJsonTypeInfoResolver();
#endif
return options;
}
///
/// Gets or sets the maximum number of messages to return in a single query batch.
/// Default is 100 for optimal performance.
///
public int MaxItemCount { get; set; } = 100;
///
/// Gets or sets the maximum number of items per transactional batch operation.
/// Default is 100, maximum allowed by Cosmos DB is 100.
///
public int MaxBatchSize { get; set; } = 100;
///
/// Gets or sets the maximum number of messages to retrieve from the store.
/// This helps prevent exceeding LLM context windows in long conversations.
/// Default is null (no limit). When set, only the most recent messages are returned.
///
public int? MaxMessagesToRetrieve { get; set; }
///
/// Gets or sets the Time-To-Live (TTL) in seconds for messages.
/// Default is 86400 seconds (24 hours). Set to null to disable TTL.
///
public int? MessageTtlSeconds { get; set; } = 86400;
///
/// Gets the conversation ID associated with this message store.
///
public string ConversationId { get; init; }
///
/// Gets the database ID associated with this message store.
///
public string DatabaseId { get; init; }
///
/// Gets the container ID associated with this message store.
///
public string ContainerId { get; init; }
///
/// Internal primary constructor used by all public constructors.
///
/// The instance to use for Cosmos DB operations.
/// The identifier of the Cosmos DB database.
/// The identifier of the Cosmos DB container.
/// The unique identifier for this conversation thread.
/// Whether this instance owns the CosmosClient and should dispose it.
/// Optional tenant identifier for hierarchical partitioning.
/// Optional user identifier for hierarchical partitioning.
internal CosmosChatMessageStore(CosmosClient cosmosClient, string databaseId, string containerId, string conversationId, bool ownsClient, string? tenantId = null, string? userId = null)
{
this._cosmosClient = Throw.IfNull(cosmosClient);
this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId));
this.ConversationId = Throw.IfNullOrWhitespace(conversationId);
this.DatabaseId = databaseId;
this.ContainerId = containerId;
this._ownsClient = ownsClient;
// Initialize partitioning mode
this._tenantId = tenantId;
this._userId = userId;
this._useHierarchicalPartitioning = tenantId != null && userId != null;
this._partitionKey = this._useHierarchicalPartitioning
? new PartitionKeyBuilder()
.Add(tenantId!)
.Add(userId!)
.Add(conversationId)
.Build()
: new PartitionKey(conversationId);
}
///
/// Initializes a new instance of the class using a connection string.
///
/// The Cosmos DB connection string.
/// The identifier of the Cosmos DB database.
/// The identifier of the Cosmos DB container.
/// Thrown when any required parameter is null.
/// Thrown when any string parameter is null or whitespace.
public CosmosChatMessageStore(string connectionString, string databaseId, string containerId)
: this(connectionString, databaseId, containerId, Guid.NewGuid().ToString("N"))
{
}
///
/// Initializes a new instance of the class using a connection string.
///
/// The Cosmos DB connection string.
/// The identifier of the Cosmos DB database.
/// The identifier of the Cosmos DB container.
/// The unique identifier for this conversation thread.
/// Thrown when any required parameter is null.
/// Thrown when any string parameter is null or whitespace.
public CosmosChatMessageStore(string connectionString, string databaseId, string containerId, string conversationId)
: this(new CosmosClient(Throw.IfNullOrWhitespace(connectionString)), databaseId, containerId, conversationId, ownsClient: true)
{
}
///
/// Initializes a new instance of the class using TokenCredential for authentication.
///
/// The Cosmos DB account endpoint URI.
/// The TokenCredential to use for authentication (e.g., DefaultAzureCredential, ManagedIdentityCredential).
/// The identifier of the Cosmos DB database.
/// The identifier of the Cosmos DB container.
/// Thrown when any required parameter is null.
/// Thrown when any string parameter is null or whitespace.
public CosmosChatMessageStore(string accountEndpoint, TokenCredential tokenCredential, string databaseId, string containerId)
: this(accountEndpoint, tokenCredential, databaseId, containerId, Guid.NewGuid().ToString("N"))
{
}
///
/// Initializes a new instance of the class using a TokenCredential for authentication.
///
/// The Cosmos DB account endpoint URI.
/// The TokenCredential to use for authentication (e.g., DefaultAzureCredential, ManagedIdentityCredential).
/// The identifier of the Cosmos DB database.
/// The identifier of the Cosmos DB container.
/// The unique identifier for this conversation thread.
/// Thrown when any required parameter is null.
/// Thrown when any string parameter is null or whitespace.
public CosmosChatMessageStore(string accountEndpoint, TokenCredential tokenCredential, string databaseId, string containerId, string conversationId)
: this(new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), Throw.IfNull(tokenCredential)), databaseId, containerId, conversationId, ownsClient: true)
{
}
///
/// Initializes a new instance of the class using an existing .
///
/// The instance to use for Cosmos DB operations.
/// The identifier of the Cosmos DB database.
/// The identifier of the Cosmos DB container.
/// Thrown when is null.
/// Thrown when any string parameter is null or whitespace.
public CosmosChatMessageStore(CosmosClient cosmosClient, string databaseId, string containerId)
: this(cosmosClient, databaseId, containerId, Guid.NewGuid().ToString("N"))
{
}
///
/// Initializes a new instance of the class using an existing .
///
/// The instance to use for Cosmos DB operations.
/// The identifier of the Cosmos DB database.
/// The identifier of the Cosmos DB container.
/// The unique identifier for this conversation thread.
/// Thrown when is null.
/// Thrown when any string parameter is null or whitespace.
public CosmosChatMessageStore(CosmosClient cosmosClient, string databaseId, string containerId, string conversationId)
: this(cosmosClient, databaseId, containerId, conversationId, ownsClient: false)
{
}
///
/// Initializes a new instance of the class using a connection string with hierarchical partition keys.
///
/// The Cosmos DB connection string.
/// The identifier of the Cosmos DB database.
/// The identifier of the Cosmos DB container.
/// The tenant identifier for hierarchical partitioning.
/// The user identifier for hierarchical partitioning.
/// The session identifier for hierarchical partitioning.
/// Thrown when any required parameter is null.
/// Thrown when any string parameter is null or whitespace.
public CosmosChatMessageStore(string connectionString, string databaseId, string containerId, string tenantId, string userId, string sessionId)
: this(new CosmosClient(Throw.IfNullOrWhitespace(connectionString)), databaseId, containerId, Throw.IfNullOrWhitespace(sessionId), ownsClient: true, Throw.IfNullOrWhitespace(tenantId), Throw.IfNullOrWhitespace(userId))
{
}
///
/// Initializes a new instance of the class using a TokenCredential for authentication with hierarchical partition keys.
///
/// The Cosmos DB account endpoint URI.
/// The TokenCredential to use for authentication (e.g., DefaultAzureCredential, ManagedIdentityCredential).
/// The identifier of the Cosmos DB database.
/// The identifier of the Cosmos DB container.
/// The tenant identifier for hierarchical partitioning.
/// The user identifier for hierarchical partitioning.
/// The session identifier for hierarchical partitioning.
/// Thrown when any required parameter is null.
/// Thrown when any string parameter is null or whitespace.
public CosmosChatMessageStore(string accountEndpoint, TokenCredential tokenCredential, string databaseId, string containerId, string tenantId, string userId, string sessionId)
: this(new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), Throw.IfNull(tokenCredential)), databaseId, containerId, Throw.IfNullOrWhitespace(sessionId), ownsClient: true, Throw.IfNullOrWhitespace(tenantId), Throw.IfNullOrWhitespace(userId))
{
}
///
/// Initializes a new instance of the class using an existing with hierarchical partition keys.
///
/// The instance to use for Cosmos DB operations.
/// The identifier of the Cosmos DB database.
/// The identifier of the Cosmos DB container.
/// The tenant identifier for hierarchical partitioning.
/// The user identifier for hierarchical partitioning.
/// The session identifier for hierarchical partitioning.
/// Thrown when is null.
/// Thrown when any string parameter is null or whitespace.
public CosmosChatMessageStore(CosmosClient cosmosClient, string databaseId, string containerId, string tenantId, string userId, string sessionId)
: this(cosmosClient, databaseId, containerId, Throw.IfNullOrWhitespace(sessionId), ownsClient: false, Throw.IfNullOrWhitespace(tenantId), Throw.IfNullOrWhitespace(userId))
{
}
///
/// Creates a new instance of the class from previously serialized state.
///
/// The instance to use for Cosmos DB operations.
/// A representing the serialized state of the message store.
/// The identifier of the Cosmos DB database.
/// The identifier of the Cosmos DB container.
/// Optional settings for customizing the JSON deserialization process.
/// A new instance of initialized from the serialized state.
/// Thrown when is null.
/// Thrown when the serialized state cannot be deserialized.
public static CosmosChatMessageStore CreateFromSerializedState(CosmosClient cosmosClient, JsonElement serializedStoreState, string databaseId, string containerId, JsonSerializerOptions? jsonSerializerOptions = null)
{
Throw.IfNull(cosmosClient);
Throw.IfNullOrWhitespace(databaseId);
Throw.IfNullOrWhitespace(containerId);
if (serializedStoreState.ValueKind is not JsonValueKind.Object)
{
throw new ArgumentException("Invalid serialized state", nameof(serializedStoreState));
}
var state = serializedStoreState.Deserialize(jsonSerializerOptions);
if (state?.ConversationIdentifier is not { } conversationId)
{
throw new ArgumentException("Invalid serialized state", nameof(serializedStoreState));
}
// Use the internal constructor with all parameters to ensure partition key logic is centralized
return state.UseHierarchicalPartitioning && state.TenantId != null && state.UserId != null
? new CosmosChatMessageStore(cosmosClient, databaseId, containerId, conversationId, ownsClient: false, state.TenantId, state.UserId)
: new CosmosChatMessageStore(cosmosClient, databaseId, containerId, conversationId, ownsClient: false);
}
///
public override async Task> GetMessagesAsync(CancellationToken cancellationToken = default)
{
#pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks
if (this._disposed)
{
throw new ObjectDisposedException(this.GetType().FullName);
}
#pragma warning restore CA1513
// Fetch most recent messages in descending order when limit is set, then reverse to ascending
var orderDirection = this.MaxMessagesToRetrieve.HasValue ? "DESC" : "ASC";
var query = new QueryDefinition($"SELECT * FROM c WHERE c.conversationId = @conversationId AND c.type = @type ORDER BY c.timestamp {orderDirection}")
.WithParameter("@conversationId", this.ConversationId)
.WithParameter("@type", "ChatMessage");
var iterator = this._container.GetItemQueryIterator(query, requestOptions: new QueryRequestOptions
{
PartitionKey = this._partitionKey,
MaxItemCount = this.MaxItemCount // Configurable query performance
});
var messages = new List();
while (iterator.HasMoreResults)
{
var response = await iterator.ReadNextAsync(cancellationToken).ConfigureAwait(false);
foreach (var document in response)
{
if (this.MaxMessagesToRetrieve.HasValue && messages.Count >= this.MaxMessagesToRetrieve.Value)
{
break;
}
if (!string.IsNullOrEmpty(document.Message))
{
var message = JsonSerializer.Deserialize(document.Message, s_defaultJsonOptions);
if (message != null)
{
messages.Add(message);
}
}
}
if (this.MaxMessagesToRetrieve.HasValue && messages.Count >= this.MaxMessagesToRetrieve.Value)
{
break;
}
}
// If we fetched in descending order (most recent first), reverse to ascending order
if (this.MaxMessagesToRetrieve.HasValue)
{
messages.Reverse();
}
return messages;
}
///
public override async Task AddMessagesAsync(IEnumerable messages, CancellationToken cancellationToken = default)
{
if (messages is null)
{
throw new ArgumentNullException(nameof(messages));
}
#pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks
if (this._disposed)
{
throw new ObjectDisposedException(this.GetType().FullName);
}
#pragma warning restore CA1513
var messageList = messages as IReadOnlyCollection ?? messages.ToList();
if (messageList.Count == 0)
{
return;
}
// Use transactional batch for atomic operations
if (messageList.Count > 1)
{
await this.AddMessagesInBatchAsync(messageList, cancellationToken).ConfigureAwait(false);
}
else
{
await this.AddSingleMessageAsync(messageList.First(), cancellationToken).ConfigureAwait(false);
}
}
///
/// Adds multiple messages using transactional batch operations for atomicity.
///
private async Task AddMessagesInBatchAsync(IReadOnlyCollection messages, CancellationToken cancellationToken)
{
var currentTimestamp = DateTimeOffset.UtcNow.ToUnixTimeSeconds();
// Process messages in optimal batch sizes
for (int i = 0; i < messages.Count; i += this.MaxBatchSize)
{
var batchMessages = messages.Skip(i).Take(this.MaxBatchSize).ToList();
await this.ExecuteBatchOperationAsync(batchMessages, currentTimestamp, cancellationToken).ConfigureAwait(false);
}
}
///
/// Executes a single batch operation with enhanced error handling.
/// Cosmos SDK handles throttling (429) retries automatically.
///
private async Task ExecuteBatchOperationAsync(List messages, long timestamp, CancellationToken cancellationToken)
{
// Create all documents upfront for validation and batch operation
var documents = new List(messages.Count);
foreach (var message in messages)
{
documents.Add(this.CreateMessageDocument(message, timestamp));
}
// Defensive check: Verify all messages share the same partition key values
// In hierarchical partitioning, this means same tenantId, userId, and sessionId
// In simple partitioning, this means same conversationId
if (documents.Count > 0)
{
if (this._useHierarchicalPartitioning)
{
// Verify all documents have matching hierarchical partition key components
var firstDoc = documents[0];
if (!documents.All(d => d.TenantId == firstDoc.TenantId && d.UserId == firstDoc.UserId && d.SessionId == firstDoc.SessionId))
{
throw new InvalidOperationException("All messages in a batch must share the same partition key values (tenantId, userId, sessionId).");
}
}
else
{
// Verify all documents have matching conversationId
var firstConversationId = documents[0].ConversationId;
if (!documents.All(d => d.ConversationId == firstConversationId))
{
throw new InvalidOperationException("All messages in a batch must share the same partition key value (conversationId).");
}
}
}
// All messages in this store share the same partition key by design
// Transactional batches require all items to share the same partition key
var batch = this._container.CreateTransactionalBatch(this._partitionKey);
foreach (var document in documents)
{
batch.CreateItem(document);
}
try
{
var response = await batch.ExecuteAsync(cancellationToken).ConfigureAwait(false);
if (!response.IsSuccessStatusCode)
{
throw new InvalidOperationException($"Batch operation failed with status: {response.StatusCode}. Details: {response.ErrorMessage}");
}
}
catch (CosmosException ex) when (ex.StatusCode == System.Net.HttpStatusCode.RequestEntityTooLarge)
{
// If batch is too large, split into smaller batches
if (messages.Count == 1)
{
// Can't split further, use single operation
await this.AddSingleMessageAsync(messages[0], cancellationToken).ConfigureAwait(false);
return;
}
// Split the batch in half and retry
var midpoint = messages.Count / 2;
var firstHalf = messages.Take(midpoint).ToList();
var secondHalf = messages.Skip(midpoint).ToList();
await this.ExecuteBatchOperationAsync(firstHalf, timestamp, cancellationToken).ConfigureAwait(false);
await this.ExecuteBatchOperationAsync(secondHalf, timestamp, cancellationToken).ConfigureAwait(false);
}
}
///
/// Adds a single message to the store.
///
private async Task AddSingleMessageAsync(ChatMessage message, CancellationToken cancellationToken)
{
var document = this.CreateMessageDocument(message, DateTimeOffset.UtcNow.ToUnixTimeSeconds());
try
{
await this._container.CreateItemAsync(document, this._partitionKey, cancellationToken: cancellationToken).ConfigureAwait(false);
}
catch (CosmosException ex) when (ex.StatusCode == System.Net.HttpStatusCode.RequestEntityTooLarge)
{
throw new InvalidOperationException(
"Message exceeds Cosmos DB's maximum item size limit of 2MB. " +
"Message ID: " + message.MessageId + ", Serialized size is too large. " +
"Consider reducing message content or splitting into smaller messages.",
ex);
}
}
///
/// Creates a message document with enhanced metadata.
///
private CosmosMessageDocument CreateMessageDocument(ChatMessage message, long timestamp)
{
return new CosmosMessageDocument
{
Id = Guid.NewGuid().ToString(),
ConversationId = this.ConversationId,
Timestamp = timestamp,
MessageId = message.MessageId,
Role = message.Role.Value,
Message = JsonSerializer.Serialize(message, s_defaultJsonOptions),
Type = "ChatMessage", // Type discriminator
Ttl = this.MessageTtlSeconds, // Configurable TTL
// Include hierarchical metadata when using hierarchical partitioning
TenantId = this._useHierarchicalPartitioning ? this._tenantId : null,
UserId = this._useHierarchicalPartitioning ? this._userId : null,
SessionId = this._useHierarchicalPartitioning ? this.ConversationId : null
};
}
///
public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null)
{
#pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks
if (this._disposed)
{
throw new ObjectDisposedException(this.GetType().FullName);
}
#pragma warning restore CA1513
var state = new StoreState
{
ConversationIdentifier = this.ConversationId,
TenantId = this._tenantId,
UserId = this._userId,
UseHierarchicalPartitioning = this._useHierarchicalPartitioning
};
var options = jsonSerializerOptions ?? s_defaultJsonOptions;
return JsonSerializer.SerializeToElement(state, options);
}
///
/// Gets the count of messages in this conversation.
/// This is an additional utility method beyond the base contract.
///
/// The cancellation token.
/// The number of messages in the conversation.
public async Task GetMessageCountAsync(CancellationToken cancellationToken = default)
{
#pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks
if (this._disposed)
{
throw new ObjectDisposedException(this.GetType().FullName);
}
#pragma warning restore CA1513
// Efficient count query
var query = new QueryDefinition("SELECT VALUE COUNT(1) FROM c WHERE c.conversationId = @conversationId AND c.Type = @type")
.WithParameter("@conversationId", this.ConversationId)
.WithParameter("@type", "ChatMessage");
var iterator = this._container.GetItemQueryIterator(query, requestOptions: new QueryRequestOptions
{
PartitionKey = this._partitionKey
});
// COUNT queries always return a result
var response = await iterator.ReadNextAsync(cancellationToken).ConfigureAwait(false);
return response.FirstOrDefault();
}
///
/// Deletes all messages in this conversation.
/// This is an additional utility method beyond the base contract.
///
/// The cancellation token.
/// The number of messages deleted.
public async Task ClearMessagesAsync(CancellationToken cancellationToken = default)
{
#pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks
if (this._disposed)
{
throw new ObjectDisposedException(this.GetType().FullName);
}
#pragma warning restore CA1513
// Batch delete for efficiency
var query = new QueryDefinition("SELECT VALUE c.id FROM c WHERE c.conversationId = @conversationId AND c.Type = @type")
.WithParameter("@conversationId", this.ConversationId)
.WithParameter("@type", "ChatMessage");
var iterator = this._container.GetItemQueryIterator(query, requestOptions: new QueryRequestOptions
{
PartitionKey = this._partitionKey,
MaxItemCount = this.MaxItemCount
});
var deletedCount = 0;
while (iterator.HasMoreResults)
{
var response = await iterator.ReadNextAsync(cancellationToken).ConfigureAwait(false);
var batch = this._container.CreateTransactionalBatch(this._partitionKey);
var batchItemCount = 0;
foreach (var itemId in response)
{
if (!string.IsNullOrEmpty(itemId))
{
batch.DeleteItem(itemId);
batchItemCount++;
deletedCount++;
}
}
if (batchItemCount > 0)
{
await batch.ExecuteAsync(cancellationToken).ConfigureAwait(false);
}
}
return deletedCount;
}
///
public void Dispose()
{
if (!this._disposed)
{
if (this._ownsClient)
{
this._cosmosClient?.Dispose();
}
this._disposed = true;
}
}
private sealed class StoreState
{
public string ConversationIdentifier { get; set; } = string.Empty;
public string? TenantId { get; set; }
public string? UserId { get; set; }
public bool UseHierarchicalPartitioning { get; set; }
}
///
/// Represents a document stored in Cosmos DB for chat messages.
///
[SuppressMessage("Performance", "CA1812:Avoid uninstantiated internal classes", Justification = "Instantiated by Cosmos DB operations")]
private sealed class CosmosMessageDocument
{
[Newtonsoft.Json.JsonProperty("id")]
public string Id { get; set; } = string.Empty;
[Newtonsoft.Json.JsonProperty("conversationId")]
public string ConversationId { get; set; } = string.Empty;
[Newtonsoft.Json.JsonProperty("timestamp")]
public long Timestamp { get; set; }
[Newtonsoft.Json.JsonProperty("messageId")]
public string? MessageId { get; set; }
[Newtonsoft.Json.JsonProperty("role")]
public string? Role { get; set; }
[Newtonsoft.Json.JsonProperty("message")]
public string Message { get; set; } = string.Empty;
[Newtonsoft.Json.JsonProperty("type")]
public string Type { get; set; } = string.Empty;
[Newtonsoft.Json.JsonProperty("ttl")]
public int? Ttl { get; set; }
///
/// Tenant ID for hierarchical partitioning scenarios (optional).
///
[Newtonsoft.Json.JsonProperty("tenantId")]
public string? TenantId { get; set; }
///
/// User ID for hierarchical partitioning scenarios (optional).
///
[Newtonsoft.Json.JsonProperty("userId")]
public string? UserId { get; set; }
///
/// Session ID for hierarchical partitioning scenarios (same as ConversationId for compatibility).
///
[Newtonsoft.Json.JsonProperty("sessionId")]
public string? SessionId { get; set; }
}
}