yuqi1129 commented on code in PR #12531: URL: https://github.com/apache/gravitino/pull/12531#discussion_r3828584269
########## mcp-server/mcp_server/core/oauth.py: ########## @@ -0,0 +1,215 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""httpx ``auth=`` hook for MCP → Gravitino OAuth2 client-credentials. + +Uses ``httpx-auth`` for fetch and cache on the existing ``httpx.AsyncClient``. +Credentials go in the form body (``client_secret_post``), matching the +Java/Python Gravitino clients. httpx-auth defaults to HTTP Basic. This class +retries once after Gravitino HTTP 401. +""" + +import asyncio +import base64 +import json +import logging +from collections.abc import AsyncGenerator, Generator +from typing import Optional, Union + +import httpx +from httpx_auth import AuthenticationFailed, OAuth2, OAuth2ClientCredentials + +_LOG = logging.getLogger(__name__) + +# Refresh this many seconds before recorded expiry. httpx-auth default is 30. +DEFAULT_REFRESH_SKEW_SECONDS = 60 + +_TokenTuple = Union[tuple[str, str], tuple[str, str, Union[int, str]]] + + +class RefreshableBearerAuth(OAuth2ClientCredentials): + """httpx-auth client-credentials with form POST and one 401 retry.""" + + requires_request_body = True + requires_response_body = True + + def __init__( + self, + *, + token_endpoint: str, + client_id: str, + client_secret: str, + scope: str = "", + refresh_skew_seconds: int = DEFAULT_REFRESH_SKEW_SECONDS, + client: Optional[httpx.Client] = None, + ): + """Build an ``auth=`` hook for the service hop. + + Args: + token_endpoint: Identity-provider token URL. + client_id: OAuth2 client id. + client_secret: OAuth2 client secret. + scope: Optional OAuth2 scope. + refresh_skew_seconds: httpx-auth ``early_expiry``. + client: Optional sync httpx client used only for token POSTs + (tests inject ``MockTransport`` here). + """ + kwargs = {"early_expiry": float(refresh_skew_seconds)} + if scope: + kwargs["scope"] = scope + if client is not None: + kwargs["client"] = client + super().__init__(token_endpoint, client_id, client_secret, **kwargs) + + def invalidate(self) -> None: + """Drop the cached token so the next call fetches a new one.""" + cache = OAuth2.token_cache + # TokenMemoryCache.clear() wipes every client; only drop ours. + with cache._forbid_concurrent_cache_access: # pylint: disable=protected-access + cache.tokens.pop(self.state, None) + + def request_new_token(self) -> _TokenTuple: + """POST ``client_credentials`` with id/secret in the form body.""" + data = self._token_form_data() + client = self.client or httpx.Client() + self._configure_client(client) + try: + response = client.post(self.token_url, data=data) + self._log_token_http_error(response) + response.raise_for_status() + body = response.json() + finally: + if self.client is None: + client.close() + return self._token_tuple(body) + + async def request_new_token_async(self) -> _TokenTuple: + """POST ``client_credentials`` without blocking the event loop.""" + if self.client is not None: + return await asyncio.to_thread(self.request_new_token) + data = self._token_form_data() + async with httpx.AsyncClient() as client: + client.timeout = self.timeout + response = await client.post(self.token_url, data=data) + self._log_token_http_error(response) + response.raise_for_status() + body = response.json() + return self._token_tuple(body) + + def auth_flow( + self, request: httpx.Request + ) -> Generator[httpx.Request, httpx.Response, None]: + """Attach a cached or freshly fetched Bearer; retry once on HTTP 401.""" + self._apply_token(request) + response = yield request + if response.status_code != 401: + return + self.invalidate() + self._apply_token(request) + yield request + + async def async_auth_flow( + self, request: httpx.Request + ) -> AsyncGenerator[httpx.Request, httpx.Response]: + """Attach a Bearer without a blocking IdP POST on the event loop.""" + if self.requires_request_body: + await request.aread() + await self._apply_token_async(request) + response = yield request + if response.status_code != 401: + return + self.invalidate() + await self._apply_token_async(request) + yield request + + def _configure_client(self, client: httpx.Client) -> None: + """Do not send HTTP Basic; id and secret go in the form body.""" + client.timeout = self.timeout + + def _token_form_data(self) -> dict: + data = dict(self.data) + data["client_id"] = self.client_id + data["client_secret"] = self.client_secret + return data + + def _token_tuple(self, body: dict) -> _TokenTuple: + token = body.get(self.token_field_name) + if not token or not isinstance(token, str): + raise ValueError("OAuth token response missing access_token") + expires_in = body.get("expires_in") + _LOG.info("Fetched OAuth access token") + if expires_in not in (None, ""): + return self.state, token, expires_in + if not self._has_jwt_exp(token): + raise ValueError( + "OAuth token response omitted expires_in and " + "access_token is not a JWT with exp" + ) + return self.state, token + + @staticmethod + def _has_jwt_exp(token: str) -> bool: + """Return True when token is a 3-part JWT whose payload has exp. + + Mirrors DefaultOAuth2TokenProvider._expires_at_millis: opaque or + reference tokens must not be handed to httpx-auth as a 2-tuple, + which splits on '.' and crashes the next tool call. + """ + parts = token.split(".") + if len(parts) != 3: + return False + try: + padded = parts[1] + "=" * (-len(parts[1]) % 4) + payload = json.loads(base64.urlsafe_b64decode(padded)) + except (ValueError, json.JSONDecodeError): + return False + return isinstance(payload.get("exp"), int) + + @staticmethod + def _log_token_http_error(response: httpx.Response) -> None: + if response.status_code >= 400: + _LOG.error("OAuth token request failed: HTTP %s", response.status_code) Review Comment: CI will fail here. `pythonCheckFormat` in `build.gradle.kts` runs `black --check`, and this line is 82 chars. I ran it locally and it reports 2 files. Please run `black --line-length 80 mcp_server tests`. ########## mcp-server/tests/unit/test_oauth.py: ########## @@ -0,0 +1,485 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Tests for OAuth2 client-credentials fetch, cache, refresh, and 401 retry.""" + +# pylint: disable=protected-access + +import asyncio +import base64 +import json +import sys +import time +import unittest +from unittest import mock + +import httpx +from httpx_auth import OAuth2 + +from mcp_server.client.factory import RESTClientFactory +from mcp_server.client.plain.plain_rest_client_operation import ( + PlainRESTClientOperation, +) +from mcp_server.core.context import ( + GravitinoContext, + service_fallback_authorization, +) +from mcp_server.core.oauth import RefreshableBearerAuth +from mcp_server.core.setting import Setting +from mcp_server.main import _parse_args + + +def _jwt_with_exp(exp: int) -> str: + header = base64.urlsafe_b64encode(b'{"alg":"none"}').rstrip(b"=").decode() + payload = ( + base64.urlsafe_b64encode(json.dumps({"exp": exp}).encode()) + .rstrip(b"=") + .decode() + ) + return f"{header}.{payload}.sig" + + +class _OAuthHttpTestCase(unittest.TestCase): + """Drive ``httpx.AsyncClient(auth=...)`` against a shared MockTransport. + + httpx-auth caches tokens in a process-global map, so every test clears it. + Tests inject a sync ``httpx.Client`` on the same transport so IdP calls + never leave the process. Production token POSTs use ``AsyncClient``. + """ + + def setUp(self): + OAuth2.token_cache.clear() + + def tearDown(self): + OAuth2.token_cache.clear() + + def _auth(self, handler, **kwargs) -> RefreshableBearerAuth: + transport = httpx.MockTransport(handler) + self.addCleanup(OAuth2.token_cache.clear) + token_client = httpx.Client(transport=transport) + self.addCleanup(token_client.close) + auth = RefreshableBearerAuth( + token_endpoint="https://idp.example/token";, + client_id="mcp", + client_secret="s3cret", + client=token_client, + **kwargs, + ) + auth._test_transport = transport + return auth + + def _get(self, auth: RefreshableBearerAuth, path: str = "/api"): + async def _run(): + async with httpx.AsyncClient( + auth=auth, + transport=auth._test_transport, + base_url="https://gravitino.example";, + ) as client: + return await client.get(path) + + return asyncio.run(_run()) + + def _get_twice(self, auth: RefreshableBearerAuth): + async def _run(): + async with httpx.AsyncClient( + auth=auth, + transport=auth._test_transport, + base_url="https://gravitino.example";, + ) as client: + first = await client.get("/api") + second = await client.get("/api") + return first, second + + return asyncio.run(_run()) + + +class TestRefreshableBearerAuth(_OAuthHttpTestCase): + def test_fetches_and_reuses_while_fresh(self): + calls = {"token": 0, "api": 0} + + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "POST": + calls["token"] += 1 + self.assertEqual(str(request.url), "https://idp.example/token";) + body = request.content.decode() + self.assertIn("grant_type=client_credentials", body) + self.assertIn("client_id=mcp", body) + self.assertIn("client_secret=s3cret", body) + self.assertIsNone(request.headers.get("authorization")) + return httpx.Response( + 200, json={"access_token": "tok-1", "expires_in": 3600} + ) + calls["api"] += 1 + self.assertEqual( + request.headers.get("authorization"), "Bearer tok-1" + ) + return httpx.Response(200, json={"ok": True}) + + first, second = self._get_twice(self._auth(handler)) + self.assertEqual(first.status_code, 200) + self.assertEqual(second.status_code, 200) + self.assertEqual(calls["token"], 1) + self.assertEqual(calls["api"], 2) + + def test_sends_scope_when_configured(self): + seen = {} + + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "POST": + seen["body"] = request.content.decode() + return httpx.Response( + 200, json={"access_token": "tok", "expires_in": 3600} + ) + return httpx.Response(200, json={"ok": True}) + + self._get(self._auth(handler, scope="gravitino")) + self.assertIn("scope=gravitino", seen["body"]) + + def test_refetches_when_expired(self): + calls = {"token": 0} + seen = [] + + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "POST": + calls["token"] += 1 + return httpx.Response( + 200, + json={ + "access_token": f"tok-{calls['token']}", + "expires_in": 1, + }, + ) + seen.append(request.headers.get("authorization")) + return httpx.Response(200, json={"ok": True}) + + # Skew > expires_in makes the cache stale immediately after fetch. + self._get_twice(self._auth(handler, refresh_skew_seconds=60)) + self.assertEqual(seen, ["Bearer tok-1", "Bearer tok-2"]) + self.assertEqual(calls["token"], 2) + + def test_accepts_string_expires_in(self): + calls = {"token": 0} + + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "POST": + calls["token"] += 1 + return httpx.Response( + 200, + json={"access_token": "tok-1", "expires_in": "3600"}, + ) + return httpx.Response(200, json={"ok": True}) + + self._get_twice(self._auth(handler)) + self.assertEqual(calls["token"], 1) + + def test_uses_jwt_exp_when_expires_in_missing(self): + token = _jwt_with_exp(int(time.time()) + 3600) + calls = {"token": 0} + + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "POST": + calls["token"] += 1 + return httpx.Response(200, json={"access_token": token}) + self.assertEqual( + request.headers.get("authorization"), f"Bearer {token}" + ) + return httpx.Response(200, json={"ok": True}) + + self._get_twice(self._auth(handler)) + self.assertEqual(calls["token"], 1) + + def test_opaque_token_without_expires_in_raises(self): + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "POST": + return httpx.Response(200, json={"access_token": "opaque-ref"}) + return httpx.Response(200) + + with self.assertRaises(ValueError) as raised: + self._get(self._auth(handler)) + self.assertIn("expires_in", str(raised.exception)) + self.assertIn("JWT", str(raised.exception)) + + def test_jwt_without_exp_and_expires_in_raises(self): + header = base64.urlsafe_b64encode(b'{"alg":"none"}').rstrip(b"=").decode() Review Comment: Same black problem. This line is 82 chars because of the indent. Line 47 has the same expression at 4 spaces so it is fine. The same black command fixes both. ########## docs/gravitino-mcp-server.md: ########## @@ -140,13 +148,17 @@ You could config Gravitino MCP server by arguments, `uv run mcp_server -h` shows | `--gravitino-uri` | The URI of Gravitino server. | `http://127.0.0.1:8090` | No | | `--transport` | Transport protocol: stdio (local), http / streamable-http (Streamable HTTP). | `stdio` | No | | `--mcp-url` | The URL of MCP server if using HTTP transport. | `http://127.0.0.1:8000/mcp` | No | -| `--token` | Static credential for Gravitino; or set `GRAVITINO_TOKEN`. See Authentication. | none (anonymous) | No | -| `--tls-cert` | PEM certificate to serve the endpoint over HTTPS. Requires `--tls-key`. | none | No | -| `--tls-key` | PEM private key to serve the endpoint over HTTPS. Requires `--tls-cert`. | none | No | +| `--token` | Static credential for Gravitino; or set `GRAVITINO_TOKEN`. See Authentication. Wins over OAuth client-credentials. | none (anonymous) | No | +| `--oauth-token-endpoint` | OAuth2 token URL for client-credentials. Or `GRAVITINO_OAUTH_TOKEN_ENDPOINT`. | none | No | Review Comment: The pipes in this table are not aligned any more. The old rows use 19 chars for the first column and the new rows use 29. Please pad the whole table to the widest cell. ########## mcp-server/mcp_server/core/oauth.py: ########## @@ -0,0 +1,215 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""httpx ``auth=`` hook for MCP → Gravitino OAuth2 client-credentials. + +Uses ``httpx-auth`` for fetch and cache on the existing ``httpx.AsyncClient``. +Credentials go in the form body (``client_secret_post``), matching the +Java/Python Gravitino clients. httpx-auth defaults to HTTP Basic. This class +retries once after Gravitino HTTP 401. +""" + +import asyncio +import base64 +import json +import logging +from collections.abc import AsyncGenerator, Generator +from typing import Optional, Union + +import httpx +from httpx_auth import AuthenticationFailed, OAuth2, OAuth2ClientCredentials + +_LOG = logging.getLogger(__name__) + +# Refresh this many seconds before recorded expiry. httpx-auth default is 30. +DEFAULT_REFRESH_SKEW_SECONDS = 60 + +_TokenTuple = Union[tuple[str, str], tuple[str, str, Union[int, str]]] + + +class RefreshableBearerAuth(OAuth2ClientCredentials): + """httpx-auth client-credentials with form POST and one 401 retry.""" + + requires_request_body = True + requires_response_body = True + + def __init__( + self, + *, + token_endpoint: str, + client_id: str, + client_secret: str, + scope: str = "", + refresh_skew_seconds: int = DEFAULT_REFRESH_SKEW_SECONDS, + client: Optional[httpx.Client] = None, + ): + """Build an ``auth=`` hook for the service hop. + + Args: + token_endpoint: Identity-provider token URL. + client_id: OAuth2 client id. + client_secret: OAuth2 client secret. + scope: Optional OAuth2 scope. + refresh_skew_seconds: httpx-auth ``early_expiry``. + client: Optional sync httpx client used only for token POSTs + (tests inject ``MockTransport`` here). + """ + kwargs = {"early_expiry": float(refresh_skew_seconds)} + if scope: + kwargs["scope"] = scope + if client is not None: + kwargs["client"] = client + super().__init__(token_endpoint, client_id, client_secret, **kwargs) + + def invalidate(self) -> None: + """Drop the cached token so the next call fetches a new one.""" + cache = OAuth2.token_cache + # TokenMemoryCache.clear() wipes every client; only drop ours. + with cache._forbid_concurrent_cache_access: # pylint: disable=protected-access + cache.tokens.pop(self.state, None) + + def request_new_token(self) -> _TokenTuple: + """POST ``client_credentials`` with id/secret in the form body.""" + data = self._token_form_data() + client = self.client or httpx.Client() + self._configure_client(client) + try: + response = client.post(self.token_url, data=data) + self._log_token_http_error(response) + response.raise_for_status() + body = response.json() + finally: + if self.client is None: + client.close() + return self._token_tuple(body) + + async def request_new_token_async(self) -> _TokenTuple: + """POST ``client_credentials`` without blocking the event loop.""" + if self.client is not None: + return await asyncio.to_thread(self.request_new_token) + data = self._token_form_data() + async with httpx.AsyncClient() as client: + client.timeout = self.timeout + response = await client.post(self.token_url, data=data) + self._log_token_http_error(response) + response.raise_for_status() + body = response.json() + return self._token_tuple(body) + + def auth_flow( + self, request: httpx.Request + ) -> Generator[httpx.Request, httpx.Response, None]: + """Attach a cached or freshly fetched Bearer; retry once on HTTP 401.""" + self._apply_token(request) + response = yield request + if response.status_code != 401: + return + self.invalidate() + self._apply_token(request) + yield request + + async def async_auth_flow( + self, request: httpx.Request + ) -> AsyncGenerator[httpx.Request, httpx.Response]: + """Attach a Bearer without a blocking IdP POST on the event loop.""" + if self.requires_request_body: + await request.aread() + await self._apply_token_async(request) + response = yield request + if response.status_code != 401: + return + self.invalidate() Review Comment: `invalidate()` drops the token from the global `OAuth2.token_cache`, which every in-flight request shares. So one 401 makes all of them lose the token, and together with the missing lock above, each one fetches its own new token. Also, when the 401 is not about expiry (wrong audience or issuer, `principalFields` not configured, clock skew), every tool call costs 2 Gravitino requests plus 1 token POST forever, with no backoff. Maybe only invalidate when the token we just used is still the cached one, and skip the retry when the token was fetched in this same flow. ########## mcp-server/mcp_server/core/oauth.py: ########## @@ -0,0 +1,215 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""httpx ``auth=`` hook for MCP → Gravitino OAuth2 client-credentials. + +Uses ``httpx-auth`` for fetch and cache on the existing ``httpx.AsyncClient``. +Credentials go in the form body (``client_secret_post``), matching the +Java/Python Gravitino clients. httpx-auth defaults to HTTP Basic. This class +retries once after Gravitino HTTP 401. +""" + +import asyncio +import base64 +import json +import logging +from collections.abc import AsyncGenerator, Generator +from typing import Optional, Union + +import httpx +from httpx_auth import AuthenticationFailed, OAuth2, OAuth2ClientCredentials + +_LOG = logging.getLogger(__name__) + +# Refresh this many seconds before recorded expiry. httpx-auth default is 30. +DEFAULT_REFRESH_SKEW_SECONDS = 60 + +_TokenTuple = Union[tuple[str, str], tuple[str, str, Union[int, str]]] + + +class RefreshableBearerAuth(OAuth2ClientCredentials): + """httpx-auth client-credentials with form POST and one 401 retry.""" + + requires_request_body = True + requires_response_body = True + + def __init__( + self, + *, + token_endpoint: str, + client_id: str, + client_secret: str, + scope: str = "", + refresh_skew_seconds: int = DEFAULT_REFRESH_SKEW_SECONDS, + client: Optional[httpx.Client] = None, + ): + """Build an ``auth=`` hook for the service hop. + + Args: + token_endpoint: Identity-provider token URL. + client_id: OAuth2 client id. + client_secret: OAuth2 client secret. + scope: Optional OAuth2 scope. + refresh_skew_seconds: httpx-auth ``early_expiry``. + client: Optional sync httpx client used only for token POSTs + (tests inject ``MockTransport`` here). + """ + kwargs = {"early_expiry": float(refresh_skew_seconds)} + if scope: + kwargs["scope"] = scope + if client is not None: + kwargs["client"] = client + super().__init__(token_endpoint, client_id, client_secret, **kwargs) + + def invalidate(self) -> None: + """Drop the cached token so the next call fetches a new one.""" + cache = OAuth2.token_cache + # TokenMemoryCache.clear() wipes every client; only drop ours. + with cache._forbid_concurrent_cache_access: # pylint: disable=protected-access + cache.tokens.pop(self.state, None) + + def request_new_token(self) -> _TokenTuple: + """POST ``client_credentials`` with id/secret in the form body.""" + data = self._token_form_data() + client = self.client or httpx.Client() + self._configure_client(client) + try: + response = client.post(self.token_url, data=data) + self._log_token_http_error(response) + response.raise_for_status() + body = response.json() + finally: + if self.client is None: + client.close() + return self._token_tuple(body) + + async def request_new_token_async(self) -> _TokenTuple: + """POST ``client_credentials`` without blocking the event loop.""" + if self.client is not None: + return await asyncio.to_thread(self.request_new_token) + data = self._token_form_data() + async with httpx.AsyncClient() as client: + client.timeout = self.timeout + response = await client.post(self.token_url, data=data) + self._log_token_http_error(response) + response.raise_for_status() + body = response.json() + return self._token_tuple(body) + + def auth_flow( + self, request: httpx.Request + ) -> Generator[httpx.Request, httpx.Response, None]: + """Attach a cached or freshly fetched Bearer; retry once on HTTP 401.""" + self._apply_token(request) + response = yield request + if response.status_code != 401: + return + self.invalidate() + self._apply_token(request) + yield request + + async def async_auth_flow( + self, request: httpx.Request + ) -> AsyncGenerator[httpx.Request, httpx.Response]: + """Attach a Bearer without a blocking IdP POST on the event loop.""" + if self.requires_request_body: + await request.aread() + await self._apply_token_async(request) + response = yield request + if response.status_code != 401: + return + self.invalidate() + await self._apply_token_async(request) + yield request + + def _configure_client(self, client: httpx.Client) -> None: + """Do not send HTTP Basic; id and secret go in the form body.""" + client.timeout = self.timeout + + def _token_form_data(self) -> dict: + data = dict(self.data) + data["client_id"] = self.client_id + data["client_secret"] = self.client_secret + return data + + def _token_tuple(self, body: dict) -> _TokenTuple: + token = body.get(self.token_field_name) + if not token or not isinstance(token, str): + raise ValueError("OAuth token response missing access_token") + expires_in = body.get("expires_in") + _LOG.info("Fetched OAuth access token") + if expires_in not in (None, ""): Review Comment: `expires_in: 0` passes this check and goes to `to_expiry(0)`, so we cache a token that is already expired. After that `_cached_bearer()` returns None every time and we fetch a new token on every call, and we never fall back to the JWT `exp`. Upstream uses `if expires_in:` and treats 0 as missing. I think we should do the same. ########## mcp-server/mcp_server/core/oauth.py: ########## @@ -0,0 +1,215 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""httpx ``auth=`` hook for MCP → Gravitino OAuth2 client-credentials. + +Uses ``httpx-auth`` for fetch and cache on the existing ``httpx.AsyncClient``. +Credentials go in the form body (``client_secret_post``), matching the +Java/Python Gravitino clients. httpx-auth defaults to HTTP Basic. This class +retries once after Gravitino HTTP 401. +""" + +import asyncio +import base64 +import json +import logging +from collections.abc import AsyncGenerator, Generator +from typing import Optional, Union + +import httpx +from httpx_auth import AuthenticationFailed, OAuth2, OAuth2ClientCredentials + +_LOG = logging.getLogger(__name__) + +# Refresh this many seconds before recorded expiry. httpx-auth default is 30. +DEFAULT_REFRESH_SKEW_SECONDS = 60 + +_TokenTuple = Union[tuple[str, str], tuple[str, str, Union[int, str]]] + + +class RefreshableBearerAuth(OAuth2ClientCredentials): + """httpx-auth client-credentials with form POST and one 401 retry.""" + + requires_request_body = True + requires_response_body = True + + def __init__( + self, + *, + token_endpoint: str, + client_id: str, + client_secret: str, + scope: str = "", + refresh_skew_seconds: int = DEFAULT_REFRESH_SKEW_SECONDS, + client: Optional[httpx.Client] = None, + ): + """Build an ``auth=`` hook for the service hop. + + Args: + token_endpoint: Identity-provider token URL. + client_id: OAuth2 client id. + client_secret: OAuth2 client secret. + scope: Optional OAuth2 scope. + refresh_skew_seconds: httpx-auth ``early_expiry``. + client: Optional sync httpx client used only for token POSTs + (tests inject ``MockTransport`` here). + """ + kwargs = {"early_expiry": float(refresh_skew_seconds)} + if scope: + kwargs["scope"] = scope + if client is not None: + kwargs["client"] = client + super().__init__(token_endpoint, client_id, client_secret, **kwargs) + + def invalidate(self) -> None: + """Drop the cached token so the next call fetches a new one.""" + cache = OAuth2.token_cache + # TokenMemoryCache.clear() wipes every client; only drop ours. + with cache._forbid_concurrent_cache_access: # pylint: disable=protected-access + cache.tokens.pop(self.state, None) + + def request_new_token(self) -> _TokenTuple: + """POST ``client_credentials`` with id/secret in the form body.""" + data = self._token_form_data() + client = self.client or httpx.Client() + self._configure_client(client) + try: + response = client.post(self.token_url, data=data) + self._log_token_http_error(response) + response.raise_for_status() + body = response.json() + finally: + if self.client is None: + client.close() + return self._token_tuple(body) + + async def request_new_token_async(self) -> _TokenTuple: + """POST ``client_credentials`` without blocking the event loop.""" + if self.client is not None: + return await asyncio.to_thread(self.request_new_token) + data = self._token_form_data() + async with httpx.AsyncClient() as client: + client.timeout = self.timeout + response = await client.post(self.token_url, data=data) + self._log_token_http_error(response) + response.raise_for_status() + body = response.json() + return self._token_tuple(body) + + def auth_flow( + self, request: httpx.Request + ) -> Generator[httpx.Request, httpx.Response, None]: + """Attach a cached or freshly fetched Bearer; retry once on HTTP 401.""" + self._apply_token(request) + response = yield request + if response.status_code != 401: + return + self.invalidate() + self._apply_token(request) + yield request + + async def async_auth_flow( + self, request: httpx.Request + ) -> AsyncGenerator[httpx.Request, httpx.Response]: + """Attach a Bearer without a blocking IdP POST on the event loop.""" + if self.requires_request_body: + await request.aread() + await self._apply_token_async(request) + response = yield request + if response.status_code != 401: + return + self.invalidate() + await self._apply_token_async(request) + yield request + + def _configure_client(self, client: httpx.Client) -> None: + """Do not send HTTP Basic; id and secret go in the form body.""" + client.timeout = self.timeout + + def _token_form_data(self) -> dict: + data = dict(self.data) + data["client_id"] = self.client_id + data["client_secret"] = self.client_secret + return data + + def _token_tuple(self, body: dict) -> _TokenTuple: + token = body.get(self.token_field_name) + if not token or not isinstance(token, str): + raise ValueError("OAuth token response missing access_token") + expires_in = body.get("expires_in") + _LOG.info("Fetched OAuth access token") + if expires_in not in (None, ""): + return self.state, token, expires_in + if not self._has_jwt_exp(token): + raise ValueError( + "OAuth token response omitted expires_in and " + "access_token is not a JWT with exp" + ) + return self.state, token + + @staticmethod + def _has_jwt_exp(token: str) -> bool: + """Return True when token is a 3-part JWT whose payload has exp. + + Mirrors DefaultOAuth2TokenProvider._expires_at_millis: opaque or + reference tokens must not be handed to httpx-auth as a 2-tuple, + which splits on '.' and crashes the next tool call. + """ + parts = token.split(".") + if len(parts) != 3: + return False + try: + padded = parts[1] + "=" * (-len(parts[1]) % 4) + payload = json.loads(base64.urlsafe_b64decode(padded)) + except (ValueError, json.JSONDecodeError): + return False + return isinstance(payload.get("exp"), int) + + @staticmethod + def _log_token_http_error(response: httpx.Response) -> None: + if response.status_code >= 400: + _LOG.error("OAuth token request failed: HTTP %s", response.status_code) + + def _cached_bearer(self) -> Optional[str]: + try: + return OAuth2.token_cache.get_token( + self.state, early_expiry=self.early_expiry + ) + except AuthenticationFailed: + return None + + def _store_and_get(self, fetched: _TokenTuple) -> str: + return OAuth2.token_cache.get_token( + self.state, + early_expiry=self.early_expiry, + on_missing_token=lambda: fetched, + ) + + def _apply_token(self, request: httpx.Request) -> None: + token = OAuth2.token_cache.get_token( + self.state, + early_expiry=self.early_expiry, + on_missing_token=self.request_new_token, + on_expired_token=self.refresh_token, + ) + self._update_user_request(request, token) + + async def _apply_token_async(self, request: httpx.Request) -> None: + token = self._cached_bearer() Review Comment: This loses the single-flight behavior of httpx-auth. Upstream `get_token` takes a lock so only one fetch happens, but here we check the cache, then `await` the fetch, then store, with no lock. So every coroutine that arrives during the fetch starts its own token POST. I tried a mock IdP with 50ms delay and 10 parallel tool calls on a cold cache, and I see 10 token POSTs. The current tests use a zero-delay mock, so they always see 1 and cannot catch this. Keycloak and Okta rate limit `client_credentials`, so a burst of tool calls at startup can get 429. Can we hold an `asyncio.Lock` around miss -> fetch -> store, and check the cache again after we get the lock? ########## docs/gravitino-mcp-server.md: ########## @@ -175,11 +187,31 @@ export GRAVITINO_TOKEN=<your-token> uv run mcp_server --metalake test --gravitino-uri http://127.0.0.1:8090 ``` -In `stdio` mode this token is used for every request. In HTTP mode it is only the fallback, used when an incoming request does not carry its own `Authorization` header. +In `stdio` mode this token is used for every request. In HTTP mode it is only the fallback, used when an incoming request does not carry its own `Authorization` header. If both `--token` and OAuth client-credentials are set, `--token` wins. + +### OAuth client credentials (service identity) + +When Gravitino uses `gravitino.authenticators = oauth`, a pasted Bearer access token in `--token` expires and is not refreshed. For the **service** identity (Cursor stdio, or HTTP when the caller sends no `Authorization` header), configure MCP as an OAuth client of the same identity provider Gravitino trusts. + +Set `--oauth-token-endpoint`, `--oauth-client-id`, and `--oauth-client-secret` together, plus optional `--oauth-scope` (or the matching `GRAVITINO_OAUTH_*` environment variables). MCP requests an access token with the `client_credentials` grant, caches it, refreshes before expiry, and retries once on HTTP 401. + +In Cursor, put the `GRAVITINO_OAUTH_*` values in the `env` block of `~/.cursor/mcp.json` (see [Usage](#usage)). `--token` / `GRAVITINO_TOKEN` overrides OAuth client-credentials and stays static (no refresh). An incoming HTTP `Authorization` header is forwarded as-is and is not refreshed by MCP. + +Gravitino maps the JWT to a metalake principal from claims configured in [`gravitino.authenticator.oauth.principalFields`](./security/how-to-authenticate.md#server-configuration) (often `sub`); that principal may differ from `--oauth-client-id`. It must exist as a metalake user with the needed grants, or tool calls fail with 403. + +```shell +uv run mcp_server --metalake test --gravitino-uri http://127.0.0.1:8090 \ + --oauth-token-endpoint https://idp.example/realms/gravitino/protocol/openid-connect/token \ + --oauth-client-id mcp-service \ + --oauth-client-secret <secret> \ Review Comment: Putting the secret on the command line makes it visible in `ps` and `/proc/<pid>/cmdline` to any user on the host, and it also goes into shell history. `GRAVITINO_OAUTH_CLIENT_SECRET` already works, so can we show the env var form first here and add a short note about `ps`? ########## mcp-server/mcp_server/main.py: ########## @@ -35,7 +35,12 @@ def do_main(): token=args.token, tls_cert=args.tls_cert, tls_key=args.tls_key, + oauth_token_endpoint=args.oauth_token_endpoint, + oauth_client_id=args.oauth_client_id, + oauth_client_secret=args.oauth_client_secret, + oauth_scope=args.oauth_scope, ) + setting.validate_oauth() Review Comment: This runs before `_init_logging()`, so a partial config exits with a raw traceback. It is easy to hit, for example only `GRAVITINO_OAUTH_CLIENT_ID` set in the Cursor `env` block. In stdio mode Cursor only shows a dead process, which is hard to debug. Can we move this after logging is ready? Also, setting only `--oauth-scope` passes the validation and is then ignored silently. ########## mcp-server/mcp_server/core/oauth.py: ########## @@ -0,0 +1,215 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""httpx ``auth=`` hook for MCP → Gravitino OAuth2 client-credentials. + +Uses ``httpx-auth`` for fetch and cache on the existing ``httpx.AsyncClient``. +Credentials go in the form body (``client_secret_post``), matching the +Java/Python Gravitino clients. httpx-auth defaults to HTTP Basic. This class +retries once after Gravitino HTTP 401. +""" + +import asyncio +import base64 +import json +import logging +from collections.abc import AsyncGenerator, Generator +from typing import Optional, Union + +import httpx +from httpx_auth import AuthenticationFailed, OAuth2, OAuth2ClientCredentials + +_LOG = logging.getLogger(__name__) + +# Refresh this many seconds before recorded expiry. httpx-auth default is 30. +DEFAULT_REFRESH_SKEW_SECONDS = 60 + +_TokenTuple = Union[tuple[str, str], tuple[str, str, Union[int, str]]] + + +class RefreshableBearerAuth(OAuth2ClientCredentials): + """httpx-auth client-credentials with form POST and one 401 retry.""" + + requires_request_body = True + requires_response_body = True + + def __init__( + self, + *, + token_endpoint: str, + client_id: str, + client_secret: str, + scope: str = "", + refresh_skew_seconds: int = DEFAULT_REFRESH_SKEW_SECONDS, + client: Optional[httpx.Client] = None, + ): + """Build an ``auth=`` hook for the service hop. + + Args: + token_endpoint: Identity-provider token URL. + client_id: OAuth2 client id. + client_secret: OAuth2 client secret. + scope: Optional OAuth2 scope. + refresh_skew_seconds: httpx-auth ``early_expiry``. + client: Optional sync httpx client used only for token POSTs + (tests inject ``MockTransport`` here). + """ + kwargs = {"early_expiry": float(refresh_skew_seconds)} + if scope: + kwargs["scope"] = scope + if client is not None: + kwargs["client"] = client + super().__init__(token_endpoint, client_id, client_secret, **kwargs) + + def invalidate(self) -> None: + """Drop the cached token so the next call fetches a new one.""" + cache = OAuth2.token_cache + # TokenMemoryCache.clear() wipes every client; only drop ours. + with cache._forbid_concurrent_cache_access: # pylint: disable=protected-access + cache.tokens.pop(self.state, None) + + def request_new_token(self) -> _TokenTuple: + """POST ``client_credentials`` with id/secret in the form body.""" + data = self._token_form_data() + client = self.client or httpx.Client() + self._configure_client(client) + try: + response = client.post(self.token_url, data=data) + self._log_token_http_error(response) + response.raise_for_status() + body = response.json() + finally: + if self.client is None: + client.close() + return self._token_tuple(body) + + async def request_new_token_async(self) -> _TokenTuple: + """POST ``client_credentials`` without blocking the event loop.""" + if self.client is not None: + return await asyncio.to_thread(self.request_new_token) + data = self._token_form_data() + async with httpx.AsyncClient() as client: + client.timeout = self.timeout + response = await client.post(self.token_url, data=data) + self._log_token_http_error(response) + response.raise_for_status() + body = response.json() + return self._token_tuple(body) + + def auth_flow( + self, request: httpx.Request + ) -> Generator[httpx.Request, httpx.Response, None]: + """Attach a cached or freshly fetched Bearer; retry once on HTTP 401.""" + self._apply_token(request) + response = yield request + if response.status_code != 401: + return + self.invalidate() + self._apply_token(request) + yield request + + async def async_auth_flow( + self, request: httpx.Request + ) -> AsyncGenerator[httpx.Request, httpx.Response]: + """Attach a Bearer without a blocking IdP POST on the event loop.""" + if self.requires_request_body: + await request.aread() + await self._apply_token_async(request) + response = yield request + if response.status_code != 401: + return + self.invalidate() + await self._apply_token_async(request) + yield request + + def _configure_client(self, client: httpx.Client) -> None: + """Do not send HTTP Basic; id and secret go in the form body.""" + client.timeout = self.timeout + + def _token_form_data(self) -> dict: + data = dict(self.data) + data["client_id"] = self.client_id + data["client_secret"] = self.client_secret + return data + + def _token_tuple(self, body: dict) -> _TokenTuple: + token = body.get(self.token_field_name) + if not token or not isinstance(token, str): + raise ValueError("OAuth token response missing access_token") + expires_in = body.get("expires_in") + _LOG.info("Fetched OAuth access token") + if expires_in not in (None, ""): + return self.state, token, expires_in + if not self._has_jwt_exp(token): + raise ValueError( + "OAuth token response omitted expires_in and " + "access_token is not a JWT with exp" + ) + return self.state, token + + @staticmethod + def _has_jwt_exp(token: str) -> bool: + """Return True when token is a 3-part JWT whose payload has exp. + + Mirrors DefaultOAuth2TokenProvider._expires_at_millis: opaque or + reference tokens must not be handed to httpx-auth as a 2-tuple, + which splits on '.' and crashes the next tool call. + """ + parts = token.split(".") + if len(parts) != 3: + return False + try: + padded = parts[1] + "=" * (-len(parts[1]) % 4) + payload = json.loads(base64.urlsafe_b64decode(padded)) + except (ValueError, json.JSONDecodeError): + return False + return isinstance(payload.get("exp"), int) Review Comment: `exp` is not always an int. RFC 7519 says NumericDate is a JSON number, so a float is valid, and some IdPs send it as a string. If that IdP also omits `expires_in`, `_token_tuple` raises `ValueError` and every tool call fails, even though the token is good. Can we accept `(int, float)` here? ########## mcp-server/pyproject.toml: ########## @@ -24,6 +24,8 @@ requires-python = ">=3.10" dependencies = [ # Pin FastMCP so breaking API changes are handled explicitly during dependency upgrades. "fastmcp==3.4.5", + # httpx.Auth plugin for hop-2 client_credentials fetch/cache. + "httpx-auth>=0.22", Review Comment: The code uses many httpx-auth internals: `self.state`, `self.data`, `self.timeout`, `self.token_field_name`, `self.early_expiry`, `OAuth2.token_cache`, `cache.tokens`, `_update_user_request`, and `cache._forbid_concurrent_cache_access` which is private, we even disable pylint for it. With `>=0.22` and no upper bound, a minor release can break auth at runtime. Can we add an upper bound, like `>=0.22,<0.24`? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
