scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 200

Merge commit '6df03f0b61' into litellm_azure_entra_token_followups

Browse source 
# Conflicts:
#	litellm/llms/azure/common_utils.py
This commit is contained in:
mateo-berri 2026-10-19 11:57:25 -07:00
commit 07b45f4cfc
2 changed files with 103 additions and 39 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

25
litellm/llms/azure/common_utils.py
View file

@ -4,6 +4,7 @@ import json
import os
import re
from collections.abc import Awaitable, Callable, Mapping
from contextlib import closing
from functools import lru_cache, partial
from types import MappingProxyType
from typing import Final, Literal, NamedTuple, cast
@ -18,6 +19,7 @@ from litellm._logging import verbose_logger
from litellm.caching.caching import DualCache
from litellm.constants import DEFAULT_MAX_RETRIES
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.llms.openai.common_utils import BaseOpenAILLM, OpenAIAsyncHTTPClient, OpenAIHTTPClient
from litellm.secret_managers.get_azure_ad_token_provider import (
get_azure_ad_token_provider,
@ -287,18 +289,17 @@ def get_azure_ad_token_from_oidc(
if azure_ad_token_access_token is not None:
return azure_ad_token_access_token
client: Final = litellm.module_level_client
req_token: Final = client.post(
f"{azure_authority_host}/{azure_tenant_id}/oauth2/v2.0/token",
data={
"client_id": azure_client_id,
"grant_type": "client_credentials",
"scope": scope,
"client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer",
"client_assertion": oidc_token,
},
)
with closing(HTTPHandler(timeout=litellm.request_timeout)) as client:
req_token: Final = client.post(
f"{azure_authority_host}/{azure_tenant_id}/oauth2/v2.0/token",
data={
"client_id": azure_client_id,
"grant_type": "client_credentials",
"scope": scope,
"client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer",
"client_assertion": oidc_token,
},
)
if req_token.status_code != 200:
raise AzureOpenAIError(

117
tests/unit/llms/azure/test_azure_common_utils.py
View file

@ -2,7 +2,9 @@ import json
import os
import threading
import traceback
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Callable, Final, Optional
from urllib.parse import parse_qs
from unittest.mock import MagicMock, patch
import pytest
@ -14,6 +16,7 @@ from litellm.llms.azure.common_utils import (
_cached_entra_id_token_provider,
get_azure_ad_token,
get_azure_ad_token_from_entra_id,
get_azure_ad_token_from_oidc,
get_azure_ad_token_from_username_password,
)
from litellm.secret_managers.get_azure_ad_token_provider import (
@ -2201,16 +2204,28 @@ async def test_non_v1_azure_token_refresh_runs_off_the_event_loop_thread() -> No
async def test_oidc_token_exchange_runs_off_the_event_loop_and_authorizes_the_request(
monkeypatch: pytest.MonkeyPatch, api_version: str
) -> None:
import asyncio
import httpx
from litellm.llms.custom_httpx.http_handler import HTTPHandler
exchange_started: Final = threading.Event()
loop_ran_during_exchange: Final = threading.Event()
exchanges: Final[list[bool]] = []
exchange_threads: Final[list[int]] = []
class TokenEndpoint(BaseHTTPRequestHandler):
protocol_version = "HTTP/1.1"
def entra_token_endpoint(request: httpx.Request) -> httpx.Response:
exchange_threads.append(threading.get_ident())
assert b"client_assertion=federated-oidc-token" in request.content
return httpx.Response(200, json={"access_token": f"exchanged-{api_version}", "expires_in": 3600})
def do_POST(self) -> None:
form: Final = parse_qs(self.rfile.read(int(self.headers["Content-Length"])).decode())
assert form["client_assertion"] == ["federated-oidc-token"]
exchange_started.set()
exchanges.append(loop_ran_during_exchange.wait(timeout=5))
body: Final = json.dumps({"access_token": f"exchanged-{api_version}", "expires_in": 3600}).encode()
self.send_response(200)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(body)))
self.end_headers()
self.wfile.write(body)
def echo_authorization(request: httpx.Request) -> httpx.Response:
return httpx.Response(
@ -2230,33 +2245,41 @@ async def test_oidc_token_exchange_runs_off_the_event_loop_and_authorizes_the_re
},
)
async def mark_loop_running_once_exchange_starts() -> None:
while not exchange_started.is_set():
await asyncio.sleep(0.01)
loop_ran_during_exchange.set()
server: Final = ThreadingHTTPServer(("127.0.0.1", 0), TokenEndpoint)
threading.Thread(target=server.serve_forever, daemon=True).start()
monkeypatch.setenv("AZURE_AUTHORITY_HOST", f"http://127.0.0.1:{server.server_port}")
monkeypatch.setenv("FEDERATED_OIDC_TOKEN", "federated-oidc-token")
monkeypatch.setattr(
litellm,
"module_level_client",
HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(entra_token_endpoint))),
)
monkeypatch.setattr(
litellm, "aclient_session", httpx.AsyncClient(transport=httpx.MockTransport(echo_authorization))
)
loop_thread: Final = threading.get_ident()
client: Final = BaseAzureLLM().get_azure_openai_client(
api_key=None,
api_base=f"https://oidc-exchange-{api_version}.openai.azure.com",
api_version=api_version,
litellm_params={
"azure_ad_token": "oidc/env/FEDERATED_OIDC_TOKEN",
"client_id": f"oidc-client-{api_version}",
"tenant_id": "oidc-tenant",
},
_is_async=True,
)
response: Final = await client.chat.completions.create(model="gpt", messages=[{"role": "user", "content": "hi"}])
loop_watcher: Final = asyncio.ensure_future(mark_loop_running_once_exchange_starts())
try:
client: Final = BaseAzureLLM().get_azure_openai_client(
api_key=None,
api_base=f"https://oidc-exchange-{api_version}.openai.azure.com",
api_version=api_version,
litellm_params={
"azure_ad_token": "oidc/env/FEDERATED_OIDC_TOKEN",
"client_id": f"oidc-client-{api_version}",
"tenant_id": "oidc-tenant",
},
_is_async=True,
)
response: Final = await client.chat.completions.create(
model="gpt", messages=[{"role": "user", "content": "hi"}]
)
finally:
loop_watcher.cancel()
server.shutdown()
server.server_close()
assert response.choices[0].message.content == f"Bearer exchanged-{api_version}"
assert len(exchange_threads) == 1
assert loop_thread not in exchange_threads
assert exchanges == [True]
@pytest.mark.asyncio
@ -2303,6 +2326,46 @@ async def test_async_azure_ad_token_provider_reaches_azure_as_bearer_token(
assert response.choices[0].message.content == "Bearer async-token"
def test_oidc_token_exchange_closes_its_token_endpoint_connection(monkeypatch: pytest.MonkeyPatch) -> None:
closed_connections: Final = threading.Semaphore(0)
class TokenEndpoint(BaseHTTPRequestHandler):
protocol_version = "HTTP/1.1"
def do_POST(self) -> None:
form: Final = parse_qs(self.rfile.read(int(self.headers["Content-Length"])).decode())
body: Final = json.dumps(
{"access_token": f"minted-for-{form['client_assertion'][0]}", "expires_in": 3600}
).encode()
self.send_response(200)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(body)))
self.end_headers()
self.wfile.write(body)
def finish(self) -> None:
super().finish()
closed_connections.release()
server: Final = ThreadingHTTPServer(("127.0.0.1", 0), TokenEndpoint)
threading.Thread(target=server.serve_forever, daemon=True).start()
monkeypatch.setenv("AZURE_AUTHORITY_HOST", f"http://127.0.0.1:{server.server_port}")
monkeypatch.setenv("AZURE_TEST_OIDC_ASSERTION", "workload-jwt")
try:
token: Final = get_azure_ad_token_from_oidc(
azure_ad_token="os.environ/AZURE_TEST_OIDC_ASSERTION",
azure_client_id="client",
azure_tenant_id="tenant",
)
connection_closed: Final = closed_connections.acquire(timeout=5)
finally:
server.shutdown()
server.server_close()
assert token == "minted-for-workload-jwt"
assert connection_closed
def test_evicting_an_azure_client_built_on_the_callers_session_leaves_it_open(monkeypatch):
"""`initialize_azure_sdk_client` puts `litellm.aclient_session` on the SDK client.