fix(mcp): OAuth 2.1 spec compliance + cookie exp + loopback edge cases

Address codex review P1 + P2 findings:

- BYOK /token now accepts OAuth 2.1 clients that omit redirect_uri
  (draft-15 §4.1.3 dropped the requirement). When the client does
  submit a value, equality is still enforced vs the /authorize record.
  PKCE + client_id binding cover the security role redirect_uri
  played under RFC 6749.
- _user_id_from_session_cookie requires the ``exp`` claim on the UI
  session JWT (PyJWT options={"require": ["exp"]}) so leaked cookies
  have a bounded lifetime.
- validate_loopback_redirect_uri rejects URIs with a fragment
  (RFC 6749 §3.1.2) and catches malformed-URI ValueError so
  unparseable input surfaces as 400 invalid_request instead of 500.

litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py

@@ -105,7 +105,15 @@ def _user_id_from_session_cookie(request: Request) -> Optional[str]:
     if not token:
         return None
     try:
-        payload = jwt.decode(token, master_key, algorithms=["HS256"])
+        payload = jwt.decode(
+            token,
+            master_key,
+            algorithms=["HS256"],
+            # Require an expiry claim so a leaked UI session cookie has a
+            # bounded lifetime. PyJWT verifies exp by default when present;
+            # require=["exp"] additionally rejects tokens that omit it.
+            options={"require": ["exp"]},
+        )
     except jwt.InvalidTokenError:
         return None
     if payload.get("type") == "byok_session":
@@ -791,10 +799,18 @@ async def byok_token(
     if not _verify_pkce(code_verifier, record["code_challenge"]):
         return _oauth_token_error("invalid_grant")
 
-    # RFC 6749 §4.1.3 / OAuth 2.1 §4.1.3: if redirect_uri was sent with the
-    # authorization request, the token request MUST include the identical
-    # value. Enforce exact match.
-    if record.get("redirect_uri") and redirect_uri != record["redirect_uri"]:
+    # RFC 6749 §4.1.3: if redirect_uri was sent with the authorization
+    # request, the token request MUST include the identical value.
+    # OAuth 2.1 draft-15 §4.1.3 drops this requirement — strict OAuth 2.1
+    # clients will omit it. Enforce equality ONLY when the client
+    # actually submitted a value, so we stay RFC 6749-backward-compatible
+    # without breaking OAuth 2.1 clients. PKCE + client_id binding
+    # (checked below) cover the security role redirect_uri played.
+    if (
+        record.get("redirect_uri")
+        and redirect_uri
+        and redirect_uri != record["redirect_uri"]
+    ):
         return _oauth_token_error("invalid_grant")
 
     # RFC 6749 §4.1.3: if the client was identified at /authorize, the

litellm/proxy/_experimental/mcp_server/oauth_utils.py

@@ -24,9 +24,17 @@ def validate_loopback_redirect_uri(redirect_uri: str) -> None:
     ``"127.0.0.1"`` alone would miss ``127.0.0.2`` and the full-form
     IPv6 loopback ``0:0:0:0:0:0:0:1``.
     """
-    parsed = urlparse(redirect_uri)
+    try:
+        parsed = urlparse(redirect_uri)
+    except ValueError:
+        raise HTTPException(status_code=400, detail="invalid_request")
     if parsed.scheme not in ("http", "https"):
         raise HTTPException(status_code=400, detail="invalid_request")
+    # Fragments are not allowed in OAuth redirect URIs (RFC 6749 §3.1.2)
+    # — rejecting them prevents a ``http://127.0.0.1/cb#frag?code=...``
+    # from silently eating the authorization code.
+    if parsed.fragment:
+        raise HTTPException(status_code=400, detail="invalid_request")
     host = (parsed.hostname or "").lower()
     if host == "localhost":
         return
@@ -34,5 +42,7 @@ def validate_loopback_redirect_uri(redirect_uri: str) -> None:
         if ip_address(host).is_loopback:
             return
     except ValueError:
+        # Unparseable host (malformed IPv6, etc.) — treat as invalid,
+        # don't let it bubble up as a 500.
         pass
     raise HTTPException(status_code=400, detail="invalid_request")

tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py

@@ -708,7 +708,11 @@ def test_authorize_post_accepts_ui_session_cookie(unauthenticated_client):
 
     with patch("litellm.proxy.proxy_server.master_key", "test-master-key"):
         cookie_jwt = _jwt.encode(
-            {"user_id": "browser-user-42", "login_method": "sso"},
+            {
+                "user_id": "browser-user-42",
+                "login_method": "sso",
+                "exp": int(time.time()) + 3600,
+            },
             "test-master-key",
             algorithm="HS256",
         )
@@ -998,3 +1002,107 @@ async def test_token_endpoint_missing_master_key_preserves_code_and_db():
     assert code in _byok_auth_codes
     # Credential never written — no inconsistent DB state.
     mock_store.assert_not_awaited()
+
+
+def test_authorize_post_rejects_cookie_without_exp(unauthenticated_client):
+    """Defense-in-depth: UI session cookies must carry an ``exp`` claim
+    so a leaked cookie has a bounded lifetime. A master-key-signed JWT
+    without ``exp`` is rejected at decode time (PyJWT
+    ``options={"require": ["exp"]}``)."""
+    import jwt as _jwt
+
+    with patch("litellm.proxy.proxy_server.master_key", "real-master-key"):
+        no_exp = _jwt.encode(
+            {"user_id": "u", "login_method": "sso"},
+            "real-master-key",
+            algorithm="HS256",
+        )
+        resp = _authorize_post_with_cookie(unauthenticated_client, no_exp, api_key="k")
+    assert resp.status_code == 401
+
+
+def test_authorize_post_rejects_expired_cookie(unauthenticated_client):
+    """An expired UI session cookie is rejected, not accepted as valid."""
+    import jwt as _jwt
+
+    with patch("litellm.proxy.proxy_server.master_key", "real-master-key"):
+        expired = _jwt.encode(
+            {
+                "user_id": "u",
+                "login_method": "sso",
+                "exp": int(time.time()) - 60,  # 1 minute ago
+            },
+            "real-master-key",
+            algorithm="HS256",
+        )
+        resp = _authorize_post_with_cookie(unauthenticated_client, expired, api_key="k")
+    assert resp.status_code == 401
+
+
+@pytest.mark.asyncio
+async def test_token_endpoint_accepts_oauth21_client_omitting_redirect_uri():
+    """OAuth 2.1 draft-15 §4.1.3 dropped the redirect_uri requirement at
+    the token endpoint. A strict OAuth 2.1 client will omit the value —
+    LiteLLM must accept that and rely on PKCE + client_id binding for
+    the security role redirect_uri played under RFC 6749.
+
+    Enforcement still fires when the client DOES submit a value that
+    disagrees with the record (see test_token_endpoint_rejects_redirect_uri_mismatch).
+    """
+    from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import (
+        byok_token,
+    )
+
+    verifier = "verifier_for_oauth21_no_redirect_uri_omit_ok!"
+    challenge = _make_challenge(verifier)
+    code = str(uuid.uuid4())
+    _byok_auth_codes[code] = {
+        "api_key": "k",
+        "server_id": "sid",
+        "code_challenge": challenge,
+        "redirect_uri": "http://127.0.0.1:3000/cb",
+        "client_id": "claude-desktop",
+        "user_id": "u",
+        "expires_at": time.time() + 60,
+    }
+
+    with (
+        patch(
+            "litellm.proxy._experimental.mcp_server.byok_oauth_endpoints.store_user_credential",
+            new=AsyncMock(),
+        ),
+        patch("litellm.proxy.proxy_server.master_key", "test-master"),
+        patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
+    ):
+        result = await byok_token(
+            request=MagicMock(),
+            grant_type="authorization_code",
+            code=code,
+            redirect_uri="",  # OAuth 2.1 client omits it
+            code_verifier=verifier,
+            client_id="claude-desktop",
+        )
+    assert result.status_code == 200
+
+
+def test_validate_loopback_redirect_uri_rejects_fragment():
+    from litellm.proxy._experimental.mcp_server.oauth_utils import (
+        validate_loopback_redirect_uri,
+    )
+
+    with pytest.raises(HTTPException) as exc:
+        validate_loopback_redirect_uri("http://127.0.0.1:3000/cb#code=1")
+    assert exc.value.status_code == 400
+
+
+def test_validate_loopback_redirect_uri_rejects_malformed_cleanly():
+    """Malformed / unparseable URIs should surface as 400 invalid_request,
+    not a 500 from an unhandled exception inside ip_address()."""
+    from litellm.proxy._experimental.mcp_server.oauth_utils import (
+        validate_loopback_redirect_uri,
+    )
+
+    # Netloc that parses but whose host is neither "localhost" nor a valid IP.
+    with pytest.raises(HTTPException) as exc:
+        validate_loopback_redirect_uri("http://[not-an-ip]/cb")
+    assert exc.value.status_code == 400