Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_exchange_provider.py: 22%

1"""Composition root for the v2-native token_exchange (OBO) exchanger. 

2 

3Wires the pure ``OboTokenExchanger`` to its runtime edges: the real httpx POST against the IdP and 

4the configured cache sizing/TTL constants. ``build_token_exchanger`` is built once at egress 

5construction and reused, so the in-process exchanged-token cache survives across requests. Unlike the 

6per-user store, nothing here reads a runtime global at build time (the httpx client is acquired per 

7call), so it needs no lazy wrapper. 

8""" 

9 

10from __future__ import annotations 

11 

12from typing import Final 

13 

14import httpx 

15 

16from litellm._logging import verbose_logger 

17from litellm.constants import ( 

18 MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL, 

19 MCP_OAUTH2_TOKEN_CACHE_MIN_TTL, 

20 MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS, 

21 MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE, 

22) 

23from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( 

24 InMemoryTokenCacheBackend, 

25) 

26from litellm.proxy._experimental.mcp_server.outbound_credentials.token_exchanger import ( 

27 OboTokenExchanger, 

28 SubjectTokenRejected, 

29 TokenExchangeClientError, 

30) 

31 

32# RFC 6749 5.2 error codes that mean the gateway's own request/credentials are wrong (not the 

33# caller's subject token), so they surface as a 500 the caller can't fix by re-authenticating. 

34_GATEWAY_FAULT_OAUTH_ERRORS: Final = frozenset( 

35 {"invalid_client", "unauthorized_client", "unsupported_grant_type", "invalid_target", "invalid_scope"} 

36) 

37 

38 

39def _oauth_error_fields(response: httpx.Response) -> tuple[str | None, str | None]: 

40 """Read the RFC 6749 5.2 ``error`` code and the IdP's step-up ``claims`` blob from a 

41 token-endpoint error body, as ``(error, claims)`` with None for whatever is absent. 

42 

43 ``claims`` is the Entra Conditional Access / CAE challenge (a JSON string the client must 

44 replay to the IdP to satisfy the step-up); it is the caller's own requirement, not an IdP 

45 internal, so it may travel to the caller. The ``error_description`` is deliberately not read: 

46 it can carry IdP internals and must never reach the caller. 

47 """ 

48 try: 

49 body: Final[object] = response.json() 

50 except Exception: # noqa: BLE001 

51 return None, None 

52 if not isinstance(body, dict): 

53 return None, None 

54 code: Final = body.get("error") 

55 claims: Final = body.get("claims") 

56 return ( 

57 code if isinstance(code, str) else None, 

58 claims if isinstance(claims, str) and claims else None, 

59 ) 

60 

61 

62async def _post_exchange_endpoint( 

63 url: str, form: dict[str, str], client_auth_headers: dict[str, str] 

64) -> dict[str, object] | None: 

65 from litellm.llms.custom_httpx.http_handler import ( # noqa: PLC0415 

66 get_async_httpx_client, # pyright: ignore 

67 ) 

68 from litellm.types.llms.custom_http import httpxSpecialProvider # noqa: PLC0415 

69 

70 # litellm's httpx handler and httpx.Response are only partially typed; the IdP returns a JSON 

71 # object and the exchanger validates each field, so the untyped boundary is contained here. 

72 # A 4xx is the IdP rejecting the subject (non-retryable -> 401 via SubjectTokenRejected); any 

73 # other failure is a miss (-> None -> upstream_unavailable -> 503), matching v1's fail-closed. 

74 headers: Final = {"Accept": "application/json", **client_auth_headers} 

75 try: 

76 client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP) # pyright: ignore 

77 response: Final = await client.post(url, headers=headers, data=form) # pyright: ignore 

78 response.raise_for_status() # pyright: ignore 

79 parsed: Final[object] = response.json() # pyright: ignore 

80 except httpx.HTTPStatusError as status_err: 

81 status_code: Final = status_err.response.status_code 

82 if 400 <= status_code < 500: 

83 oauth_error, claims = _oauth_error_fields(status_err.response) 

84 if oauth_error in _GATEWAY_FAULT_OAUTH_ERRORS: 

85 verbose_logger.warning( 

86 "MCP token exchange rejected as %s (HTTP %d); check the gateway client credentials, " 

87 "audience, and scope for this server", 

88 oauth_error, 

89 status_code, 

90 ) 

91 raise TokenExchangeClientError(oauth_error) from status_err 

92 raise SubjectTokenRejected( 

93 f"IdP rejected the subject token (HTTP {status_code})", 

94 claims=claims, 

95 ) from status_err 

96 verbose_logger.warning("MCP token exchange request failed: %s", status_err) 

97 return None 

98 except Exception as exc: # noqa: BLE001 

99 verbose_logger.warning("MCP token exchange request failed: %s", exc) 

100 return None 

101 if not isinstance(parsed, dict): 

102 # A valid-but-non-object JSON body (list/string/number) would crash the field parsing; map it 

103 # to a miss so it surfaces as a typed upstream_unavailable, not a 500. 

104 verbose_logger.warning("MCP token exchange returned non-object JSON (%s)", type(parsed).__name__) 

105 return None 

106 return parsed # pyright: ignore 

107 

108 

109def build_token_exchanger() -> OboTokenExchanger: 

110 return OboTokenExchanger( 

111 _post_exchange_endpoint, 

112 cache=InMemoryTokenCacheBackend(max_size=MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE), 

113 default_ttl_seconds=MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL, 

114 min_ttl_seconds=MCP_OAUTH2_TOKEN_CACHE_MIN_TTL, 

115 expiry_buffer_seconds=MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS, 

116 )