Fix: Lazy initialization of _InferenceService to prevent startup crash with AAD credentials (#46243)

aayush3011 · VenkataAnilKumar · commit bea59a7a3836 · 2026-04-18T10:07:41.000-04:00
diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md
@@ -7,6 +7,7 @@
 #### Breaking Changes
 
 #### Bugs Fixed
+* Fixed bug where `CosmosClient` construction with AAD credentials would crash at startup if the semantic reranking inference endpoint environment variable was not set, even when semantic reranking was not being used. The inference service is now lazily initialized on first use. See [PR 46243](https://github.com/Azure/azure-sdk-for-python/pull/46243)
 
 #### Other Changes
 
diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py
@@ -25,6 +25,7 @@
 """
 import logging
 import os
+import threading
 import urllib.parse
 import uuid
 from concurrent.futures.thread import ThreadPoolExecutor
@@ -262,8 +263,7 @@ def __init__( # pylint: disable=too-many-statements
         )
 
         self._inference_service: Optional[_InferenceService] = None
-        if self.aad_credentials:
-            self._inference_service = _InferenceService(self)
+        self._inference_service_lock = threading.Lock()
 
         # Query compatibility mode.
         # Allows to specify compatibility mode used by client when making query requests. Should be removed when
@@ -332,7 +332,18 @@ def _set_client_consistency_level(
             self.session = None
 
     def _get_inference_service(self) -> Optional[_InferenceService]:
-        """Get inference service instance"""
+        """Get inference service instance, lazily initializing on first access."""
+        if self._inference_service is None and self.aad_credentials:
+            with self._inference_service_lock:
+                if self._inference_service is None:
+                    try:
+                        self._inference_service = _InferenceService(self)
+                    except ValueError as e:
+                        raise exceptions.CosmosHttpResponseError(
+                            message=f"Failed to initialize inference service: {e}",
+                            response=None,
+                            status_code=http_constants.StatusCodes.BAD_REQUEST
+                        ) from e
         return self._inference_service
 
     @property
diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py
@@ -259,8 +259,6 @@ def __init__( # pylint: disable=too-many-statements
         )
 
         self._inference_service: Optional[_InferenceService] = None
-        if self.aad_credentials:
-            self._inference_service = _InferenceService(self)
 
         self._setup_kwargs: dict[str, Any] = kwargs
         self.session: Optional[_session.Session] = None
@@ -295,7 +293,16 @@ def _set_container_properties_cache(self, container_link: str, properties: Optio
             self.__container_properties_cache[container_link] = {}
 
     def _get_inference_service(self) -> Optional[_InferenceService]:
-        """Get async inference service instance"""
+        """Get async inference service instance, lazily initializing on first access."""
+        if self._inference_service is None and self.aad_credentials:
+            try:
+                self._inference_service = _InferenceService(self)
+            except ValueError as e:
+                raise exceptions.CosmosHttpResponseError(
+                    message=f"Failed to initialize inference service: {e}",
+                    response=None,
+                    status_code=http_constants.StatusCodes.BAD_REQUEST
+                ) from e
         return self._inference_service
 
     @property
diff --git a/sdk/cosmos/azure-cosmos/tests/test_aad_inference_service_lazy_init.py b/sdk/cosmos/azure-cosmos/tests/test_aad_inference_service_lazy_init.py
@@ -0,0 +1,135 @@
+# The MIT License (MIT)
+# Copyright (c) Microsoft Corporation. All rights reserved.
+
+# cspell:ignore reranker
+"""Regression test for lazy initialization of _InferenceService with AAD credentials.
+
+Verifies that constructing a CosmosClient with AAD credentials does not crash
+when AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT is not set.
+"""
+import base64
+import json
+import os
+import time
+import unittest
+from io import StringIO
+
+import pytest
+from azure.core.credentials import AccessToken
+
+import azure.cosmos.cosmos_client as cosmos_client
+import test_config
+from azure.cosmos import DatabaseProxy, ContainerProxy
+
+
+def _remove_padding(encoded_string):
+    while encoded_string.endswith("="):
+        encoded_string = encoded_string[0:len(encoded_string) - 1]
+    return encoded_string
+
+
+def get_test_item(num):
+    test_item = {
+        'pk': 'pk',
+        'id': 'LazyInit_' + str(num),
+        'test_object': True,
+    }
+    return test_item
+
+
+class CosmosEmulatorCredential(object):
+    """Fake AAD credential for the Cosmos emulator."""
+    def get_token(self, *scopes, **kwargs):
+        aad_header_cosmos_emulator = "{\"typ\":\"JWT\",\"alg\":\"RS256\",\"x5t\":\"" \
+                                     "CosmosEmulatorPrimaryMaster\",\"kid\":\"CosmosEmulatorPrimaryMaster\"}"
+        aad_claim_cosmos_emulator_format = {
+            "aud": "https://localhost.localhost",
+            "iss": "https://sts.fake-issuer.net/7b1999a1-dfd7-440e-8204-00170979b984",
+            "iat": int(time.time()), "nbf": int(time.time()),
+            "exp": int(time.time() + 7200), "aio": "", "appid": "localhost",
+            "appidacr": "1", "idp": "https://localhost:8081/",
+            "oid": "96313034-4739-43cb-93cd-74193adbe5b6", "rh": "", "sub": "localhost",
+            "tid": "EmulatorFederation", "uti": "", "ver": "1.0",
+            "scp": "user_impersonation",
+            "groups": ["7ce1d003-4cb3-4879-b7c5-74062a35c66e",
+                       "e99ff30c-c229-4c67-ab29-30a6aebc3e58",
+                       "5549bb62-c77b-4305-bda9-9ec66b85d9e4",
+                       "c44fd685-5c58-452c-aaf7-13ce75184f65",
+                       "be895215-eab5-43b7-9536-9ef8fe130330"]}
+        emulator_key = test_config.TestConfig.masterKey
+        first = _remove_padding(str(base64.urlsafe_b64encode(aad_header_cosmos_emulator.encode("utf-8")), "utf-8"))
+        str_io_obj = StringIO()
+        json.dump(aad_claim_cosmos_emulator_format, str_io_obj)
+        second = _remove_padding(
+            str(base64.urlsafe_b64encode(str(str_io_obj.getvalue()).replace(" ", "").encode("utf-8")), "utf-8"))
+        third = _remove_padding(str(base64.urlsafe_b64encode(emulator_key.encode("utf-8")), "utf-8"))
+        return AccessToken(first + "." + second + "." + third, int(time.time() + 7200))
+
+
+@pytest.mark.cosmosEmulator
+@pytest.mark.cosmosLong
+class TestAADInferenceServiceLazyInit(unittest.TestCase):
+    """Verify AAD client construction succeeds without the semantic reranker env var.
+
+    Before the fix, _InferenceService was eagerly initialized during CosmosClient
+    construction whenever AAD credentials were used, causing a ValueError if
+    AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT was missing.
+    """
+
+    client: cosmos_client.CosmosClient = None
+    database: DatabaseProxy = None
+    container: ContainerProxy = None
+    configs = test_config.TestConfig
+    host = configs.host
+    masterKey = configs.masterKey
+    credential = CosmosEmulatorCredential() if configs.is_emulator else configs.credential
+
+    def setUp(self):
+        """Save the env var state before each test."""
+        self._saved_endpoint = os.environ.get("AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT")
+
+    def tearDown(self):
+        """Ensure the env var is always unset after each test to prevent leakage."""
+        os.environ.pop("AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT", None)
+
+    def test_aad_client_construction_without_inference_endpoint(self):
+        """Constructing a CosmosClient with AAD creds must not raise when
+        AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT is unset."""
+        os.environ.pop("AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT", None)
+
+        client = cosmos_client.CosmosClient(self.host, self.credential)
+        db = client.get_database_client(self.configs.TEST_DATABASE_ID)
+        container = db.get_container_client(self.configs.TEST_SINGLE_PARTITION_CONTAINER_ID)
+
+        item = get_test_item(0)
+        container.create_item(item)
+        read_result = container.read_item(item=item['id'], partition_key='pk')
+        assert read_result['id'] == item['id']
+
+        query_results = list(container.query_items(
+            query='SELECT * FROM c WHERE c.id=@id',
+            parameters=[{"name": "@id", "value": item['id']}],
+            partition_key='pk'
+        ))
+        assert len(query_results) == 1
+
+        container.delete_item(item=item['id'], partition_key='pk')
+
+    def test_aad_client_construction_with_inference_endpoint(self):
+        """Constructing a CosmosClient with AAD creds must also work when
+        AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT IS set."""
+        os.environ["AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT"] = "https://example.com"
+
+        client = cosmos_client.CosmosClient(self.host, self.credential)
+        db = client.get_database_client(self.configs.TEST_DATABASE_ID)
+        container = db.get_container_client(self.configs.TEST_SINGLE_PARTITION_CONTAINER_ID)
+
+        item = get_test_item(1)
+        container.create_item(item)
+        read_result = container.read_item(item=item['id'], partition_key='pk')
+        assert read_result['id'] == item['id']
+        container.delete_item(item=item['id'], partition_key='pk')
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/sdk/cosmos/azure-cosmos/tests/test_aad_inference_service_lazy_init_async.py b/sdk/cosmos/azure-cosmos/tests/test_aad_inference_service_lazy_init_async.py
@@ -0,0 +1,138 @@
+# The MIT License (MIT)
+# Copyright (c) Microsoft Corporation. All rights reserved.
+
+# cspell:ignore reranker
+"""Async regression test for lazy initialization of _InferenceService with AAD credentials.
+
+Verifies that constructing an async CosmosClient with AAD credentials does not crash
+when AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT is not set.
+"""
+import base64
+import json
+import os
+import time
+import unittest
+from io import StringIO
+
+import pytest
+from azure.core.credentials import AccessToken
+
+import test_config
+from azure.cosmos.aio import CosmosClient, DatabaseProxy, ContainerProxy
+
+
+def _remove_padding(encoded_string):
+    while encoded_string.endswith("="):
+        encoded_string = encoded_string[0:len(encoded_string) - 1]
+    return encoded_string
+
+
+def get_test_item(num):
+    test_item = {
+        'pk': 'pk',
+        'id': 'LazyInitAsync_' + str(num),
+        'test_object': True,
+    }
+    return test_item
+
+
+class CosmosEmulatorCredential(object):
+    """Fake async AAD credential for the Cosmos emulator."""
+    async def get_token(self, *scopes, **kwargs):
+        aad_header_cosmos_emulator = "{\"typ\":\"JWT\",\"alg\":\"RS256\",\"x5t\":\"" \
+                                     "CosmosEmulatorPrimaryMaster\",\"kid\":\"CosmosEmulatorPrimaryMaster\"}"
+        aad_claim_cosmos_emulator_format = {
+            "aud": "https://localhost.localhost",
+            "iss": "https://sts.fake-issuer.net/7b1999a1-dfd7-440e-8204-00170979b984",
+            "iat": int(time.time()), "nbf": int(time.time()),
+            "exp": int(time.time() + 7200), "aio": "", "appid": "localhost",
+            "appidacr": "1", "idp": "https://localhost:8081/",
+            "oid": "96313034-4739-43cb-93cd-74193adbe5b6", "rh": "", "sub": "localhost",
+            "tid": "EmulatorFederation", "uti": "", "ver": "1.0",
+            "scp": "user_impersonation",
+            "groups": ["7ce1d003-4cb3-4879-b7c5-74062a35c66e",
+                       "e99ff30c-c229-4c67-ab29-30a6aebc3e58",
+                       "5549bb62-c77b-4305-bda9-9ec66b85d9e4",
+                       "c44fd685-5c58-452c-aaf7-13ce75184f65",
+                       "be895215-eab5-43b7-9536-9ef8fe130330"]}
+        emulator_key = test_config.TestConfig.masterKey
+        first = _remove_padding(str(base64.urlsafe_b64encode(aad_header_cosmos_emulator.encode("utf-8")), "utf-8"))
+        str_io_obj = StringIO()
+        json.dump(aad_claim_cosmos_emulator_format, str_io_obj)
+        second = _remove_padding(
+            str(base64.urlsafe_b64encode(str(str_io_obj.getvalue()).replace(" ", "").encode("utf-8")), "utf-8"))
+        third = _remove_padding(str(base64.urlsafe_b64encode(emulator_key.encode("utf-8")), "utf-8"))
+        return AccessToken(first + "." + second + "." + third, int(time.time() + 7200))
+
+
+@pytest.mark.cosmosEmulator
+@pytest.mark.cosmosLong
+class TestAADInferenceServiceLazyInitAsync(unittest.IsolatedAsyncioTestCase):
+    """Verify async AAD client construction succeeds without the semantic reranker env var.
+
+    Before the fix, _InferenceService was eagerly initialized during CosmosClient
+    construction whenever AAD credentials were used, causing a ValueError if
+    AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT was missing.
+    """
+
+    client: CosmosClient = None
+    database: DatabaseProxy = None
+    container: ContainerProxy = None
+    configs = test_config.TestConfig
+    host = configs.host
+    masterKey = configs.masterKey
+    credential = CosmosEmulatorCredential() if configs.is_emulator else configs.credential_async
+
+    def setUp(self):
+        """Save the env var state before each test."""
+        self._saved_endpoint = os.environ.get("AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT")
+
+    def tearDown(self):
+        """Ensure the env var is always unset after each test to prevent leakage."""
+        os.environ.pop("AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT", None)
+
+    async def test_aad_client_construction_without_inference_endpoint(self):
+        """Constructing an async CosmosClient with AAD creds must not raise when
+        AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT is unset."""
+        os.environ.pop("AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT", None)
+
+        client = CosmosClient(self.host, self.credential)
+        db = client.get_database_client(self.configs.TEST_DATABASE_ID)
+        container = db.get_container_client(self.configs.TEST_SINGLE_PARTITION_CONTAINER_ID)
+
+        item = get_test_item(0)
+        await container.create_item(item)
+        read_result = await container.read_item(item=item['id'], partition_key='pk')
+        assert read_result['id'] == item['id']
+
+        query_results = [
+            doc async for doc in container.query_items(
+                query='SELECT * FROM c WHERE c.id=@id',
+                parameters=[{"name": "@id", "value": item['id']}],
+                partition_key='pk'
+            )
+        ]
+        assert len(query_results) == 1
+
+        await container.delete_item(item=item['id'], partition_key='pk')
+        await client.close()
+
+    async def test_aad_client_construction_with_inference_endpoint(self):
+        """Constructing an async CosmosClient with AAD creds must also work when
+        AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT IS set."""
+        os.environ["AZURE_COSMOS_SEMANTIC_RERANKER_INFERENCE_ENDPOINT"] = "https://placeholder.example.com"
+
+        client = CosmosClient(self.host, self.credential)
+        db = client.get_database_client(self.configs.TEST_DATABASE_ID)
+        container = db.get_container_client(self.configs.TEST_SINGLE_PARTITION_CONTAINER_ID)
+
+        item = get_test_item(1)
+        await container.create_item(item)
+        read_result = await container.read_item(item=item['id'], partition_key='pk')
+        assert read_result['id'] == item['id']
+        await container.delete_item(item=item['id'], partition_key='pk')
+        await client.close()
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/sdk/cosmos/azure-cosmos/tests/test_semantic_reranker_unit.py b/sdk/cosmos/azure-cosmos/tests/test_semantic_reranker_unit.py
