This is an automated email from the ASF dual-hosted git repository.
yufei pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/polaris-tools.git
The following commit(s) were added to refs/heads/main by this push:
new c3c1f5f MCP: Add configurable refresh buffer (#74)
c3c1f5f is described below
commit c3c1f5ffc563bbe49387239578a8857c44b3625b
Author: Yong Zheng <[email protected]>
AuthorDate: Mon Nov 24 19:06:55 2025 -0600
MCP: Add configurable refresh buffer (#74)
---
mcp-server/README.md | 17 +++++-----
mcp-server/polaris_mcp/authorization.py | 11 ++++--
mcp-server/polaris_mcp/server.py | 9 +++++
mcp-server/tests/test_authorization.py | 60 +++++++++++++++++++++++++++++++++
mcp-server/tests/test_server.py | 1 +
5 files changed, 87 insertions(+), 11 deletions(-)
diff --git a/mcp-server/README.md b/mcp-server/README.md
index a1c57c2..042b5e4 100644
--- a/mcp-server/README.md
+++ b/mcp-server/README.md
@@ -65,14 +65,15 @@ Please note: `--directory` specifies a local directory. It
is not needed when we
## Configuration
-| Variable | Description
| Default
|
-|----------------------------------------------------------------|----------------------------------------------------------|--------------------------------------------------|
-| `POLARIS_BASE_URL` | Base URL
for all Polaris REST calls. | `http://localhost:8181/`
|
-| `POLARIS_API_TOKEN` / `POLARIS_BEARER_TOKEN` / `POLARIS_TOKEN` | Static
bearer token (if supplied, overrides other auth). | _unset_
|
-| `POLARIS_CLIENT_ID` | OAuth
client id for client-credential flow. | _unset_
|
-| `POLARIS_CLIENT_SECRET` | OAuth
client secret. | _unset_
|
-| `POLARIS_TOKEN_SCOPE` | OAuth scope
string. | _unset_
|
-| `POLARIS_TOKEN_URL` | Optional
override for the token endpoint URL. |
`${POLARIS_BASE_URL}api/catalog/v1/oauth/tokens` |
+| Variable | Description
| Default
|
+|----------------------------------------------------------------|----------------------------------------------------------------|--------------------------------------------------|
+| `POLARIS_BASE_URL` | Base URL
for all Polaris REST calls. |
`http://localhost:8181/` |
+| `POLARIS_API_TOKEN` / `POLARIS_BEARER_TOKEN` / `POLARIS_TOKEN` | Static
bearer token (if supplied, overrides other auth). | _unset_
|
+| `POLARIS_CLIENT_ID` | OAuth
client id for client-credential flow. | _unset_
|
+| `POLARIS_CLIENT_SECRET` | OAuth
client secret. | _unset_
|
+| `POLARIS_TOKEN_SCOPE` | OAuth scope
string. | _unset_
|
+| `POLARIS_TOKEN_URL` | Optional
override for the token endpoint URL. |
`${POLARIS_BASE_URL}api/catalog/v1/oauth/tokens` |
+| `POLARIS_TOKEN_REFRESH_BUFFER_SECONDS` | Minimum
remaining token lifetime before refreshing in seconds. | `60.0`
|
When OAuth variables are supplied, the server automatically acquires and
refreshes tokens using the client credentials flow; otherwise a static bearer
token is used if provided.
diff --git a/mcp-server/polaris_mcp/authorization.py
b/mcp-server/polaris_mcp/authorization.py
index e58c8c8..0545ab4 100644
--- a/mcp-server/polaris_mcp/authorization.py
+++ b/mcp-server/polaris_mcp/authorization.py
@@ -59,6 +59,7 @@ class
ClientCredentialsAuthorizationProvider(AuthorizationProvider):
client_secret: str,
scope: Optional[str],
http: urllib3.PoolManager,
+ refresh_buffer_seconds: float,
) -> None:
self._token_endpoint = token_endpoint
self._client_id = client_id
@@ -67,6 +68,7 @@ class
ClientCredentialsAuthorizationProvider(AuthorizationProvider):
self._http = http
self._lock = threading.Lock()
self._cached: Optional[tuple[str, float]] = None # (token,
expires_at_epoch)
+ self._refresh_buffer_seconds = max(refresh_buffer_seconds, 0.0)
def authorization_header(self) -> Optional[str]:
token = self._current_token()
@@ -75,10 +77,13 @@ class
ClientCredentialsAuthorizationProvider(AuthorizationProvider):
def _current_token(self) -> Optional[str]:
now = time.time()
cached = self._cached
- if not cached or cached[1] - 60 <= now:
+ if not cached or cached[1] - self._refresh_buffer_seconds <= now:
with self._lock:
cached = self._cached
- if not cached or cached[1] - 60 <= time.time():
+ if (
+ not cached
+ or cached[1] - self._refresh_buffer_seconds <= time.time()
+ ):
self._cached = cached = self._fetch_token()
return cached[0] if cached else None
@@ -119,7 +124,7 @@ class
ClientCredentialsAuthorizationProvider(AuthorizationProvider):
ttl = float(expires_in)
except (TypeError, ValueError):
ttl = 3600.0
- ttl = max(ttl, 60.0)
+ ttl = max(ttl, self._refresh_buffer_seconds)
expires_at = time.time() + ttl
return token, expires_at
diff --git a/mcp-server/polaris_mcp/server.py b/mcp-server/polaris_mcp/server.py
index bdbacc1..c73c401 100644
--- a/mcp-server/polaris_mcp/server.py
+++ b/mcp-server/polaris_mcp/server.py
@@ -59,6 +59,7 @@ OUTPUT_SCHEMA = {
"required": ["isError"],
"additionalProperties": True,
}
+DEFAULT_TOKEN_REFRESH_BUFFER_SECONDS = 60.0
def create_server() -> FastMCP:
@@ -429,12 +430,20 @@ def _resolve_authorization_provider(
scope = _first_non_blank(os.getenv("POLARIS_TOKEN_SCOPE"))
token_url = _first_non_blank(os.getenv("POLARIS_TOKEN_URL"))
endpoint = token_url or urljoin(base_url,
"api/catalog/v1/oauth/tokens")
+ refresh_buffer_seconds = DEFAULT_TOKEN_REFRESH_BUFFER_SECONDS
+ refresh_buffer_seconds_str =
os.getenv("POLARIS_TOKEN_REFRESH_BUFFER_SECONDS")
+ if refresh_buffer_seconds_str:
+ try:
+ refresh_buffer_seconds =
float(refresh_buffer_seconds_str.strip())
+ except ValueError:
+ pass
return ClientCredentialsAuthorizationProvider(
token_endpoint=endpoint,
client_id=client_id,
client_secret=client_secret,
scope=scope,
http=http,
+ refresh_buffer_seconds=refresh_buffer_seconds,
)
return none()
diff --git a/mcp-server/tests/test_authorization.py
b/mcp-server/tests/test_authorization.py
index 4d57356..e120de6 100644
--- a/mcp-server/tests/test_authorization.py
+++ b/mcp-server/tests/test_authorization.py
@@ -65,6 +65,7 @@ def test_client_credentials_fetches_and_caches_tokens(
client_secret="secret",
scope=None,
http=http,
+ refresh_buffer_seconds=0.0,
)
with mock.patch("time.time", return_value=now):
@@ -94,6 +95,63 @@ def test_client_credentials_fetches_and_caches_tokens(
http.request.assert_called_once()
+def test_client_credentials_refresh_buffer() -> None:
+ http = mock.Mock()
+ now = time.time()
+ expires_in = 120
+ refresh_buffer = 30.0
+
+ response = SimpleNamespace(
+ status=200,
+ data=json.dumps({"access_token": "initial", "expires_in":
expires_in}).encode(
+ "utf-8"
+ ),
+ )
+ http.request.return_value = response
+
+ provider = ClientCredentialsAuthorizationProvider(
+ token_endpoint="https://auth/token";,
+ client_id="client",
+ client_secret="secret",
+ scope=None,
+ http=http,
+ refresh_buffer_seconds=refresh_buffer,
+ )
+
+ # Initial valid token
+ with mock.patch("time.time", return_value=now):
+ header1 = provider.authorization_header()
+ assert header1 == "Bearer initial"
+ http.request.assert_called_once()
+ http.request.reset_mock()
+
+ # Valid token before refresh
+ with mock.patch("time.time", return_value=now + (expires_in -
refresh_buffer - 1)):
+ header2 = provider.authorization_header()
+ assert header2 == "Bearer initial"
+ http.request.assert_not_called()
+
+ # Refresh token once reached refresh buffer
+ refreshed_response = SimpleNamespace(
+ status=200,
+ data=json.dumps({"access_token": "refreshed", "expires_in":
expires_in}).encode(
+ "utf-8"
+ ),
+ )
+ http.request.return_value = refreshed_response
+ with mock.patch("time.time", return_value=now + (expires_in -
refresh_buffer)):
+ header3 = provider.authorization_header()
+ assert header3 == "Bearer refreshed"
+ http.request.assert_called_once()
+ http.request.reset_mock()
+
+ # Force expiry to trigger a refresh
+ with mock.patch("time.time", return_value=now + expires_in + 1):
+ header4 = provider.authorization_header()
+ assert header4 == "Bearer refreshed"
+ http.request.assert_not_called()
+
+
@pytest.mark.parametrize(
"payload,expected_message",
[
@@ -118,6 +176,7 @@ def test_client_credentials_rejects_invalid_responses(
client_secret="secret",
scope=None,
http=http,
+ refresh_buffer_seconds=0.0,
)
with pytest.raises(RuntimeError, match=expected_message):
@@ -134,6 +193,7 @@ def test_client_credentials_errors_on_non_200_status() ->
None:
client_secret="secret",
scope=None,
http=http,
+ refresh_buffer_seconds=0.0,
)
with pytest.raises(RuntimeError, match="500"):
diff --git a/mcp-server/tests/test_server.py b/mcp-server/tests/test_server.py
index 9680cc8..06e4f1d 100644
--- a/mcp-server/tests/test_server.py
+++ b/mcp-server/tests/test_server.py
@@ -247,4 +247,5 @@ class TestAuthorizationProviderResolution:
client_secret="secret",
scope="scope",
http=fake_http,
+ refresh_buffer_seconds=60.0,
)
