[Internal] DTS: Adds partitionKeyRangeId to operation result to assemble valid session tokens in DTS

Microsoft.Azure.Cosmos/src/DistributedTransaction/DistributedTransactionCommitter.cs

@@ -208,11 +208,31 @@ internal static void MergeSessionTokens(
                     continue;
                 }
 
+                // SessionContainer.SetSessionToken expects the token in {pkRangeId}:{lsn} format.
+                // The backend may send pkRangeId separately or the token may already be assembled.
+                string tokenForHeader = result.SessionToken;
+                if (tokenForHeader.IndexOf(':') < 0)
+                {
+                    // Token is LSN-only; we need pkRangeId to assemble the full format.
+                    if (string.IsNullOrWhiteSpace(result.PartitionKeyRangeId))
+                    {
+                        // Cannot form a valid session token without pkRangeId; silently skip merging
+                        // this operation's token but continue processing the rest of the response.
+                        DefaultTrace.TraceWarning(
+                            "DTC operation index {0} (collection {1}) returned LSN-only session token without partitionKeyRangeId; skipping session token merge.",
+                            result.Index,
+                            operation.CollectionResourceId);
+                        continue;
+                    }
+
+                    tokenForHeader = result.PartitionKeyRangeId + ":" + tokenForHeader;
+                }
+
                 // Note: each SetSessionToken call acquires a write lock on the SessionContainer.
                 // For a future optimization, consider a batch-update API on ISessionContainer to
                 // reduce lock acquisitions when multiple operations target the same collection.
                 headers.Clear();
-                headers[HttpConstants.HttpHeaders.SessionToken] = result.SessionToken;
+                headers[HttpConstants.HttpHeaders.SessionToken] = tokenForHeader;
 
                 sessionContainer.SetSessionToken(
                     operation.CollectionResourceId,

Microsoft.Azure.Cosmos/src/DistributedTransaction/DistributedTransactionOperationResult.cs

@@ -35,6 +35,7 @@ internal DistributedTransactionOperationResult(DistributedTransactionOperationRe
             this.ETag = other.ETag;
             this.ResourceStream = other.ResourceStream;
             this.SessionToken = other.SessionToken;
+            this.PartitionKeyRangeId = other.PartitionKeyRangeId;
             this.RequestCharge = other.RequestCharge;
             this.ActivityId = other.ActivityId;
             this.Trace = other.Trace;
@@ -89,6 +90,15 @@ public DistributedTransactionOperationResult()
         [JsonPropertyName("sessionToken")]
         public virtual string SessionToken { get; internal set; }
 
+        /// <summary>
+        /// Gets the partition key range ID associated with the operation result.
+        /// When present, it is combined with <see cref="SessionToken"/> to form the
+        /// full session token in the format {partitionKeyRangeId}:{lsn}.
+        /// </summary>
+        [JsonInclude]
+        [JsonPropertyName("partitionKeyRangeId")]
+        public virtual string PartitionKeyRangeId { get; internal set; }
+
         /// <summary>
         /// Gets the resource stream associated with the operation result.
         /// The stream contains the raw response payload returned by the operation.

Microsoft.Azure.Cosmos/src/DistributedTransaction/DistributedTransactionSerializer.cs

@@ -26,6 +26,7 @@ internal static class DistributedTransactionSerializer
         internal const string Index = "index";
         internal const string ResourceBody = "resourceBody";
         internal const string SessionToken = "sessionToken";
+        internal const string PartitionKeyRangeId = "partitionKeyRangeId";
         internal const string ETag = "ifMatch";
         internal const string OperationType = "operationType";
         internal const string ResourceType = "resourceType";

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.Tests/DistributedTransaction/DistributedTransactionCommitterTests.cs

@@ -275,6 +275,199 @@ public async Task CommitTransactionAsync_MergesSessionTokens_OnFailureResponse()
                 "Session token should still be merged even when the DTC response indicates a failure.");
         }
 
+        [TestMethod]
+        [Description("When session token is LSN-only and partitionKeyRangeId is present, the token is assembled as {pkRangeId}:{lsn}")]
+        public async Task CommitTransactionAsync_AssemblesSessionToken_WhenPartitionKeyRangeIdIsPresent()
+        {
+            const string lsnOnly = "1#9#4=8#5=7";
+            const string pkRangeId = "0";
+            const string expectedToken = "0:1#9#4=8#5=7";
+
+            SessionContainer sessionContainer = new SessionContainer("testhost");
+
+            string responseJson = BuildDtcResponseJson(
+                new[] { (statusCode: 201, subStatusCode: (int?)null, sessionToken: lsnOnly, partitionKeyRangeId: pkRangeId) });
+
+            Mock<CosmosClientContext> mockContext = this.CreateMockContext(
+                sessionContainer,
+                responseContent: responseJson,
+                statusCode: HttpStatusCode.OK);
+
+            List<DistributedTransactionOperation> operations = new List<DistributedTransactionOperation>
+            {
+                new DistributedTransactionOperation(
+                    OperationType.Create,
+                    operationIndex: 0,
+                    DatabaseName,
+                    ContainerName,
+                    new PartitionKey("pk1"),
+                    id: "doc1")
+            };
+
+            DistributedTransactionCommitter committer = new DistributedTransactionCommitter(
+                operations, mockContext.Object);
+
+            await committer.CommitTransactionAsync(CancellationToken.None);
+
+            string storedToken = sessionContainer.GetSessionToken(DistributedTransactionConstants.GetCollectionFullName(DatabaseName, ContainerName));
+            Assert.AreEqual(expectedToken, storedToken,
+                "Session token should be assembled as {pkRangeId}:{lsn} when partitionKeyRangeId is present.");
+        }
+
+        [TestMethod]
+        [Description("When session token is LSN-only and partitionKeyRangeId is absent, merge is silently skipped")]
+        public async Task CommitTransactionAsync_SkipsMerge_WhenLsnOnlyAndPartitionKeyRangeIdIsAbsent()
+        {
+            const string lsnOnly = "1#9#4=8#5=7";
+
+            SessionContainer sessionContainer = new SessionContainer("testhost");
+
+            // No partitionKeyRangeId, and sessionToken has no ':' (LSN-only)
+            string responseJson = BuildDtcResponseJson(
+                new[] { (statusCode: 201, subStatusCode: (int?)null, sessionToken: lsnOnly, partitionKeyRangeId: (string)null) });
+
+            Mock<CosmosClientContext> mockContext = this.CreateMockContext(
+                sessionContainer,
+                responseContent: responseJson,
+                statusCode: HttpStatusCode.OK);
+
+            List<DistributedTransactionOperation> operations = new List<DistributedTransactionOperation>
+            {
+                new DistributedTransactionOperation(
+                    OperationType.Create,
+                    operationIndex: 0,
+                    DatabaseName,
+                    ContainerName,
+                    new PartitionKey("pk1"),
+                    id: "doc1")
+            };
+
+            DistributedTransactionCommitter committer = new DistributedTransactionCommitter(
+                operations, mockContext.Object);
+
+            await committer.CommitTransactionAsync(CancellationToken.None);
+
+            string storedToken = sessionContainer.GetSessionToken(DistributedTransactionConstants.GetCollectionFullName(DatabaseName, ContainerName));
+            Assert.IsTrue(string.IsNullOrEmpty(storedToken),
+                "SessionContainer should not be updated when partitionKeyRangeId is absent and session token is LSN-only.");
+        }
+
+        [TestMethod]
+        [Description("When session token already contains ':' (fully assembled), it is used as-is even without partitionKeyRangeId")]
+        public async Task CommitTransactionAsync_UsesPreAssembledSessionToken_WhenAlreadyContainsColon()
+        {
+            const string preAssembledToken = "0:1#9#4=8#5=7";
+
+            SessionContainer sessionContainer = new SessionContainer("testhost");
+
+            // Token already contains ':', so no pkRangeId needed
+            string responseJson = BuildDtcResponseJson(
+                new[] { (statusCode: 201, subStatusCode: (int?)null, sessionToken: preAssembledToken, partitionKeyRangeId: (string)null) });
+
+            Mock<CosmosClientContext> mockContext = this.CreateMockContext(
+                sessionContainer,
+                responseContent: responseJson,
+                statusCode: HttpStatusCode.OK);
+
+            List<DistributedTransactionOperation> operations = new List<DistributedTransactionOperation>
+            {
+                new DistributedTransactionOperation(
+                    OperationType.Create,
+                    operationIndex: 0,
+                    DatabaseName,
+                    ContainerName,
+                    new PartitionKey("pk1"),
+                    id: "doc1")
+            };
+
+            DistributedTransactionCommitter committer = new DistributedTransactionCommitter(
+                operations, mockContext.Object);
+
+            await committer.CommitTransactionAsync(CancellationToken.None);
+
+            string storedToken = sessionContainer.GetSessionToken(DistributedTransactionConstants.GetCollectionFullName(DatabaseName, ContainerName));
+            Assert.AreEqual(preAssembledToken, storedToken,
+                "Pre-assembled session token (with ':') should be used as-is.");
+        }
+
+        [TestMethod]
+        [Description("When session token already contains ':' AND partitionKeyRangeId is also present, the pre-assembled token takes precedence")]
+        public async Task CommitTransactionAsync_PreAssembledTokenTakesPrecedence_WhenBothTokenAndPartitionKeyRangeIdPresent()
+        {
+            const string preAssembledToken = "0:1#9#4=8#5=7";
+            const string differentPkRangeId = "5"; // would produce "5:1#9#4=8#5=7" if incorrectly used
+
+            SessionContainer sessionContainer = new SessionContainer("testhost");
+
+            string responseJson = BuildDtcResponseJson(
+                new[] { (statusCode: 201, subStatusCode: (int?)null, sessionToken: preAssembledToken, partitionKeyRangeId: differentPkRangeId) });
+
+            Mock<CosmosClientContext> mockContext = this.CreateMockContext(
+                sessionContainer,
+                responseContent: responseJson,
+                statusCode: HttpStatusCode.OK);
+
+            List<DistributedTransactionOperation> operations = new List<DistributedTransactionOperation>
+            {
+                new DistributedTransactionOperation(
+                    OperationType.Create,
+                    operationIndex: 0,
+                    DatabaseName,
+                    ContainerName,
+                    new PartitionKey("pk1"),
+                    id: "doc1")
+            };
+
+            DistributedTransactionCommitter committer = new DistributedTransactionCommitter(
+                operations, mockContext.Object);
+
+            await committer.CommitTransactionAsync(CancellationToken.None);
+
+            string storedToken = sessionContainer.GetSessionToken(DistributedTransactionConstants.GetCollectionFullName(DatabaseName, ContainerName));
+            Assert.AreEqual(preAssembledToken, storedToken,
+                "Pre-assembled session token (containing ':') must be used as-is; partitionKeyRangeId should not be prepended again.");
+        }
+
+        [DataTestMethod]
+        [DataRow("", DisplayName = "Empty string partitionKeyRangeId")]
+        [DataRow(" ", DisplayName = "Whitespace-only partitionKeyRangeId")]
+        [DataRow("   ", DisplayName = "Multiple whitespace partitionKeyRangeId")]
+        [Description("When partitionKeyRangeId is empty or whitespace-only and session token is LSN-only, merge is silently skipped")]
+        public async Task CommitTransactionAsync_SkipsMerge_WhenPartitionKeyRangeIdIsEmptyOrWhitespace(string pkRangeId)
+        {
+            const string lsnOnly = "1#9#4=8#5=7";
+
+            SessionContainer sessionContainer = new SessionContainer("testhost");
+
+            string responseJson = BuildDtcResponseJson(
+                new[] { (statusCode: 201, subStatusCode: (int?)null, sessionToken: lsnOnly, partitionKeyRangeId: pkRangeId) });
+
+            Mock<CosmosClientContext> mockContext = this.CreateMockContext(
+                sessionContainer,
+                responseContent: responseJson,
+                statusCode: HttpStatusCode.OK);
+
+            List<DistributedTransactionOperation> operations = new List<DistributedTransactionOperation>
+            {
+                new DistributedTransactionOperation(
+                    OperationType.Create,
+                    operationIndex: 0,
+                    DatabaseName,
+                    ContainerName,
+                    new PartitionKey("pk1"),
+                    id: "doc1")
+            };
+
+            DistributedTransactionCommitter committer = new DistributedTransactionCommitter(
+                operations, mockContext.Object);
+
+            await committer.CommitTransactionAsync(CancellationToken.None);
+
+            string storedToken = sessionContainer.GetSessionToken(DistributedTransactionConstants.GetCollectionFullName(DatabaseName, ContainerName));
+            Assert.IsTrue(string.IsNullOrEmpty(storedToken),
+                $"SessionContainer should not be updated when partitionKeyRangeId is '{pkRangeId}' (empty/whitespace).");
+        }
+
         // ─── Retry / Spec-Compliance Tests ─────────────────────────────────────
 
         [TestMethod]
@@ -768,11 +961,18 @@ private static string BuildDtcResponseJson(
             (int statusCode, string sessionToken)[] operations)
         {
             return BuildDtcResponseJson(
-                operations.Select(o => (o.statusCode, subStatusCode: (int?)null, o.sessionToken)).ToArray());
+                operations.Select(o => (o.statusCode, subStatusCode: (int?)null, o.sessionToken, partitionKeyRangeId: (string)null)).ToArray());
         }
 
         private static string BuildDtcResponseJson(
             (int statusCode, int? subStatusCode, string sessionToken)[] operations)
+        {
+            return BuildDtcResponseJson(
+                operations.Select(o => (o.statusCode, o.subStatusCode, o.sessionToken, partitionKeyRangeId: (string)null)).ToArray());
+        }
+
+        private static string BuildDtcResponseJson(
+            (int statusCode, int? subStatusCode, string sessionToken, string partitionKeyRangeId)[] operations)
         {
             StringBuilder sb = new StringBuilder();
             sb.Append(@"{""operationResponses"":[");
@@ -794,6 +994,11 @@ private static string BuildDtcResponseJson(
                     sb.Append($@",""{DistributedTransactionSerializer.SessionToken}"":""{operations[i].sessionToken}""");
                 }
 
+                if (operations[i].partitionKeyRangeId != null)
+                {
+                    sb.Append($@",""{DistributedTransactionSerializer.PartitionKeyRangeId}"":""{operations[i].partitionKeyRangeId}""");
+                }
+
                 sb.Append('}');
             }