From 58ead7f6531482c028835ad6ca35279e4f2c8ee7 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Tue, 4 Aug 2026 16:06:41 -0700 Subject: [PATCH] fix(azure_storage): honor AZURE_STORAGE_ENDPOINT_SUFFIX for sovereign clouds (#35806) The azure_storage logging callback and the azure blob files backend built every storage URL against the hardcoded commercial host, so an Azure Government account was unreachable with no way to override it. Read AZURE_STORAGE_ENDPOINT_SUFFIX (default core.windows.net) once in AzureBlobStorageLogger and derive the Data Lake and Blob hosts from it, so all seven previously hardcoded sites follow the configured cloud. Parse stored blob URLs with urlparse instead of matching the commercial host, so URLs persisted before the suffix was configured still resolve, and pin the resulting host-validation boundary with tests. --- litellm/constants.py | 1 + .../azure_storage/azure_storage.py | 21 +- .../files/azure_blob_storage_backend.py | 27 ++- .../azure_storage/test_azure_storage.py | 80 +++++++ .../llms/base_llm/files/__init__.py | 0 .../files/test_azure_blob_storage_backend.py | 226 ++++++++++++++++++ 6 files changed, 339 insertions(+), 16 deletions(-) create mode 100644 tests/test_litellm/llms/base_llm/files/__init__.py create mode 100644 tests/test_litellm/llms/base_llm/files/test_azure_blob_storage_backend.py diff --git a/litellm/constants.py b/litellm/constants.py index 73f0d1e160e..0c7316455d6 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1305,6 +1305,7 @@ RESPONSE_FORMAT_TOOL_NAME = "json_tool_call" # default tool name used when conv ########################### Logging Callback Constants ########################### AZURE_STORAGE_MSFT_VERSION: Final = "2019-07-07" +AZURE_STORAGE_DEFAULT_ENDPOINT_SUFFIX: Final = "core.windows.net" PROMETHEUS_BUDGET_METRICS_REFRESH_INTERVAL_MINUTES: Final = int( os.getenv("PROMETHEUS_BUDGET_METRICS_REFRESH_INTERVAL_MINUTES", 5) ) diff --git a/litellm/integrations/azure_storage/azure_storage.py b/litellm/integrations/azure_storage/azure_storage.py index d2181dbbb38..cb7175691df 100644 --- a/litellm/integrations/azure_storage/azure_storage.py +++ b/litellm/integrations/azure_storage/azure_storage.py @@ -6,7 +6,11 @@ from typing import Final from litellm._logging import verbose_logger from litellm._uuid import uuid -from litellm.constants import _DEFAULT_TTL_FOR_HTTPX_CLIENTS, AZURE_STORAGE_MSFT_VERSION +from litellm.constants import ( + _DEFAULT_TTL_FOR_HTTPX_CLIENTS, + AZURE_STORAGE_DEFAULT_ENDPOINT_SUFFIX, + AZURE_STORAGE_MSFT_VERSION, +) from litellm.integrations.custom_batch_logger import CustomBatchLogger from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.azure.common_utils import get_azure_ad_token_from_entra_id @@ -41,6 +45,9 @@ class AzureBlobStorageLogger(CustomBatchLogger): if not _azure_storage_file_system: raise ValueError("Missing required environment variable: AZURE_STORAGE_FILE_SYSTEM") self.azure_storage_file_system: str = _azure_storage_file_system + self.azure_storage_endpoint_suffix: str = ( + os.getenv("AZURE_STORAGE_ENDPOINT_SUFFIX") or AZURE_STORAGE_DEFAULT_ENDPOINT_SUFFIX + ) self._service_client = None # Time that the azure service client expires, in order to reset the connection pool and keep it fresh self._service_client_timeout: float | None = None @@ -59,6 +66,14 @@ class AzureBlobStorageLogger(CustomBatchLogger): ) raise e + @property + def azure_storage_dfs_endpoint(self) -> str: + return f"https://{self.azure_storage_account_name}.dfs.{self.azure_storage_endpoint_suffix}" + + @property + def azure_storage_blob_endpoint(self) -> str: + return f"https://{self.azure_storage_account_name}.blob.{self.azure_storage_endpoint_suffix}" + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): """ Async Log success events to Azure Blob Storage @@ -144,7 +159,7 @@ class AzureBlobStorageLogger(CustomBatchLogger): json_payload: Final = safe_dumps(payload) + "\n" # Add newline for each log entry payload_bytes: Final = json_payload.encode("utf-8") filename: Final = f"{payload.get('id') or str(uuid.uuid4())}.json" - base_url = f"https://{self.azure_storage_account_name}.dfs.core.windows.net/{self.azure_storage_file_system}/{filename}" + base_url = f"{self.azure_storage_dfs_endpoint}/{self.azure_storage_file_system}/{filename}" # Execute the 3-step upload process await self._create_file(async_client, base_url) @@ -296,7 +311,7 @@ class AzureBlobStorageLogger(CustomBatchLogger): self._service_client = None if not self._service_client: self._service_client = DataLakeServiceClient( - account_url=f"https://{self.azure_storage_account_name}.dfs.core.windows.net", + account_url=self.azure_storage_dfs_endpoint, credential=self.azure_storage_account_key, ) self._service_client_timeout = time.time() + _DEFAULT_TTL_FOR_HTTPX_CLIENTS diff --git a/litellm/llms/base_llm/files/azure_blob_storage_backend.py b/litellm/llms/base_llm/files/azure_blob_storage_backend.py index b8d27b71636..e22cf528856 100644 --- a/litellm/llms/base_llm/files/azure_blob_storage_backend.py +++ b/litellm/llms/base_llm/files/azure_blob_storage_backend.py @@ -8,7 +8,7 @@ to reuse all authentication and Azure Storage operations. import time from typing import Final -from urllib.parse import quote +from urllib.parse import quote, urlparse from litellm._logging import verbose_logger from litellm._uuid import uuid @@ -47,6 +47,8 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): - AZURE_STORAGE_TENANT_ID (optional, if using Azure AD) - AZURE_STORAGE_CLIENT_ID (optional, if using Azure AD) - AZURE_STORAGE_CLIENT_SECRET (optional, if using Azure AD) + - AZURE_STORAGE_ENDPOINT_SUFFIX (optional, defaults to core.windows.net; set to + core.usgovcloudapi.net or another sovereign-cloud suffix as needed) Note: We skip periodic_flush since we're not using this as a logger. """ @@ -103,7 +105,7 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): """ Upload a file to Azure Blob Storage. - Returns the blob URL in format: https://{account}.blob.core.windows.net/{container}/{path} + Returns the blob URL in format: https://{account}.blob.{endpoint_suffix}/{container}/{path} """ try: # Generate file name @@ -172,7 +174,7 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): await file_client.flush_data(position=len(file_content), offset=0) # Return blob URL (not DFS URL) - blob_url = f"https://{self.azure_storage_account_name}.blob.core.windows.net/{self.azure_storage_file_system}/{full_path}" + blob_url = f"{self.azure_storage_blob_endpoint}/{self.azure_storage_file_system}/{full_path}" return blob_url async def _upload_file_with_azure_ad(self, file_content: bytes, full_path: str) -> str: @@ -188,7 +190,7 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): async_client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) # Use DFS endpoint for upload - base_url = f"https://{self.azure_storage_account_name}.dfs.core.windows.net/{self.azure_storage_file_system}/{full_path}" + base_url = f"{self.azure_storage_dfs_endpoint}/{self.azure_storage_file_system}/{full_path}" # Execute 3-step upload process: create, append, flush # Reuse the logger's helper methods @@ -198,7 +200,7 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): await self._flush_data(async_client, base_url, len(file_content)) # Return blob URL (not DFS URL) - blob_url = f"https://{self.azure_storage_account_name}.blob.core.windows.net/{self.azure_storage_file_system}/{full_path}" + blob_url = f"{self.azure_storage_blob_endpoint}/{self.azure_storage_file_system}/{full_path}" return blob_url async def _append_data_bytes(self, client, base_url: str, file_content: bytes): @@ -222,23 +224,22 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): Download a file from Azure Blob Storage. Args: - storage_url: Blob URL in format: https://{account}.blob.core.windows.net/{container}/{path} + storage_url: Blob URL in format: https://{account}.blob.{endpoint_suffix}/{container}/{path} Returns: bytes: File content """ try: # Parse blob URL to extract path - # URL format: https://{account}.blob.core.windows.net/{container}/{path} - if ".blob.core.windows.net/" not in storage_url: + # URL format: https://{account}.blob.{endpoint_suffix}/{container}/{path} + parsed_url: Final = urlparse(storage_url) + if ".blob." not in (parsed_url.hostname or ""): raise ValueError(f"Invalid Azure Blob Storage URL: {storage_url}") # Extract path after container name - container_and_path: Final = storage_url.split(".blob.core.windows.net/", 1)[1] - path_parts: Final = container_and_path.split("/", 1) - if len(path_parts) < 2: + _, _, file_path = parsed_url.path.lstrip("/").partition("/") + if not file_path: raise ValueError(f"Invalid Azure Blob Storage URL format: {storage_url}") - file_path: Final = path_parts[1] # Path after container name if self.azure_storage_account_key: # Use Azure SDK (reuse logger's service client) @@ -279,7 +280,7 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): async_client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) # Use blob endpoint for download (simpler than DFS) - blob_url = f"https://{self.azure_storage_account_name}.blob.core.windows.net/{self.azure_storage_file_system}/{file_path}" + blob_url = f"{self.azure_storage_blob_endpoint}/{self.azure_storage_file_system}/{file_path}" headers: Final = { "x-ms-version": AZURE_STORAGE_MSFT_VERSION, diff --git a/tests/test_litellm/integrations/azure_storage/test_azure_storage.py b/tests/test_litellm/integrations/azure_storage/test_azure_storage.py index d6c9d7a5c92..5d7c55e81af 100644 --- a/tests/test_litellm/integrations/azure_storage/test_azure_storage.py +++ b/tests/test_litellm/integrations/azure_storage/test_azure_storage.py @@ -20,6 +20,13 @@ def mock_env_vars(monkeypatch): monkeypatch.setenv("AZURE_STORAGE_TENANT_ID", "test-tenant-id") monkeypatch.setenv("AZURE_STORAGE_CLIENT_ID", "test-client-id") monkeypatch.setenv("AZURE_STORAGE_CLIENT_SECRET", "test-client-secret") + monkeypatch.delenv("AZURE_STORAGE_ENDPOINT_SUFFIX", raising=False) + + +@pytest.fixture +def mock_gov_env_vars(mock_env_vars, monkeypatch): + """Point the logger at an Azure Government storage account""" + monkeypatch.setenv("AZURE_STORAGE_ENDPOINT_SUFFIX", "core.usgovcloudapi.net") @pytest.mark.asyncio @@ -99,3 +106,76 @@ async def test_async_upload_payload_to_azure_blob_storage(mock_env_vars): # Verify raise_for_status was called on all responses assert mock_response.raise_for_status.call_count == 3 + + +@pytest.mark.asyncio +async def test_async_upload_payload_uses_configured_endpoint_suffix(mock_gov_env_vars): + """ + AZURE_STORAGE_ENDPOINT_SUFFIX must reach the Entra-ID REST upload path so a + sovereign-cloud account is addressed instead of the commercial dfs host. + """ + with patch( + "litellm.integrations.azure_storage.azure_storage.get_async_httpx_client" + ) as mock_get_client: + mock_http_client = AsyncMock() + mock_response = MagicMock() + mock_http_client.put.return_value = mock_response + mock_http_client.patch.return_value = mock_response + mock_get_client.return_value = mock_http_client + + logger = AzureBlobStorageLogger() + logger.azure_auth_token = "mock-azure-ad-token" + logger.token_expiry = None + + test_payload: StandardLoggingPayload = {"id": "gov-log-id"} + + await logger.async_upload_payload_to_azure_blob_storage(test_payload) + + expected_base_url = ( + "https://test-account.dfs.core.usgovcloudapi.net/test-container/gov-log-id.json" + ) + assert mock_http_client.put.call_args[0][0] == f"{expected_base_url}?resource=file" + assert ( + mock_http_client.patch.call_args_list[0][0][0] + == f"{expected_base_url}?action=append&position=0" + ) + assert mock_http_client.patch.call_args_list[1][0][0].startswith( + f"{expected_base_url}?action=flush" + ) + + +@pytest.mark.asyncio +async def test_service_client_uses_configured_endpoint_suffix(mock_gov_env_vars): + """ + The account key path builds its own account_url; the Azure SDK derives the blob + host from it, so the suffix has to be applied here too. + """ + fake_aio_module = MagicMock() + + with patch.dict( + sys.modules, {"azure.storage.filedatalake.aio": fake_aio_module} + ): + logger = AzureBlobStorageLogger() + await logger.get_service_client() + + assert ( + fake_aio_module.DataLakeServiceClient.call_args.kwargs["account_url"] + == "https://test-account.dfs.core.usgovcloudapi.net" + ) + + +@pytest.mark.asyncio +async def test_service_client_defaults_to_commercial_endpoint(mock_env_vars): + """Unset AZURE_STORAGE_ENDPOINT_SUFFIX keeps the pre-existing commercial host""" + fake_aio_module = MagicMock() + + with patch.dict( + sys.modules, {"azure.storage.filedatalake.aio": fake_aio_module} + ): + logger = AzureBlobStorageLogger() + await logger.get_service_client() + + assert ( + fake_aio_module.DataLakeServiceClient.call_args.kwargs["account_url"] + == "https://test-account.dfs.core.windows.net" + ) diff --git a/tests/test_litellm/llms/base_llm/files/__init__.py b/tests/test_litellm/llms/base_llm/files/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/base_llm/files/test_azure_blob_storage_backend.py b/tests/test_litellm/llms/base_llm/files/test_azure_blob_storage_backend.py new file mode 100644 index 00000000000..b924ea8f93f --- /dev/null +++ b/tests/test_litellm/llms/base_llm/files/test_azure_blob_storage_backend.py @@ -0,0 +1,226 @@ +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from litellm.llms.base_llm.files.azure_blob_storage_backend import ( + AzureBlobStorageBackend, +) + +GOV_SUFFIX = "core.usgovcloudapi.net" + + +@pytest.fixture +def mock_env_vars(monkeypatch): + """Azure AD (no account key) configuration for the files backend""" + monkeypatch.setenv("AZURE_STORAGE_ACCOUNT_NAME", "test-account") + monkeypatch.setenv("AZURE_STORAGE_FILE_SYSTEM", "test-container") + monkeypatch.setenv("AZURE_STORAGE_TENANT_ID", "test-tenant-id") + monkeypatch.setenv("AZURE_STORAGE_CLIENT_ID", "test-client-id") + monkeypatch.setenv("AZURE_STORAGE_CLIENT_SECRET", "test-client-secret") + monkeypatch.delenv("AZURE_STORAGE_ACCOUNT_KEY", raising=False) + monkeypatch.delenv("AZURE_STORAGE_ENDPOINT_SUFFIX", raising=False) + + +@pytest.fixture +def mock_gov_env_vars(mock_env_vars, monkeypatch): + monkeypatch.setenv("AZURE_STORAGE_ENDPOINT_SUFFIX", GOV_SUFFIX) + + +def _make_backend() -> AzureBlobStorageBackend: + backend = AzureBlobStorageBackend() + backend.azure_auth_token = "mock-azure-ad-token" + backend.token_expiry = None + return backend + + +def _mock_upload_client() -> AsyncMock: + client = AsyncMock() + response = MagicMock() + client.put = AsyncMock(return_value=response) + client.patch = AsyncMock(return_value=response) + return client + + +@pytest.mark.parametrize( + "env_fixture, expected_suffix", + [("mock_env_vars", "core.windows.net"), ("mock_gov_env_vars", GOV_SUFFIX)], +) +@pytest.mark.asyncio +async def test_upload_file_with_azure_ad_honors_endpoint_suffix(request, env_fixture, expected_suffix): + """ + The REST upload targets the dfs host and the returned handle is a blob URL, so both + have to follow AZURE_STORAGE_ENDPOINT_SUFFIX or a sovereign-cloud account is unreachable. + """ + request.getfixturevalue(env_fixture) + client = _mock_upload_client() + + with patch( + "litellm.llms.custom_httpx.http_handler.get_async_httpx_client", + return_value=client, + ): + backend = _make_backend() + storage_url = await backend.upload_file( + file_content=b"hello", + filename="report.json", + content_type="application/json", + path_prefix="logs", + file_naming_strategy="original_filename", + ) + + expected_dfs = f"https://test-account.dfs.{expected_suffix}/test-container/logs/report.json" + assert client.put.call_args[0][0] == f"{expected_dfs}?resource=file" + assert client.patch.call_args_list[0][0][0] == f"{expected_dfs}?action=append&position=0" + assert storage_url == f"https://test-account.blob.{expected_suffix}/test-container/logs/report.json" + + +@pytest.mark.parametrize( + "env_fixture, expected_suffix", + [("mock_env_vars", "core.windows.net"), ("mock_gov_env_vars", GOV_SUFFIX)], +) +@pytest.mark.asyncio +async def test_download_file_honors_endpoint_suffix(request, env_fixture, expected_suffix): + """ + download_file both validates and splits the stored blob URL on the host, so a + sovereign-cloud URL must parse and round-trip back to the same host. + """ + request.getfixturevalue(env_fixture) + response = MagicMock() + response.content = b"file-bytes" + client = AsyncMock() + client.get = AsyncMock(return_value=response) + + storage_url = f"https://test-account.blob.{expected_suffix}/test-container/logs/report.json" + + with patch( + "litellm.llms.custom_httpx.http_handler.get_async_httpx_client", + return_value=client, + ): + backend = _make_backend() + content = await backend.download_file(storage_url) + + assert content == b"file-bytes" + assert client.get.call_args[0][0] == storage_url + + +@pytest.mark.asyncio +async def test_download_file_accepts_url_persisted_before_the_suffix_was_set(mock_gov_env_vars): + """ + storage_url is persisted in the managed files table while the suffix is process config, + so rows written before the suffix was configured must still resolve. Only the path after + the container is taken from the stored URL; the host comes from the current config. + """ + response = MagicMock() + response.content = b"file-bytes" + client = AsyncMock() + client.get = AsyncMock(return_value=response) + + with patch( + "litellm.llms.custom_httpx.http_handler.get_async_httpx_client", + return_value=client, + ): + backend = _make_backend() + content = await backend.download_file( + "https://old-account.blob.core.windows.net/old-container/logs/report.json" + ) + + assert content == b"file-bytes" + assert ( + client.get.call_args[0][0] + == f"https://test-account.blob.{GOV_SUFFIX}/test-container/logs/report.json" + ) + + +@pytest.mark.parametrize( + "storage_url", + [ + "https://example-bucket.s3.amazonaws.com/container/report.json", + "https://example.com/download?u=.blob.core.windows.net/container/report.json", + "mygovacct.blob.core.windows.net/container/report.json", + ], + ids=["other-provider", "blob-host-only-in-query", "no-scheme"], +) +@pytest.mark.asyncio +async def test_download_file_rejects_url_whose_host_is_not_an_azure_blob_host(mock_env_vars, storage_url): + """ + The host is checked on the parsed hostname, so a blob host appearing anywhere else in the + string no longer passes. No first-party producer emits these, and rejecting beats issuing a + request built from a mis-split path. + """ + client = AsyncMock() + + with patch( + "litellm.llms.custom_httpx.http_handler.get_async_httpx_client", + return_value=client, + ): + backend = _make_backend() + + with pytest.raises(ValueError, match="Invalid Azure Blob Storage URL"): + await backend.download_file(storage_url) + + client.get.assert_not_called() + + +@pytest.mark.asyncio +async def test_download_file_drops_query_string_from_the_stored_url(mock_env_vars): + """A query string on the stored URL is not part of the blob path and must not reach the request""" + response = MagicMock() + response.content = b"file-bytes" + client = AsyncMock() + client.get = AsyncMock(return_value=response) + + with patch( + "litellm.llms.custom_httpx.http_handler.get_async_httpx_client", + return_value=client, + ): + backend = _make_backend() + await backend.download_file( + "https://test-account.blob.core.windows.net/test-container/logs/report.json?sig=redacted&se=2026" + ) + + assert ( + client.get.call_args[0][0] + == "https://test-account.blob.core.windows.net/test-container/logs/report.json" + ) + + +@pytest.mark.parametrize( + "env_fixture, expected_suffix", + [("mock_env_vars", "core.windows.net"), ("mock_gov_env_vars", GOV_SUFFIX)], +) +@pytest.mark.asyncio +async def test_upload_file_with_account_key_honors_endpoint_suffix(request, env_fixture, expected_suffix, monkeypatch): + """The account key path returns its own blob URL, built independently of the REST path""" + request.getfixturevalue(env_fixture) + monkeypatch.setenv("AZURE_STORAGE_ACCOUNT_KEY", "dGVzdC1rZXk=") + + file_client = MagicMock() + file_client.create_file = AsyncMock() + file_client.append_data = AsyncMock() + file_client.flush_data = AsyncMock() + + directory_client = MagicMock() + directory_client.exists = AsyncMock(return_value=True) + directory_client.get_file_client = MagicMock(return_value=file_client) + + file_system_client = MagicMock() + file_system_client.exists = AsyncMock(return_value=True) + file_system_client.get_directory_client = MagicMock(return_value=directory_client) + + service_client = MagicMock() + service_client.get_file_system_client = MagicMock(return_value=file_system_client) + + fake_aio_module = MagicMock() + fake_aio_module.DataLakeServiceClient = MagicMock(return_value=service_client) + + with patch.dict(sys.modules, {"azure.storage.filedatalake.aio": fake_aio_module}): + backend = AzureBlobStorageBackend() + storage_url = await backend.upload_file( + file_content=b"hello", + filename="report.json", + content_type="application/json", + path_prefix="logs", + file_naming_strategy="original_filename", + ) + + assert storage_url == f"https://test-account.blob.{expected_suffix}/test-container/logs/report.json"