yuqi1129 commented on code in PR #12960: URL: https://github.com/apache/gravitino/pull/12960#discussion_r3965261522
########## mcp-server/mcp_server/core/middleware.py: ########## @@ -0,0 +1,121 @@ +# 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. + +from typing import Any, Dict, Sequence + +import mcp.types as mt +from fastmcp.server.middleware.middleware import ( + CallNext, + Middleware, + MiddlewareContext, +) +from fastmcp.tools.base import Tool, ToolResult + +from mcp_server.core.context import ( + METALAKE_ARGUMENT, + reset_request_metalake, + set_request_metalake, +) + +# Tools that never resolve a metalake, so advertising the argument on them +# would offer the agent a knob that does nothing. `list_metalakes` is the +# discovery tool itself (its whole point is working without a metalake) and +# `metadata_type_to_fullname_formats` is pure computation that never calls +# Gravitino. A tool missing from this set only gets a harmless no-op argument. +TOOLS_WITHOUT_METALAKE = frozenset( + {"list_metalakes", "metadata_type_to_fullname_formats"} +) + +_METALAKE_ARGUMENT_DESCRIPTION = ( + "Metalake to operate on. Omit to use the server's configured default " + "metalake. Call 'list_metalakes' to discover which metalakes are " + "available to you." +) + + +def _schema_with_metalake(parameters: Dict[str, Any]) -> Dict[str, Any]: + """Return ``parameters`` with an optional ``metalake`` property added. + + Copied rather than mutated so the registered Tool objects keep the schema + their functions actually declare; the argument exists only on the wire. + ``required`` is deliberately left alone - omitting the argument is what + every single-metalake deployment does. + """ + schema = dict(parameters) + properties = dict(schema.get("properties") or {}) + # Never shadow a parameter a tool declares itself. + if METALAKE_ARGUMENT in properties: + return parameters + properties[METALAKE_ARGUMENT] = { + "type": "string", + "description": _METALAKE_ARGUMENT_DESCRIPTION, + } + schema["properties"] = properties + return schema + + +class MetalakeArgumentMiddleware(Middleware): + """Lets any tool call name the metalake it operates on. + + Every tool gains an optional ``metalake`` argument without declaring it: + this middleware advertises it in each tool's input schema, strips it from + the incoming arguments before the tool function runs, and publishes it for + ``GravitinoContext.rest_client()`` to resolve against. + + The value lives in a context variable for the duration of one tool call + only, so no metalake state is carried between calls or shared between + server replicas. + """ + + async def on_list_tools( + self, + context: MiddlewareContext[mt.ListToolsRequest], + call_next: CallNext[mt.ListToolsRequest, Sequence[Tool]], + ) -> Sequence[Tool]: + tools = await call_next(context) + return [ + ( + tool + if tool.name in TOOLS_WITHOUT_METALAKE + else tool.model_copy( + update={ + "parameters": _schema_with_metalake(tool.parameters) + } + ) + ) + for tool in tools + ] + + async def on_call_tool( + self, + context: MiddlewareContext[mt.CallToolRequestParams], + call_next: CallNext[mt.CallToolRequestParams, ToolResult], + ) -> ToolResult: + arguments = context.message.arguments + # Popped so the tool function never sees an argument it cannot accept. + metalake = ( + arguments.pop(METALAKE_ARGUMENT, "") + if isinstance(arguments, dict) + else "" + ) + token = set_request_metalake(metalake) Review Comment: [P1] Validate the metalake argument before removing it from framework validation The advertised schema requires a string, but this middleware removes the argument before FastMCP validates it. `set_request_metalake()` then uses `(metalake or "").strip()`: explicit values such as `false`, `0`, and `[]` silently become an omitted argument and route to the default metalake. Through the MCP protocol, each of these values successfully issued `/api/metalakes/prod/catalogs?details=true` on a server configured with default `prod`. The same middleware handles mutation tools, so malformed tenant input can also direct a write to the default tenant. A truthy non-string such as `42` instead raises `AttributeError` before the audit/error middleware runs; the probe recorded no audit entry. Please distinguish absence from invalid input, reject explicitly supplied non-string values before any REST call, and ensure validation failures pass through error handling and audit. Add protocol tests for these values, including a mutation tool that must issue no REST request. This should be fixed in this PR. ########## mcp-server/mcp_server/core/context.py: ########## @@ -190,23 +255,77 @@ def rest_client(self): "HTTP request omitted Authorization and " "--no-service-identity-fallback is set" ) - return self._default_client + return self._service_client(metalake) - cached = self._clients_by_auth.get(authorization) + key = (authorization, metalake) + cached = self._clients_by_auth.get(key) if cached is not None: - self._clients_by_auth.move_to_end(authorization) + self._clients_by_auth.move_to_end(key) return cached client = RESTClientFactory.create_rest_client( - self._setting.metalake, + metalake, self._setting.gravitino_uri, authorization, ) - self._clients_by_auth[authorization] = client + self._cache_put(key, client) + return client + + def _resolve_metalake(self) -> str: + """Resolve the metalake for the current call, tool argument first. + + Raises ``ValueError`` (an invalid/missing request parameter, mapped + by FastMCP's error middleware to a client-facing "Invalid params" + error rather than an internal-error code) when the call names none and + no startup default (``--metalake``) is configured. The message is + written for the agent that will read it: it names the recovery path so + a model can correct itself instead of just reporting the failure. + """ + metalake = get_request_metalake() or self._setting.metalake + if not metalake: + # Only point at the discovery tool when this deployment actually + # exposes it; a tag filter can hide it, and naming a tool the + # agent cannot call leaves it with no way forward. + recovery = ( + "Call 'list_metalakes' to see the metalakes you can access, " + f"then retry this call with the '{METALAKE_ARGUMENT}' " + "argument set to one of them." + if self._setting.exposes_metalake_discovery() + else f"Retry this call with the '{METALAKE_ARGUMENT}' argument " + "set to the metalake to use, or ask the user which one to use." + ) + raise ValueError(f"No metalake specified. {recovery}") + return metalake + + def _service_client(self, metalake: str): + """Return the service-identity client (static token / OAuth) for ``metalake``.""" + if ( + metalake == self._setting.metalake + and self._default_client is not None + ): + return self._default_client + + key = ("", metalake) + cached = self._clients_by_auth.get(key) + if cached is not None: + self._clients_by_auth.move_to_end(key) + return cached + + client = RESTClientFactory.create_rest_client( + metalake, + self._setting.gravitino_uri, + startup_authorization(self._setting), + auth=_service_auth(self._setting), + ) + self._cache_put(key, client) + return client + + def _cache_put(self, key: "tuple[str, str]", client) -> None: + """Cache a client, evicting (and closing) the oldest past the cap.""" + self._clients_by_auth[key] = client if len(self._clients_by_auth) > _MAX_CACHED_CLIENTS: _, evicted = self._clients_by_auth.popitem(last=False) self._schedule_close(evicted) Review Comment: [P1] Do not close a cached client while a tool call is still using it Eviction immediately schedules `client.close()` without checking for active requests. A slow call can therefore lose its connection when other calls populate the cache. I reproduced this with a real local HTTP server: hold one response open, create clients for 128 other metalakes, and the original request fails with `httpx.ReadError`; its client is already closed. For writes, the backend may have committed before the response is interrupted, leaving an ambiguous outcome and a risk of duplicate execution on retry. The previous per-principal cache already had this hazard, but this PR extends it to service-identity clients and a single caller accessing multiple metalakes. Please track active borrows and defer closing an evicted client until its last user releases it. Add a regression test that triggers eviction while a real request is in flight. This should be fixed in this PR. ########## mcp-server/mcp_server/tools/statistic.py: ########## @@ -22,7 +22,6 @@ def load_statistic_tools(mcp: FastMCP): @mcp.tool(tags={"statistic"}) async def list_statistics_for_metadata( ctx: Context, - metalake_name: str, metadata_type: str, Review Comment: [P2] Preserve a compatibility alias for the statistic tools' metalake_name argument Removing `metalake_name` from both statistic tools immediately breaks existing tool calls and clients with cached schemas. A protocol call using the previous arguments now fails with `Unexpected keyword argument` before reaching Gravitino. I see that the PR description explicitly labels this breaking, but unifying metalake resolution does not require an immediate caller migration. Please accept `metalake_name` as a deprecated alias for these two tools and normalize it into the shared metalake context in the middleware. Reject conflicting old/new argument values explicitly. This keeps the lower-level operations on the new shared resolution path while preserving existing callers. Add old/new/conflicting-argument tests. This should be handled here unless the project has explicitly agreed to a breaking release with coordinated caller migration. ########## mcp-server/mcp_server/core/context.py: ########## @@ -190,23 +255,77 @@ def rest_client(self): "HTTP request omitted Authorization and " "--no-service-identity-fallback is set" ) - return self._default_client + return self._service_client(metalake) - cached = self._clients_by_auth.get(authorization) + key = (authorization, metalake) + cached = self._clients_by_auth.get(key) if cached is not None: - self._clients_by_auth.move_to_end(authorization) + self._clients_by_auth.move_to_end(key) return cached client = RESTClientFactory.create_rest_client( - self._setting.metalake, + metalake, self._setting.gravitino_uri, authorization, ) - self._clients_by_auth[authorization] = client + self._cache_put(key, client) + return client + + def _resolve_metalake(self) -> str: + """Resolve the metalake for the current call, tool argument first. + + Raises ``ValueError`` (an invalid/missing request parameter, mapped + by FastMCP's error middleware to a client-facing "Invalid params" + error rather than an internal-error code) when the call names none and + no startup default (``--metalake``) is configured. The message is + written for the agent that will read it: it names the recovery path so + a model can correct itself instead of just reporting the failure. + """ + metalake = get_request_metalake() or self._setting.metalake + if not metalake: + # Only point at the discovery tool when this deployment actually + # exposes it; a tag filter can hide it, and naming a tool the + # agent cannot call leaves it with no way forward. + recovery = ( + "Call 'list_metalakes' to see the metalakes you can access, " + f"then retry this call with the '{METALAKE_ARGUMENT}' " + "argument set to one of them." + if self._setting.exposes_metalake_discovery() + else f"Retry this call with the '{METALAKE_ARGUMENT}' argument " + "set to the metalake to use, or ask the user which one to use." + ) + raise ValueError(f"No metalake specified. {recovery}") + return metalake + + def _service_client(self, metalake: str): + """Return the service-identity client (static token / OAuth) for ``metalake``.""" + if ( + metalake == self._setting.metalake + and self._default_client is not None + ): + return self._default_client + + key = ("", metalake) + cached = self._clients_by_auth.get(key) + if cached is not None: + self._clients_by_auth.move_to_end(key) + return cached + + client = RESTClientFactory.create_rest_client( + metalake, + self._setting.gravitino_uri, + startup_authorization(self._setting), + auth=_service_auth(self._setting), Review Comment: [P2] Reuse the service OAuth auth object across metalake clients Calling `_service_auth()` for each metalake creates separate `RefreshableBearerAuth` instances for identical service credentials. They share the token-cache key, but each instance owns its own `_token_lock` and 401 retry state. Concurrent cold-cache or expired-token requests across metalakes therefore bypass the existing refresh coalescing and send duplicate token requests to the IdP. With the actual auth/cache implementation and a mocked token fetch, four concurrent calls sharing one auth instance fetched once; calls through four metalake clients fetched four times despite having only one unique cache key. The existing OAuth concurrency tests cover a single auth instance and miss this path. Please construct one service auth object per `GravitinoContext` and reuse it for the default, discovery, and other metalake clients. Add cross-metalake cold-cache, expiry, and 401 concurrency coverage. This is a small regression fix worth including in this PR. -- 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]
