fix: update LiteAdmin MCP integration

README.md

@@ -5,6 +5,14 @@ manage teams and budgets, and inspect spending.
 
 **Currently, this MCP requires a `proxy_admin` identity for all tools, including read-only tools.**
 
+The connector also checks the authenticated credential's current route permissions
+before each action and before returning results. Restricted virtual keys only see
+tools their routes permit. Full Access keys (`allowed_routes` unset or empty) keep
+their existing access, and SSO sessions retain the user's role-based access.
+The gateway must support the `admin_mcp_permissions` response from
+`/user/info?admin_mcp=true`; older gateways are rejected rather than falling back to
+an owner-role check. Upgrade the gateway before updating this connector.
+
 This is a standalone MCP server. It calls your gateway's management API using
 your own proxy-admin credential. Your MCP client supplies the agent and model.
 

SECURITY.md

@@ -10,10 +10,17 @@ to start HTTP mode with a process-wide gateway key. The gateway's live proxy-adm
 role is checked for discovery, before execution and before returning results.
 The caller's own key remains the execution identity and audit actor.
 
+The gateway evaluates the authenticated credential's route permissions and returns
+them with the current identity. The connector requires this permission response,
+filters discovery, and checks the concrete resource route before and after each
+action. It does not infer an SSO session from a client header, token prefix or key
+metadata. Full Access keys keep the gateway's unrestricted-route semantics.
+
 Large sanitized results can be retained as short-lived in-memory snapshots.
 Handles are unguessable and bound to both a keyed digest of the requesting
 credential and the verified principal; another admin cannot read them. Each
-read checks the current role before and after access. Raw caller credentials
+read checks the current role and the source operation's current route permission.
+Raw caller credentials
 are not stored in the cache. Secret-returning key operations and results with
 recognizable credentials bypass it and return their complete response inline.
 Snapshot reads never re-execute the original operation. Expiration and capacity

src/litellm_admin_mcp/gateway.py

@@ -1,5 +1,6 @@
 """Authenticated, non-retrying calls to a fixed LiteLLM management API."""
 import asyncio
+from dataclasses import dataclass
 import json
 import re
 import time
@@ -34,6 +35,12 @@ def __init__(self, message: str, status: int = 400):
         self.status = status
 
 
+@dataclass(frozen=True)
+class Authorization:
+    principal: str
+    tools: frozenset[str]
+
+
 def valid_credential(value: str) -> str:
     if not value or len(value) > 8192 or any(not 33 <= ord(c) <= 126 for c in value):
         raise AdminError("A personal LiteLLM admin credential is required.", 401)
@@ -117,8 +124,17 @@ async def _request(self, method: str, path: str, credential: str, *,
             raise AdminError("Gateway returned an invalid JSON response." + uncertain, 502) from None
 
     async def authorize(self, credential: str) -> str:
+        return (await self._authorization(credential)).principal
+
+    async def _authorization(self, credential: str, *, tool: str | None = None,
+                             path: str | None = None) -> Authorization:
         try:
-            data = await self._request("GET", "/user/info", credential)
+            params = {"admin_mcp": "true"}
+            if tool is not None:
+                params["mcp_tool"] = tool
+            if path is not None:
+                params["mcp_path"] = path
+            data = await self._request("GET", "/user/info", credential, params=params)
         except AdminError as exc:
             if exc.status != 404:
                 raise
@@ -129,13 +145,21 @@ async def authorize(self, credential: str) -> str:
                 or data.get("user_id") != user["user_id"] or user.get("blocked")
                 or user.get("deleted") or user.get("is_active") is False):
             raise AdminError("Only current LiteLLM proxy admins may use these tools.", 403)
-        return user["user_id"]
+        permissions = data.get("admin_mcp_permissions")
+        if (not isinstance(permissions, dict) or not isinstance(permissions.get("tools"), list)
+                or any(not isinstance(name, str) for name in permissions["tools"])):
+            raise AdminError("Gateway cannot verify this credential's admin permissions; upgrade LiteLLM.", 403)
+        if tool is not None and permissions.get("action_allowed") is not True:
+            raise AdminError("This credential is not allowed to use this admin tool.", 403)
+        return Authorization(user["user_id"], frozenset(permissions["tools"]))
 
     async def tools(self, credential: str) -> dict[str, types.Tool]:
-        await self.authorize(credential)
+        authorization = await self._authorization(credential)
         canonical = await self._load_tools(credential)
         visible = {}
         for name, tool in canonical.items():
+            if name not in authorization.tools:
+                continue
             visible[name] = (tool.model_copy(update={"input_schema": discovery_schema(tool.input_schema, name),
                 "description": tool.description + " Call describe_admin_tool for the complete argument schema before use."})
                 if self.config.schema_mode == "discovery" else tool)
@@ -186,11 +210,14 @@ async def _load_tools(self, credential: str) -> dict[str, types.Tool]:
     async def call(self, name: str, arguments: dict, credential: str) -> Any:
         # No redundant call to tools(): execution and disclosure each have a
         # fresh authorization check. Roles/identities are never cached.
-        principal = await self.authorize(credential)
+        authorization = await self._authorization(credential)
+        principal = authorization.principal
         loaded_at = self._loaded_at
         available = await self._load_tools(credential)
         if name not in available and name not in HELPERS:
             raise AdminError("This tool is not enabled for this connector.", 403)
+        if name not in authorization.tools and name not in HELPERS:
+            raise AdminError("This credential is not allowed to use this admin tool.", 403)
         try:
             schema = HELPERS[name][0] if name in HELPERS else available[name].input_schema
             jsonschema.Draft202012Validator(schema).validate(arguments)
@@ -206,13 +233,15 @@ async def call(self, name: str, arguments: dict, credential: str) -> Any:
         view = arguments.get("response", {}).get("view", self.config.response_view)
         return await self._execute(name, arguments, credential, principal, view)
 
-    async def _verify(self, credential: str, principal: str):
-        if await self.authorize(credential) != principal:
+    async def _verify(self, credential: str, principal: str, *, tool: str | None = None,
+                      path: str | None = None):
+        if (await self._authorization(credential, tool=tool, path=path)).principal != principal:
             raise AdminError("Admin identity changed during the request. Inspect gateway state before retrying.", 403)
 
-    def _view(self, name: str, value: Any, credential: str, principal: str, view: str) -> Any:
+    def _view(self, name: str, value: Any, credential: str, principal: str, view: str,
+              path: str | None = None) -> Any:
         if view == "compact" and len(encode(value)) > PAGE_BYTES:
-            token = self.results.put(name, value, credential, principal)
+            token = self.results.put(name, value, credential, principal, route=path)
             if token is not None:
                 return page(self.results.get(token, credential, principal), token)
         return value
@@ -224,7 +253,7 @@ async def _helper(self, name, arguments, credential, principal, available):
                 if target not in available:
                     raise AdminError("This tool is not enabled for this connector.", 403)
                 result = available[target].model_dump(by_alias=True, exclude_none=True, mode="json")
-                await self._verify(credential, principal)
+                await self._verify(credential, principal, tool=target)
                 return self._view("schema:" + target, result, credential, principal,
                                   arguments.get("response", {}).get("view", "full"))
             token = arguments["result_id"]
@@ -237,7 +266,7 @@ async def _helper(self, name, arguments, credential, principal, available):
                 result = select(json.loads(item.raw), arguments.get("path", ""))
             else:
                 result = page(item, token, **{k: v for k, v in arguments.items() if k not in {"result_id", "view"}})
-            await self._verify(credential, principal)
+            await self._verify(credential, principal, tool=item.source.removeprefix("schema:"), path=item.route)
             return result
         except ResultError as exc:
             raise AdminError(str(exc), 404) from None
@@ -264,8 +293,9 @@ async def _execute(self, name, arguments, credential, principal, view):
                 headers["Content-Type"] = "application/json"
             else:
                 body = {"json": arguments["body"]}
+        await self._verify(credential, principal, tool=name, path=path)
         result = await self._request(item.method, path, credential, params=query, **body,
             read_only=item.read_only, max_bytes=MAX_RESULT_BYTES,
             headers=headers)
-        await self._verify(credential, principal)
-        return self._view(name, safe_result(result, credential), credential, principal, view)
+        await self._verify(credential, principal, tool=name, path=path)
+        return self._view(name, safe_result(result, credential), credential, principal, view, path=path)

src/litellm_admin_mcp/results.py

@@ -107,6 +107,7 @@ class Snapshot:
     raw: bytes
     deadline: float
     expires_at: float
+    route: str | None = None
 
 
 class ResultStore:
@@ -130,12 +131,12 @@ def clear(self):
         self._items.clear()
         self._bytes = 0
 
-    def put(self, source: str, value: Any, credential: str, principal: str) -> str | None:
+    def put(self, source: str, value: Any, credential: str, principal: str, *, route: str | None = None) -> str | None:
         self._expire()
         raw = encode(value)
         # New credentials are delivered in the original full response, never
         # retained for later reads. Also exclude recognizable keys in log data.
-        if source in SECRET_RESULTS or contains_credential(value, source=source):
+        if source in SECRET_RESULTS or contains_credential(value, source=source) or contains_credential(route):
             return None
         owner = self._owner(credential)
         if (len(raw) > MAX_RESULT_BYTES or self._bytes + len(raw) > self.max_bytes
@@ -145,7 +146,7 @@ def put(self, source: str, value: Any, credential: str, principal: str) -> str |
             return None
         token = secrets.token_urlsafe(24)
         self._items[token] = Snapshot(owner, principal, source, raw,
-            time.monotonic() + self.ttl, time.time() + self.ttl)
+            time.monotonic() + self.ttl, time.time() + self.ttl, route)
         self._bytes += len(raw)
         return token
 

tests/conftest.py

@@ -8,6 +8,7 @@
 import pytest
 
 from litellm_admin_mcp.config import Config
+from litellm_admin_mcp.catalog import OPERATIONS
 from litellm_admin_mcp.gateway import Gateway
 from litellm_admin_mcp.server import create_server
 
@@ -53,6 +54,8 @@ def __init__(self):
         self.revoke = False
         self.revoked = set()
         self.spec = SPEC
+        self.denied_tools = set()
+        self.denied_paths = set()
 
     def handle(self, request):
         self.requests.append(request)
@@ -61,7 +64,15 @@ def handle(self, request):
             return httpx2.Response(404, json={"error": {"message": "User None not found"}})
         if request.url.path == "/user/info":
             role = "proxy_admin" if token in {"alice-admin", "bob-admin"} and token not in self.revoked else "internal_user"
-            return httpx2.Response(200, json={"user_id": token, "user_info": {"user_id": token, "user_role": role}})
+            tools = [op.name for op in OPERATIONS if op.name not in self.denied_tools]
+            return httpx2.Response(200, json={
+                "user_id": token, "user_info": {"user_id": token, "user_role": role},
+                "admin_mcp_permissions": {
+                    "tools": tools,
+                    "action_allowed": request.url.params.get("mcp_tool") in tools
+                    and request.url.params.get("mcp_path") not in self.denied_paths,
+                },
+            })
         if request.url.path == "/openapi.json":
             return httpx2.Response(200, json=self.spec)
         assert token in {"alice-admin", "bob-admin"}

tests/test_connector.py

@@ -112,6 +112,43 @@ def test_http_server_refuses_shared_process_credential():
         create_http_app(Gateway(Config("https://gateway.example.com", api_key="shared")))
 
 
+async def test_key_permissions_filter_discovery_and_block_direct_calls(stub):
+    stub.denied_tools = {"create_key", "add_model"}
+    async with session_for(stub) as session:
+        assert {t.name for t in (await session.list_tools()).tools} == {
+            "list_keys", "get_model", "reset_key_spend", *HELPERS,
+        }
+        assert (await session.call_tool("create_key", {"body": {"key_alias": "forbidden"}})).is_error
+        assert (await session.call_tool("describe_admin_tool", {"name": "create_key"})).is_error
+        assert data(await session.call_tool("list_keys", {})) == {"keys": []}
+    assert not stub.writes
+
+
+async def test_concrete_resource_permission_is_checked_before_write(stub):
+    stub.denied_paths = {"/key/other-key/reset_spend"}
+    async with session_for(stub) as session:
+        assert data(await session.call_tool("reset_key_spend", {"path": {"key": "my-key"}})) == {"status": "ok"}
+        assert (await session.call_tool("reset_key_spend", {"path": {"key": "other-key"}})).is_error
+    assert [request.url.path for request in stub.writes] == ["/key/my-key/reset_spend"]
+
+
+async def test_legacy_gateway_without_credential_permissions_fails_closed(stub):
+    original = stub.handle
+    def handle(request):
+        response = original(request)
+        if request.url.path == "/user/info":
+            body = response.json()
+            body.pop("admin_mcp_permissions")
+            return httpx2.Response(200, json=body)
+        return response
+    stub.handle = handle
+    async with session_for(stub) as session:
+        with pytest.raises(MCPError, match="upgrade LiteLLM"):
+            await session.list_tools()
+        assert (await session.call_tool("create_key", {"body": {"key_alias": "forbidden"}})).is_error
+    assert not stub.writes
+
+
 def test_recursive_refs_and_overridden_parameters_have_valid_schema():
     import jsonschema
     spec = {"components": {"schemas": {"Node": {"type": "object", "properties": {

tests/test_progressive.py

@@ -227,7 +227,7 @@ async def test_warm_call_retains_pre_post_checks_without_discovery_auth(stub):
         await session.list_tools()
         stub.requests.clear()
         await session.call_tool("list_keys", {})
-        assert [r.url.path for r in stub.requests] == ["/user/info", "/key/list", "/user/info"]
+        assert [r.url.path for r in stub.requests] == ["/user/info", "/user/info", "/key/list", "/user/info"]
 
 
 async def test_identity_change_during_cold_discovery_prevents_write(stub):
@@ -239,7 +239,7 @@ def handle(request):
         if request.url.path == "/openapi.json":
             changed = True
         elif request.url.path == "/user/info" and changed:
-            return httpx2.Response(200, json={"user_id": "new-identity", "user_info": {"user_id": "new-identity", "user_role": "proxy_admin"}})
+            return httpx2.Response(200, json={**result.json(), "user_id": "new-identity", "user_info": {"user_id": "new-identity", "user_role": "proxy_admin"}})
         return result
     stub.handle = handle
     async with session_for(stub) as session:
@@ -324,7 +324,10 @@ def handle(request):
         if request.url.path == "/openapi.json":
             return httpx2.Response(200, json=spec)
         if request.url.path == "/user/info" and "litellm-changed-by" not in request.headers:
-            return httpx2.Response(200, json={"user_id": "admin", "user_info": {"user_id": "admin", "user_role": "proxy_admin"}})
+            return httpx2.Response(200, json={
+                "user_id": "admin", "user_info": {"user_id": "admin", "user_role": "proxy_admin"},
+                "admin_mcp_permissions": {"tools": [op.name for op in OPERATIONS], "action_allowed": True},
+            })
         calls.append(request)
         return httpx2.Response(200, json=result)
     async with httpx2.AsyncClient(transport=httpx2.MockTransport(handle)) as client:
@@ -357,3 +360,29 @@ def handle(request):
 def test_unknown_modes_fail_instead_of_silently_changing_behavior(config):
     with pytest.raises(ValueError):
         Config("https://gateway.example.com", **config)
+
+
+async def test_cached_result_is_denied_after_source_permission_is_revoked(stub):
+    stub.keys = [{"data": "private" * 4000}]
+    async with session_for(stub) as session:
+        result = data(await session.call_tool("list_keys", {}))
+        stub.denied_tools.add("list_keys")
+        denied = await session.call_tool("read_admin_result", result["full_read"])
+        assert denied.is_error
+        assert "private" not in str(denied)
+    assert sum(r.url.path == "/key/list" for r in stub.requests) == 1
+
+
+async def test_cached_result_rechecks_its_original_resource_route(stub):
+    original = stub.handle
+    def handle(request):
+        response = original(request)
+        return httpx2.Response(200, json={"history": "private" * 4000}) if request.url.path.endswith("/reset_spend") else response
+    stub.handle = handle
+    async with session_for(stub) as session:
+        result = data(await session.call_tool("reset_key_spend", {"path": {"key": "fixture-key-id"}}))
+        stub.denied_paths.add("/key/fixture-key-id/reset_spend")
+        denied = await session.call_tool("read_admin_result", result["full_read"])
+        assert denied.is_error
+        assert "private" not in str(denied)
+    assert len(stub.writes) == 1