[Cosmos] Session container fixes new branch (#41678)

* session container fixes

* async changes

* sync changes

* mypy/pylint

* Update _session.py

* mark query plan as fetched for query

* adjust logic after merging

* Update _base.py

* Update _base.py

* tests - missing sync mwr

* sync mwr tests, test fixes

* Update test_session_async.py

* Update test_backwards_compatibility.py

* Update test_backwards_compatibility.py

* Update test_backwards_compatibility_async.py

* Update test_backwards_compatibility_async.py

* Update test_backwards_compatibility.py

* Update execution_dispatcher.py

* merge leftovers

* slip

* Update test_backwards_compatibility.py

* Update test_backwards_compatibility.py

* Update test_query_hybrid_search_async.py

* further changes, changelog

* add tests

* typehint

* address comments

* Update _session.py

* Update _base.py

* ci tests

* merging main

* Update _cosmos_client_connection.py

* tests

* change session token logic

* Update _session.py

* small fixes

* Update _session.py

* Update _session.py

* updates

* update tests that checked for compound tokens - this will no longer be the case

* Update test_config.py

* Update execution_dispatcher.py

* remove query logic

* remove partition split testing

* oylint

* Update dev_requirements.txt

* delete duplicate

* update tests for partition split

* change timeout of split partition key tests

reduces timeout from 25 minutes to 7 minutes

* update test

* Test fixes

Fixes tests to accomodate the extra readfeed that may  happen during a read item operation. This also includes other general test fixes.

* update tests

* update tests

Updates tests
and fixes a key error issue in session.py

* add logger to session token

* Update _session.py

* Update _session.py

* Update test_session.py

* async changes

* small test fixes - align sync and async

* small updates

* Update documents.py

* test updates

* Update _session.py

* async tests

* pylint

* tests

* logic fixes, test updates

* leftover changes

* match session logic exactly

* propagate headers to request params for retry logic

* Update CHANGELOG.md

* revert session token dropping

* last test updates

* Update _session.py

* update session token on errors

* reformat internal exceptions

* change string for constant

* address comments

---------

Co-authored-by: Simon Moreno <[email protected]>

Author: bambriz

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -8,6 +8,8 @@
 #### Breaking Changes
 
 #### Bugs Fixed
+* Fixed session container session token logic. The SDK will now only send the relevant partition-local session tokens for read document requests and write requests when multi-region writes are enabled, as opposed to the entire compound session token for the container for every document request. See [PR 41678](https://github.com/Azure/azure-sdk-for-python/pull/41678).
+* Write requests for single-write region accounts will no longer send session tokens when using session consistency. See [PR 41678](https://github.com/Azure/azure-sdk-for-python/pull/41678).
 * Fixed bug where container cache was not being properly updated resulting in unnecessary extra requests. See [PR 42143](https://github.com/Azure/azure-sdk-for-python/pull/42143).
 
 #### Other Changes
@@ -16,7 +18,7 @@
 ### 4.14.0b1 (2025-07-14)
 
 #### Features Added
-* Added option to enable automatic retries for write operations. See [PR 41272](https://github.com/Azure/azure-sdk-for-python/pull/41272).
+* Added option to enable automatic retries for write document operations. See [PR 41272](https://github.com/Azure/azure-sdk-for-python/pull/41272).
 
 #### Breaking Changes
 * Adds cross region retries when no preferred locations are set. This is only a breaking change for customers using bounded staleness consistency. See [PR 39714](https://github.com/Azure/azure-sdk-for-python/pull/39714)

sdk/cosmos/azure-cosmos/azure/cosmos/_base.py

@@ -45,7 +45,9 @@
 if TYPE_CHECKING:
     from ._cosmos_client_connection import CosmosClientConnection
     from .aio._cosmos_client_connection_async import CosmosClientConnection as AsyncClientConnection
+    from ._request_object import RequestObject
 
+# pylint: disable=protected-access
 
 _COMMON_OPTIONS = {
     'initial_headers': 'initialHeaders',
@@ -174,37 +176,9 @@ def GetHeaders(  # pylint: disable=too-many-statements,too-many-branches
     if options.get("indexingDirective"):
         headers[http_constants.HttpHeaders.IndexingDirective] = options["indexingDirective"]
 
-    consistency_level = None
-
-    # get default client consistency level
-    default_client_consistency_level = headers.get(http_constants.HttpHeaders.ConsistencyLevel)
-
-    # set consistency level. check if set via options, this will override the default
+    # set request consistency level - if session consistency, the client should be setting this on its own
     if options.get("consistencyLevel"):
-        consistency_level = options["consistencyLevel"]
-        # TODO: move this line outside of if-else cause to remove the code duplication
-        headers[http_constants.HttpHeaders.ConsistencyLevel] = consistency_level
-    elif default_client_consistency_level is not None:
-        consistency_level = default_client_consistency_level
-        headers[http_constants.HttpHeaders.ConsistencyLevel] = consistency_level
-
-    # figure out if consistency level for this request is session
-    is_session_consistency = consistency_level == documents.ConsistencyLevel.Session
-
-    # set session token if required
-    if is_session_consistency is True and not IsMasterResource(resource_type):
-        # if there is a token set via option, then use it to override default
-        if options.get("sessionToken"):
-            headers[http_constants.HttpHeaders.SessionToken] = options["sessionToken"]
-        else:
-            # check if the client's default consistency is session (and request consistency level is same),
-            # then update from session container
-            if default_client_consistency_level == documents.ConsistencyLevel.Session and \
-                    cosmos_client_connection.session:
-                # populate session token from the client's session container
-                headers[http_constants.HttpHeaders.SessionToken] = cosmos_client_connection.session.get_session_token(
-                    path
-                )
+        headers[http_constants.HttpHeaders.ConsistencyLevel] = options["consistencyLevel"]
 
     if options.get("enableScanInQuery"):
         headers[http_constants.HttpHeaders.EnableScanInQuery] = options["enableScanInQuery"]
@@ -349,6 +323,77 @@ def GetHeaders(  # pylint: disable=too-many-statements,too-many-branches
 
     return headers
 
+def _is_session_token_request(
+        cosmos_client_connection: Union["CosmosClientConnection", "AsyncClientConnection"],
+        headers: dict,
+        request_object: "RequestObject") -> bool:
+    consistency_level = headers.get(http_constants.HttpHeaders.ConsistencyLevel)
+    # Figure out if consistency level for this request is session
+    is_session_consistency = consistency_level == documents.ConsistencyLevel.Session
+
+    # Verify that it is not a metadata request, and that it is either a read request, batch request, or an account
+    # configured to use multiple write regions. Batch requests are special-cased because they can contain both read and
+    # write operations, and we want to use session consistency for the read operations.
+    return (is_session_consistency is True and cosmos_client_connection.session is not None
+            and not IsMasterResource(request_object.resource_type)
+            and (documents._OperationType.IsReadOnlyOperation(request_object.operation_type)
+                 or request_object.operation_type == "Batch"
+                 or cosmos_client_connection._global_endpoint_manager.can_use_multiple_write_locations(request_object)))
+
+
+def set_session_token_header(
+        cosmos_client_connection: Union["CosmosClientConnection", "AsyncClientConnection"],
+        headers: dict,
+        path: str,
+        request_object: "RequestObject",
+        options: Mapping[str, Any],
+        partition_key_range_id: Optional[str] = None) -> None:
+    # set session token if required
+    if _is_session_token_request(cosmos_client_connection, headers, request_object):
+        # if there is a token set via option, then use it to override default
+        if options.get("sessionToken"):
+            headers[http_constants.HttpHeaders.SessionToken] = options["sessionToken"]
+        else:
+            # check if the client's default consistency is session (and request consistency level is same),
+            # then update from session container
+            if headers[http_constants.HttpHeaders.ConsistencyLevel] == documents.ConsistencyLevel.Session and \
+                    cosmos_client_connection.session:
+                # populate session token from the client's session container
+                session_token = (
+                    cosmos_client_connection.session.get_session_token(path,
+                                                                options.get('partitionKey'),
+                                                                cosmos_client_connection._container_properties_cache,
+                                                                cosmos_client_connection._routing_map_provider,
+                                                                partition_key_range_id))
+                if session_token != "":
+                    headers[http_constants.HttpHeaders.SessionToken] = session_token
+
+async def set_session_token_header_async(
+        cosmos_client_connection: Union["CosmosClientConnection", "AsyncClientConnection"],
+        headers: dict,
+        path: str,
+        request_object: "RequestObject",
+        options: Mapping[str, Any],
+        partition_key_range_id: Optional[str] = None) -> None:
+    # set session token if required
+    if _is_session_token_request(cosmos_client_connection, headers, request_object):
+        # if there is a token set via option, then use it to override default
+        if options.get("sessionToken"):
+            headers[http_constants.HttpHeaders.SessionToken] = options["sessionToken"]
+        else:
+            # check if the client's default consistency is session (and request consistency level is same),
+            # then update from session container
+            if headers[http_constants.HttpHeaders.ConsistencyLevel] == documents.ConsistencyLevel.Session and \
+                    cosmos_client_connection.session:
+                # populate session token from the client's session container
+                session_token = \
+                    await cosmos_client_connection.session.get_session_token_async(path,
+                                                                options.get('partitionKey'),
+                                                                cosmos_client_connection._container_properties_cache,
+                                                                cosmos_client_connection._routing_map_provider,
+                                                                partition_key_range_id)
+                if session_token != "":
+                    headers[http_constants.HttpHeaders.SessionToken] = session_token
 
 def GetResourceIdOrFullNameFromLink(resource_link: str) -> str:
     """Gets resource id or full name from resource link.