.NET: [Feature]: Allow customizing partition key strategy in CosmosChatHistoryProvider #5601

Activity

1. Mark-Good

   added the

   Feature

   issue type on May 1, 2026
2. Evan Mattson (moonbox3)

   added

   .NETUsage: [Issues, PRs], Target: .NetUsage: [Issues, PRs], Target: .Net

   triageUsage: [Issues], Target: All issues that still need to be triagedUsage: [Issues], Target: All issues that still need to be triaged

   on May 1, 2026
3. chetantoshniwal

   assigned

   westey (westey-m)

   on May 5, 2026
4. chetantoshniwal

   removed

   triageUsage: [Issues], Target: All issues that still need to be triagedUsage: [Issues], Target: All issues that still need to be triaged

   on May 5, 2026
5. chetantoshniwal

   added this to Agent Frameworkon May 5, 2026
6. westey-m commented on Oct 8, 2026

   

   westey (westey-m)

   on Oct 8, 2026

   Collaborator

   More actions

   /plan

   👀React with 👀1Reacted by github-actions
7. github-actions commented on Oct 8, 2026

   

   github-actionsbot

   on Oct 8, 2026 – with GitHub Actions

   Contributor

   More actions

   Summary

   The issue is about CosmosChatHistoryProvider (in dotnet/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:

   • A flat key (/conversationId) — used when State.TenantId or State.UserId is not set.
   • A "hierarchical" key (tenantId → userId → conversationId) — used only when both TenantId and UserId are set on State.

   I confirmed in the code (lines 205–222 and 453–470) that the same condition, UseHierarchicalPartitioning(state), controls two different things at once:

   1. Which partition-key shape is built (BuildPartitionKey).
   2. Whether TenantId, UserId, and SessionId are written onto each stored document at all — when hierarchical partitioning is not used, those fields are forced to null (see CreateMessageDocument, lines 453–470).

   So today there is no way to keep a flat /conversationId partition 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 recording TenantId/UserId on 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):

   1. Always persist tenant/user/session fields when present. In CreateMessageDocument, stop gating TenantId/UserId/SessionId on UseHierarchicalPartitioning(state). Instead, write them directly from state.TenantId / state.UserId / state.ConversationId whenever they are set, independent of which partition key shape is used. This decouples "what's stored" from "how it's partitioned."
   2. Add an optional partition-key override delegate. Add a constructor parameter such as Func<State, PartitionKey>? partitionKeyBuilder to CosmosChatHistoryProvider, stored in a field and used wherever BuildPartitionKey(state) is currently called (the GetMessagesAsync, StoreChatHistoryAsync, and count/other helper methods). When the delegate is null, fall back to today's default logic, so existing callers see no behavior change.
   3. Adjust the batch-validation check. ExecuteBatchOperationAsync currently branches its "all messages share the same partition key" check on UseHierarchicalPartitioning(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.
   4. Update XML doc comments on the constructor(s), State class, and the changed methods to describe the new, decoupled behavior.
   5. Tests (likely added to dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatHistoryProviderTests.cs):
      • A case with flat partitioning (no tenant/user) confirming the document still omits those fields (current behavior preserved).
      • A new case with a custom partitionKeyBuilder plus TenantId/UserId set, confirming the flat key is used for routing and the fields are still written to the document.
      • A case confirming default behavior (hierarchical partitioning) is unchanged when no override is supplied.
      • A batch-write case exercising the updated validation logic with a custom partition key builder.
   6. Documentation/samples: if any sample or README under 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 /conversationId key while still setting TenantId/UserId on State — 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/sessionId fields where they were null before. 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 being null for flat-partitioned documents.

   The smaller interim alternative mentioned in the issue (removing sealed from the class and marking BuildPartitionKey, UseHierarchicalPartitioning, and CreateMessageDocument as protected 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 the protected virtual option only as a fallback if a constructor-based delegate design is rejected.

   Generated by Issue Planning for #5601 · copilot · auto · 29.1 AIC · ⌖ 6.09 AIC · ⊞ 8.9K · ◷
8. westey-m commented on Oct 8, 2026

   

   westey (westey-m)

   on Oct 8, 2026

   Collaborator

   More actions

   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.
9. westey (westey-m)

   assigned

   Copilot

   on Oct 8, 2026
10. Copilot

    linked a pull request that will close this issue.NET: [BREAKING] Decouple Cosmos chat metadata from partition-key strategy #9214on Oct 8, 2026
11. Vincent Biret (baywet)

    closed this as completedin #9214on Oct 9, 2026