Repository navigation
.NET: [Feature]: Allow customizing partition key strategy in CosmosChatHistoryProvider #5601
Description
Description
Problem
CosmosChatHistoryProvider couples two concerns: the partition key shape and which fields land on the document. Setting State.TenantId and State.UserId is what triggers hierarchical partitioning and what causes those values to be persisted on documents — they can't be decoupled. So there's no way to use flat /conversationId partitioning while still storing tenant/user metadata for filtering, RBAC scoping, or change-feed routing.
Motivation
Multi-tenant chatbots that support unauthenticated → authenticated user transitions need both: flat PK for continuity across the auth boundary, and tenant/user fields on documents for tenant isolation and admin queries. Currently you have to give up one.
Proposal
Two separable, backward-compatible changes:
- Always persist
TenantId/UserId/SessionIdon documents when set onState, regardless of partition key shape. - Add an optional
Func<State, PartitionKey>constructor parameter to overrideBuildPartitionKey. Defaults to current behavior when null.
Smaller interim alternative: remove sealed and mark BuildPartitionKey, UseHierarchicalPartitioning, and CreateMessageDocument as protected virtual.
Alternatives considered
- Storing tenant/user inside
ChatMessagemetadata — values get buried in the serializedmessageblob; no efficient filtering. - Writing a fully custom
ChatHistoryProvider— viable but requires reimplementing transactional batching, request-too-large splitting, TTL, and batch validation.
Source: dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs
Language/SDK
.NET
Metadata
Metadata
Assignees
Type
Projects
- Status
Activity
westey-m commentedon Oct 8, 2026
/plan
github-actions commentedon Oct 8, 2026
Summary
The issue is about
CosmosChatHistoryProvider(indotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs), the class that stores and loads chat messages in Azure Cosmos DB. A "partition key" is the value Cosmos DB uses to group and route documents; this provider currently supports two shapes:/conversationId) — used whenState.TenantIdorState.UserIdis not set.tenantId→userId→conversationId) — used only when bothTenantIdandUserIdare set onState.I confirmed in the code (lines 205–222 and 453–470) that the same condition,
UseHierarchicalPartitioning(state), controls two different things at once:BuildPartitionKey).TenantId,UserId, andSessionIdare written onto each stored document at all — when hierarchical partitioning is not used, those fields are forced tonull(seeCreateMessageDocument, lines 453–470).So today there is no way to keep a flat
/conversationIdpartition key (useful when a conversation starts anonymously and a user/tenant is attached later, since changing the partition key for existing documents isn't possible) while still recordingTenantId/UserIdon the documents for filtering, access scoping, or Cosmos DB's change feed. This matches the reporter's description, and I did not find any existing mechanism for overriding this behavior or any related open pull request.Proposed plan
This is a proposal only — no code has been changed. Assuming the maintainers favor the full solution (option 1 in the issue) over the smaller interim alternative (just loosening access modifiers):
CreateMessageDocument, stop gatingTenantId/UserId/SessionIdonUseHierarchicalPartitioning(state). Instead, write them directly fromstate.TenantId/state.UserId/state.ConversationIdwhenever they are set, independent of which partition key shape is used. This decouples "what's stored" from "how it's partitioned."Func<State, PartitionKey>? partitionKeyBuildertoCosmosChatHistoryProvider, stored in a field and used whereverBuildPartitionKey(state)is currently called (theGetMessagesAsync,StoreChatHistoryAsync, and count/other helper methods). When the delegate isnull, fall back to today's default logic, so existing callers see no behavior change.ExecuteBatchOperationAsynccurrently branches its "all messages share the same partition key" check onUseHierarchicalPartitioning(state). This needs to be updated so the check is based on whichever partition key strategy is active (the custom delegate's output, or default), not on whether hierarchical fields are populated.Stateclass, and the changed methods to describe the new, decoupled behavior.dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatHistoryProviderTests.cs):partitionKeyBuilderplusTenantId/UserIdset, confirming the flat key is used for routing and the fields are still written to the document.dotnet/samples/or package docs describes the tenant/user partitioning behavior, update it to mention the new override parameter and the decoupled metadata behavior.How to confirm this resolves the issue: after the change, someone could configure the provider with
partitionKeyBuilder: null(default, unchanged), or with a custom delegate returning a flat/conversationIdkey while still settingTenantId/UserIdonState— and see those fields persisted on every message document, enabling tenant-scoped queries and change-feed filtering without being forced into hierarchical partitioning.Trade-off to note: this is a breaking change in document shape only in the sense that more documents (those previously flat-partitioned with tenant/user state) will now carry
tenantId/userId/sessionIdfields where they werenullbefore. This seems acceptable and is what the issue is requesting, but it's worth maintainers confirming there's no reliance elsewhere (e.g., a query or index) on those fields beingnullfor flat-partitioned documents.The smaller interim alternative mentioned in the issue (removing
sealedfrom the class and markingBuildPartitionKey,UseHierarchicalPartitioning, andCreateMessageDocumentasprotected virtual) is simpler but pushes the maintenance burden of overriding onto each consumer and doesn't solve the "always persist metadata" piece by itself — I'd recommend the delegate-based approach above as the primary fix, with theprotected virtualoption only as a fallback if a constructor-based delegate design is rejected.westey-m commentedon Oct 8, 2026
I think the plan makes sense, but the list of optional parameters on CosmosChatHistoryProvider is getting too long, so we should also refactor it to use an options class, with required parameters on the constructor and optional going to options, similar to e.g. ChatHistoryMemoryProvider.