Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py: 54%
105 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1"""Producer and consumer helpers for the gateway-level DCR session token.
3The aggregate ``/mcp`` front door (``mcp_gateway_dcr``) issues the identity-only session
4tokens defined in :mod:`.session_token`. The gateway token endpoint mints them (producer)
5after SSO sign-in, and at the MCP admission edge the gateway derives the session signing
6key from the proxy ``master_key``, opens the bearer, and admits the request under the
7recovered litellm user (consumer), reloading the live user record and policy before
8anything runs. This module is the pure surface for both sides; the token-endpoint and
9admission wiring live in their respective call sites.
11The signing key is derived with the same memory-hard scrypt construction as
12:func:`~.bridge_credentials.envelope_keys_from_master_key` but under a distinct domain
13label, so session tokens and bridge envelopes never share key material: a token of one
14family is unverifiable in the other by key separation, on top of the distinct issuers,
15prefixes, and claim shapes.
16"""
18import hashlib
19from datetime import datetime
20from functools import lru_cache
21from typing import Final, Literal, TypeAlias
23from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError
25from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import (
26 AsymmetricSessionKeys,
27 OpenedSessionToken,
28 SessionExpired,
29 SessionKeys,
30 SessionPrincipal,
31 SessionRotatedPublicKey,
32 SessionSigningKeys,
33 is_session_refresh_token,
34 is_session_token,
35 open_session_refresh_token,
36 open_session_token,
37)
39_SESSION_SIGNING_KEY_DOMAIN: Final = b"litellm-mcp-gateway:session-signing:"
41# scrypt work factors (RFC 7914), identical to the envelope KDF: memory-hard so a captured
42# session token is not a cheap offline oracle for the master key.
43_SCRYPT_N: Final = 2**15
44_SCRYPT_R: Final = 8
45_SCRYPT_P: Final = 1
46_SCRYPT_MAXMEM: Final = 128 * _SCRYPT_N * _SCRYPT_R * _SCRYPT_P * 2
47_DERIVED_KEY_BYTES: Final = 32
50@lru_cache(maxsize=8)
51def session_keys_from_master_key(master_key: str) -> SessionKeys:
52 """Derive the session signing key from the proxy master key.
54 A memory-hard scrypt KDF (RFC 7914) over a session-specific domain-label salt yields a
55 256-bit subkey from the one secret, so the producer (mint) and consumer (open) agree on
56 the key without persisting any. The domain label differs from both envelope labels in
57 :mod:`.bridge_credentials`, so compromise or misuse of one token family never crosses
58 into the other. The result is cached (the master key is fixed for a process); rotating
59 ``master_key`` invalidates every outstanding session, which is the intended behavior
60 for a signing-key change.
61 """
62 signing: Final = hashlib.scrypt(
63 master_key.encode(),
64 salt=_SESSION_SIGNING_KEY_DOMAIN,
65 n=_SCRYPT_N,
66 r=_SCRYPT_R,
67 p=_SCRYPT_P,
68 maxmem=_SCRYPT_MAXMEM,
69 dklen=_DERIVED_KEY_BYTES,
70 ).hex()
71 return SessionKeys(signing_key=SecretStr(signing))
74class SessionSigningPreviousKey(BaseModel):
75 """One retired key in ``mcp_session_token_signing.previous_public_keys``: its ``kid``
76 and the PEM public half (inline or an ``os.environ/`` reference)."""
78 model_config = ConfigDict(frozen=True, extra="forbid")
79 kid: str = Field(min_length=1)
80 public_key: str = Field(min_length=1)
83class MCPSessionTokenSigningSettings(BaseModel):
84 """The ``general_settings.mcp_session_token_signing`` block: opt-in asymmetric signing
85 for the gateway session tokens. Absent, the gateway keeps the backward-compatible
86 HS256 key derived from ``master_key``. ``private_key`` and each ``public_key`` accept
87 a PEM string inline or an ``os.environ/<NAME>`` (or secret manager) reference."""
89 model_config = ConfigDict(frozen=True, extra="forbid")
90 algorithm: Literal["RS256"]
91 kid: str = Field(min_length=1)
92 private_key: str = Field(min_length=1)
93 previous_public_keys: tuple[SessionSigningPreviousKey, ...] = ()
96class SessionSigningConfigError(BaseModel):
97 """``mcp_session_token_signing`` is present but unusable (bad shape, unresolvable
98 secret reference, or a key that is not a loadable RSA PEM); the caller fails closed
99 with a server error instead of silently falling back to HS256."""
101 model_config = ConfigDict(frozen=True)
102 tag: Literal["session_signing_config_error"] = "session_signing_config_error"
103 detail: str
106def _resolve_key_material(value: str) -> str | None:
107 if not value.startswith("os.environ/"):
108 return value
109 from litellm.secret_managers.main import get_secret_str # noqa: PLC0415 # heavy import kept off the pure path
111 return get_secret_str(value)
114def resolve_session_signing_keys(
115 master_key: str,
116 raw_settings: object | None,
117) -> SessionSigningKeys | SessionSigningConfigError:
118 """Turn the operator's ``mcp_session_token_signing`` setting into signing key material.
120 ``None`` (the setting absent) keeps the backward-compatible HS256 key derived from
121 ``master_key``. A present setting must fully validate into RS256 material; any defect
122 is a ``SessionSigningConfigError`` value so token issuance and admission fail closed
123 rather than minting under a key the operator did not intend.
124 """
125 if raw_settings is None: 125 ↛ 127line 125 didn't jump to line 127 because the condition on line 125 was always true
126 return session_keys_from_master_key(master_key)
127 try:
128 settings: Final = MCPSessionTokenSigningSettings.model_validate(raw_settings)
129 except ValidationError as exc:
130 return SessionSigningConfigError(detail=f"mcp_session_token_signing is malformed: {exc}")
131 private_pem: Final = _resolve_key_material(settings.private_key)
132 if private_pem is None:
133 return SessionSigningConfigError(detail="mcp_session_token_signing.private_key reference did not resolve")
134 resolved_previous: Final = tuple(
135 (previous.kid, _resolve_key_material(previous.public_key)) for previous in settings.previous_public_keys
136 )
137 unresolved: Final = tuple(kid for kid, pem in resolved_previous if pem is None)
138 if unresolved:
139 return SessionSigningConfigError(
140 detail=f"mcp_session_token_signing.previous_public_keys reference did not resolve for kid(s): {', '.join(unresolved)}"
141 )
142 try:
143 return AsymmetricSessionKeys(
144 private_key_pem=SecretStr(private_pem),
145 kid=settings.kid,
146 previous_public_keys=tuple(
147 SessionRotatedPublicKey(kid=kid, public_key_pem=pem)
148 for kid, pem in resolved_previous
149 if pem is not None
150 ),
151 )
152 except ValidationError as exc:
153 return SessionSigningConfigError(
154 detail=f"mcp_session_token_signing keys are not usable RSA PEM material: {exc}"
155 )
158def active_session_signing_keys(master_key: str) -> SessionSigningKeys | SessionSigningConfigError:
159 """Wiring helper for the token endpoint and the admission edge: resolve the signing
160 keys from the live ``general_settings.mcp_session_token_signing`` block, or derive the
161 default HS256 key from ``master_key`` when the block is absent."""
162 from litellm.proxy.proxy_server import general_settings # noqa: PLC0415 # circular import at module load
164 return resolve_session_signing_keys(master_key, general_settings.get("mcp_session_token_signing"))
167class NotSessionBearer(BaseModel):
168 """The bearer is not session-shaped; admission continues on its normal path."""
170 model_config = ConfigDict(frozen=True)
171 tag: Literal["not_session_bearer"] = "not_session_bearer"
174class SessionBearerAdmitted(BaseModel):
175 """A valid session access token: the principal to admit under after a live reload."""
177 model_config = ConfigDict(frozen=True)
178 tag: Literal["admitted"] = "admitted"
179 principal: SessionPrincipal
182class SessionBearerInvalid(BaseModel):
183 """The bearer is session-shaped but must not admit (expired, tampered, wrong key, or a
184 refresh token presented at the tool-call edge); admission fails closed with the
185 ``invalid_token`` challenge rather than falling through to another arm. ``expired``
186 distinguishes a routine expiry (debug-log worthy) from a tampered or foreign token."""
188 model_config = ConfigDict(frozen=True)
189 tag: Literal["invalid"] = "invalid"
190 expired: bool = False
193SessionBearerResult: TypeAlias = NotSessionBearer | SessionBearerAdmitted | SessionBearerInvalid
196def _strip_bearer(value: str) -> str:
197 parts: Final = value.split(None, 1)
198 if len(parts) == 2 and parts[0].lower() == "bearer": 198 ↛ 199line 198 didn't jump to line 199 because the condition on line 198 was never true
199 return parts[1]
200 return value
203def is_session_bearer_shaped(authorization_value: str) -> bool:
204 """Cheap, keyless test that an ``Authorization`` value carries a session token of either
205 kind (optional ``Bearer`` scheme stripped). The admission edge engages the session arm
206 for an access token (to admit) and for a refresh token (to reject it explicitly, since
207 a refresh credential is never usable at the tool-call edge); anything else falls
208 through to normal admission."""
209 candidate: Final = _strip_bearer(authorization_value)
210 return is_session_token(candidate) or is_session_refresh_token(candidate)
213def resolve_session_bearer(
214 authorization_value: str,
215 keys: SessionSigningKeys,
216 now: datetime,
217) -> SessionBearerResult:
218 """Classify an ``Authorization`` value presented at the aggregate MCP edge.
220 Strips an optional ``Bearer`` scheme, then returns ``NotSessionBearer`` for a
221 non-session bearer (normal admission continues), ``SessionBearerAdmitted`` with the
222 recovered principal for a valid access token, and ``SessionBearerInvalid`` for a
223 session-shaped bearer that must not admit. Never raises: total over hostile input via
224 :func:`~.session_token.open_session_token`.
226 A refresh token is ``SessionBearerInvalid`` here: it is a valid gateway credential but
227 only ever presented back to the token endpoint, so admission must fail it closed rather
228 than let it fall through to another arm.
229 """
230 candidate: Final = _strip_bearer(authorization_value)
231 if is_session_refresh_token(candidate):
232 return SessionBearerInvalid()
233 if not is_session_token(candidate):
234 return NotSessionBearer()
235 opened: Final = open_session_token(candidate, keys, now)
236 if isinstance(opened, OpenedSessionToken):
237 return SessionBearerAdmitted(principal=opened.principal)
238 return SessionBearerInvalid(expired=isinstance(opened, SessionExpired))
241class SessionRefreshOpened(BaseModel):
242 """A valid session refresh token presented to the token endpoint: the principal to
243 re-validate and renew under."""
245 model_config = ConfigDict(frozen=True)
246 tag: Literal["opened"] = "opened"
247 principal: SessionPrincipal
248 jti: str
251class SessionRefreshInvalid(BaseModel):
252 """The presented refresh grant is not a valid session refresh token for this client
253 (not refresh-shaped, will not open, or bound to a different ``client_id``); the token
254 endpoint fails the refresh closed."""
256 model_config = ConfigDict(frozen=True)
257 tag: Literal["invalid"] = "invalid"
260SessionRefreshResult: TypeAlias = SessionRefreshOpened | SessionRefreshInvalid
263def open_session_refresh_bearer(
264 refresh_value: str,
265 keys: SessionSigningKeys,
266 now: datetime,
267 expected_client_id: str,
268) -> SessionRefreshResult:
269 """Open a session refresh token presented on a ``refresh_token`` grant.
271 The token-endpoint mirror of :func:`resolve_session_bearer`: strips an optional
272 ``Bearer`` scheme, then returns ``SessionRefreshOpened`` with the recovered principal,
273 or ``SessionRefreshInvalid`` for anything that is not a valid session refresh token
274 issued to ``expected_client_id``. Never raises. The client binding (RFC 6749 section 6)
275 stops a refresh token stolen from one DCR client from being renewed through another;
276 ``client_id`` is not a secret (the caller presents it), so a plain equality check is
277 sufficient and, unlike ``hmac.compare_digest`` on ``str``, does not raise on non-ASCII.
278 """
279 candidate: Final = _strip_bearer(refresh_value)
280 if not is_session_refresh_token(candidate):
281 return SessionRefreshInvalid()
282 opened: Final = open_session_refresh_token(candidate, keys, now)
283 if not isinstance(opened, OpenedSessionToken):
284 return SessionRefreshInvalid()
285 if opened.principal.client_id != expected_client_id:
286 return SessionRefreshInvalid()
287 return SessionRefreshOpened(principal=opened.principal, jti=opened.jti)