session token data optimization

sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py

@@ -3336,6 +3336,10 @@ def __GetBodiesFromQueryResult(result: dict[str, Any]) -> list[dict[str, Any]]:
                 EPK_sub_range = routing_range.Range(range_min=max(single_range.min, feed_range_epk.min),
                                                     range_max=min(single_range.max, feed_range_epk.max),
                                                     isMinInclusive=True, isMaxInclusive=False)
+
+                # set the session token for this specific partition to avoid sending compound token for all partitions
+                base.set_session_token_header(self, req_headers, path, request_params, options,
+                                              over_lapping_range["id"])
                 if single_range.min == EPK_sub_range.min and EPK_sub_range.max == single_range.max:
                     # The Epk Sub Range spans exactly one physical partition
                     # In this case we can route to the physical pk range id

sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py

@@ -3128,6 +3128,10 @@ def __GetBodiesFromQueryResult(result: dict[str, Any]) -> list[dict[str, Any]]:
                 EPK_sub_range = routing_range.Range(range_min=max(single_range.min, feed_range_epk.min),
                                                     range_max=min(single_range.max, feed_range_epk.max),
                                                     isMinInclusive=True, isMaxInclusive=False)
+
+                # set the session token for this specific partition to avoid sending compound token for all partitions
+                await base.set_session_token_header_async(self, req_headers, path, request_params, options,
+                                              over_lapping_range["id"])
                 if single_range.min == EPK_sub_range.min and EPK_sub_range.max == single_range.max:
                     # The Epk Sub Range spans exactly one physical partition
                     # In this case we can route to the physical pk range id

sdk/cosmos/azure-cosmos/tests/test_session.py

@@ -153,6 +153,51 @@ def test_session_token_sm_for_ops(self):
         assert self.created_db.client_connection.last_response_headers.get(HttpHeaders.SessionToken) is not None
         assert self.created_db.client_connection.last_response_headers.get(HttpHeaders.SessionToken) != batch_response_token
 
+    def test_session_token_compound_not_sent_for_single_partition_query(self):
+        """
+        Verify that when querying with a feed range (single physical partition),
+        only that partition's session token is sent, not the entire compound token.
+        """
+        test_container = self.created_db.create_container(
+            "Container query test" + str(uuid.uuid4()),
+            PartitionKey(path="/pk"),
+            offer_throughput=11000
+        )
+
+        try:
+            # Create items across multiple partition keys
+            for i in range(100):
+                test_container.create_item({
+                    'id': str(uuid.uuid4()),
+                    'pk': f"pk_{i:04d}"
+                })
+
+            # Get feed ranges and verify multiple exist
+            feed_ranges = list(test_container.read_feed_ranges())
+            self.assertGreater(len(feed_ranges), 1, "Expected multiple feed ranges")
+
+            # Capture session token sent with feed range query
+            captured_session_token = {}
+
+            def capture_session_token(request):
+                captured_session_token['token'] = request.http_request.headers.get(HttpHeaders.SessionToken)
+
+            # Query with single feed range
+            list(test_container.query_items(
+                query="SELECT * FROM c",
+                feed_range=feed_ranges[0],
+                raw_request_hook=capture_session_token
+            ))
+
+            # Verify only single partition token was sent
+            token = captured_session_token.get('token')
+            self.assertIsNotNone(token, "Session token should be present")
+            self.assertNotIn(',', token,
+                             f"Expected single partition token, got compound token: {token}")
+
+        finally:
+            self.created_db.delete_container(test_container)
+
     def test_session_token_with_space_in_container_name(self):
 
         # Session token should not be sent for control plane operations

sdk/cosmos/azure-cosmos/tests/test_session_async.py

@@ -86,6 +86,51 @@ def manual_token_hook(request):
             raw_request_hook=manual_token_hook
         )
 
+    async def test_session_token_compound_not_sent_for_single_partition_query_async(self):
+        """
+        Verify that when querying with a feed range (single physical partition),
+        only that partition's session token is sent, not the entire compound token.
+        """
+        test_container = await self.created_db.create_container(
+            "Container query test" + str(uuid.uuid4()),
+            PartitionKey(path="/pk"),
+            offer_throughput=11000
+        )
+
+        try:
+            # Create items across multiple partition keys
+            for i in range(100):
+                await test_container.create_item({
+                    'id': str(uuid.uuid4()),
+                    'pk': f"pk_{i:04d}"
+                })
+
+            # Get feed ranges and verify multiple exist
+            feed_ranges = [feed_range async for feed_range in test_container.read_feed_ranges()]
+            self.assertGreater(len(feed_ranges), 1, "Expected multiple feed ranges")
+
+            # Capture session token sent with feed range query
+            captured_session_token = {}
+
+            def capture_session_token(request):
+                captured_session_token['token'] = request.http_request.headers.get(HttpHeaders.SessionToken)
+
+            # Query with single feed range
+            _ = [item async for item in test_container.query_items(
+                query="SELECT * FROM c",
+                feed_range=feed_ranges[0],
+                raw_request_hook=capture_session_token
+            )]
+
+            # Verify only single partition token was sent
+            token = captured_session_token.get('token')
+            self.assertIsNotNone(token, "Session token should be present")
+            self.assertNotIn(',', token,
+                             f"Expected single partition token, got compound token: {token}")
+
+        finally:
+            await self.created_db.delete_container(test_container)
+
     async def test_manual_session_token_override_async(self):
         # Create an item to get a valid session token from the response
         created_document = await self.created_container.create_item(