You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tooling.apache.org by gs...@apache.org on 2026/10/13 03:44:56 UTC
(tooling-llmao) 04/09: tiers: mint keys on the right tier, with the tier's current limits
This is an automated email from the ASF dual-hosted git repository.
gstein pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tooling-llmao.git
commit 33724d8bb23e8e91a6df0e4547667c16edbb4c51
Author: Greg Stein <gs...@gmail.com>
AuthorDate: Mon Oct 12 22:27:06 2026 -0500
tiers: mint keys on the right tier, with the tier's current limits
Four bugs in how a tier reaches a key.
A tier change never reached new keys. The Tiers page tells an admin that
rate limits "apply to keys issued from now on", but minting read the limits
from config alone and ignored what the admin had set. KeyManager.create now
asks TierManager, through the new key_limits(), for the tier as it stands.
Automation keys were minted on the project tier. Minting defaulted to it,
and create_automation_key did not say otherwise. It names the service tier
now, which is what that tier is for. The tier argument has no default any
more, so a caller that forgets fails rather than quietly getting project
limits. This raises the limits on automation keys created from here on.
A key request could name any tier. It arrives in a form post and was stored
unchecked, so an invented name would later mint a key on the default limits
under a tier no page lists. request_key and set_tier refuse a tier that is
not one of the three.
A key's recorded tier could disagree with its limits. Metadata that already
named a tier kept its name while the limits came from the argument. The
argument decides both.
tests/test_keys.py is the first test of the real KeyManager: what
key/generate is sent. The seam tests check which tier each path uses and
that unknown tiers are refused.
Not addressed: keys already minted keep the limits they were given.
Co-Authored-By: Claude Opus 5.5 <no...@anthropic.com>
---
llmao/proxy/client.py | 34 +++++++++++++++-----
llmao/seam.py | 18 +++++++++--
tests/mock_backend.py | 3 ++
tests/test_access_requests.py | 10 ++++++
tests/test_keys.py | 75 +++++++++++++++++++++++++++++++++++++++++++
tests/test_phase1.py | 55 +++++++++++++++++++++++++++++++
6 files changed, 184 insertions(+), 11 deletions(-)
diff --git a/llmao/proxy/client.py b/llmao/proxy/client.py
index 12263ed..9e33b07 100644
--- a/llmao/proxy/client.py
+++ b/llmao/proxy/client.py
@@ -54,6 +54,7 @@ GRANTOR_FREE_TIER = "Free Tier"
TIER_FREE = "free_tier"
TIER_PROJECT = "project"
TIER_SERVICE = "service"
+TIERS = (TIER_FREE, TIER_PROJECT, TIER_SERVICE)
# LiteLLM fields these map to. Named here so the mapping is in one place
# rather than spread across call sites.
@@ -316,7 +317,7 @@ class LiteLLMBackend:
self.tiers = TierManager(cfg, memory)
self.teams = TeamManager(cfg, self._transport)
self.metering = UsageManager(self._transport, fleet, self.teams, self.tiers)
- self.keys = KeyManager(cfg, self._transport, self.teams)
+ self.keys = KeyManager(self._transport, self.teams, self.tiers)
async def aclose(self) -> None:
await self._transport.aclose()
@@ -353,8 +354,8 @@ class LiteLLMBackend:
class KeyManager:
"""LiteLLM's virtual keys: listing, minting, revoking, and what each may do.
- Handed the config, the transport and the teams -- every key is minted
- into its project's team.
+ Handed the transport, the teams -- every key is minted into its
+ project's team -- and the tiers, which say what a new key may do.
"""
# The team a key with no project is minted into. LiteLLM wants every
@@ -362,10 +363,10 @@ class KeyManager:
# is from when tier overrides were kept on this team's metadata.)
PROJECTLESS_TEAM = "llmao-entitlement-defaults"
- def __init__(self, cfg: Any, transport: Any, teams: Any):
- self._cfg = cfg
+ def __init__(self, transport: Any, teams: Any, tiers: Any):
self._transport = transport
self.teams = teams
+ self.tiers = tiers
async def update_entitlement(
self,
@@ -444,10 +445,15 @@ class KeyManager:
project: str,
purpose: str,
user: str | None = None,
+ tier: str,
metadata: dict | None = None,
- tier: str = TIER_PROJECT,
token_cap: int | None = None,
) -> CreatedKey:
+ """Mint a key on `tier`, which decides its rate limits.
+
+ The tier has no default. It used to be the project tier, and a
+ caller that forgot to say got project limits for a pipeline key.
+ """
project = (project or "").strip()
if not project:
raise BackendUnavailableError("create_key requires project")
@@ -463,7 +469,7 @@ class KeyManager:
# price and no budget can ever be exhausted by them. A key minted
# without limits can saturate the fleet, and LiteLLM's default is
# unlimited.
- meta.setdefault("tier", tier)
+ meta["tier"] = tier
if token_cap is not None:
# A per-key allowance, set by an admin when approving. Only a
# projectless service key really needs one: a project key draws
@@ -471,7 +477,10 @@ class KeyManager:
# a key belonging to neither has nothing to draw on, and
# unattended automation would consume a week in an afternoon.
meta["token_cap"] = int(token_cap)
- limits = resolve_entitlement(self._cfg, tier)
+ # From the tiers rather than from config alone: an admin who has
+ # changed a tier's limits is told they apply to keys issued from
+ # then on, and this is where that is made true.
+ limits = await self.tiers.key_limits(tier)
payload: dict[str, Any] = {
"team_id": team_id,
@@ -1140,6 +1149,15 @@ class TierManager:
effective.update({k: v for k, v in stored.items() if v is not None})
return effective
+ async def key_limits(self, tier: str) -> dict[str, Any]:
+ """The rate limits a key minted on this tier is given.
+
+ Config with the admin's changes on top, like limits(), but only what
+ LiteLLM holds on a key. The token cap is not one of those: it is
+ metered here, against usage, and LiteLLM would not know the field.
+ """
+ return {k: v for k, v in (await self.limits(tier)).items() if k in _LIMIT_FIELDS}
+
async def set_limits(
self,
tier: str,
diff --git a/llmao/seam.py b/llmao/seam.py
index 08bf2f2..71dffe8 100644
--- a/llmao/seam.py
+++ b/llmao/seam.py
@@ -33,9 +33,9 @@ from typing import Any
from . import access_requests as _access
from .proxy.client import (
GRANTOR_FREE_TIER,
- TIER_FREE,
TIER_PROJECT,
TIER_SERVICE,
+ TIERS,
BackendUnavailableError,
CreatedKey,
KeyInfo,
@@ -319,12 +319,18 @@ class Seam:
raise AuthzError("Say what the key is for.")
if project and project not in (identity.committees or []):
raise AuthzError(f"You are not on the {project} PMC.")
+ tier = (tier or "").strip() or TIER_SERVICE
+ if tier not in TIERS:
+ # The tier arrives in a form post and ends up deciding a key's
+ # limits. An unknown one would mint a key on the defaults, under
+ # a name no page lists.
+ raise AuthzError(f"There is no {tier} tier.")
req = _access.AccessRequest.new(
kind=_access.KIND_KEY,
requester=identity.uid,
project=project,
reason=reason,
- wanted={"purpose": purpose, "tier": tier or TIER_SERVICE},
+ wanted={"purpose": purpose, "tier": tier},
)
return await self._backend.requests.add(req)
@@ -339,7 +345,7 @@ class Seam:
if not identity.is_site_admin:
raise AuthzError("Tiers are limited to site admins.")
out = []
- for tier in (TIER_FREE, TIER_PROJECT, TIER_SERVICE):
+ for tier in TIERS:
out.append({"tier": tier, **(await self._backend.tiers.limits(tier))})
return out
@@ -351,6 +357,8 @@ class Seam:
"""
if not identity.is_site_admin:
raise AuthzError("Changing tiers is limited to site admins.")
+ if tier not in TIERS:
+ raise AuthzError(f"There is no {tier} tier.")
named = {k: v for k, v in limits.items() if v is not None}
if not named:
raise AuthzError("Nothing to change.")
@@ -634,6 +642,8 @@ class Seam:
project=project,
purpose=purpose,
user=identity.uid,
+ # A person's key in a project draws on the project's share.
+ tier=TIER_PROJECT,
)
@require_admin
@@ -650,6 +660,8 @@ class Seam:
purpose=purpose,
user=None,
metadata={"created_by": identity.uid},
+ # No person at the keyboard: this is what the service tier is for.
+ tier=TIER_SERVICE,
)
async def revoke_key(self, identity: Identity, token_id: str) -> None:
diff --git a/tests/mock_backend.py b/tests/mock_backend.py
index a4f889e..f756bf9 100644
--- a/tests/mock_backend.py
+++ b/tests/mock_backend.py
@@ -116,8 +116,10 @@ class MockBackend:
*,
project: str,
purpose: str,
+ tier: str,
user: str | None = None,
metadata: dict | None = None,
+ token_cap: int | None = None,
) -> CreatedKey:
project = (project or "").strip()
purpose = (purpose or "").strip()
@@ -126,6 +128,7 @@ class MockBackend:
token_id = f"tok-{uuid.uuid4().hex[:16]}"
meta = dict(metadata or {})
meta["project"] = project
+ meta["tier"] = tier
if purpose:
meta["purpose"] = purpose
row = {
diff --git a/tests/test_access_requests.py b/tests/test_access_requests.py
index 940f0d5..c3fbaa0 100644
--- a/tests/test_access_requests.py
+++ b/tests/test_access_requests.py
@@ -434,3 +434,13 @@ def test_changing_one_tier_leaves_another_alone():
def test_tier_limits_fall_back_to_config_when_litellm_is_down():
"""A page must still have limits to show during a proxy restart."""
assert _run(_tiers(down=True)[0].limits("free_tier"))["token_cap"] == 1_000_000
+
+
+def test_a_new_keys_limits_include_what_an_admin_has_changed():
+ """The Tiers page says a change applies to keys issued from then on."""
+ t, _ = _tiers()
+ assert _run(t.key_limits("service")) == {"rpm_limit": 60, "tpm_limit": 60000, "max_parallel_requests": 4}
+ _run(t.set_limits("service", rpm_limit=10, token_cap=5_000_000))
+ limits = _run(t.key_limits("service"))
+ assert limits == {"rpm_limit": 10, "tpm_limit": 60000, "max_parallel_requests": 4}
+ assert "token_cap" not in limits, "LiteLLM holds rate limits on a key, not our token cap"
diff --git a/tests/test_keys.py b/tests/test_keys.py
new file mode 100644
index 0000000..3ca1844
--- /dev/null
+++ b/tests/test_keys.py
@@ -0,0 +1,75 @@
+# 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.
+"""Minting a key: what LiteLLM is actually sent.
+
+The seam tests run against MockBackend and never reach this code.
+"""
+
+import asyncio
+
+from easydict import EasyDict as edict # noqa: N813
+
+from llmao.proxy.client import KeyManager
+from tests.fake_transport import FakeTransport
+
+
+class _Teams:
+ """Stands in for TeamManager: every project has a team."""
+
+ async def ensure(self, project):
+ return edict(team_id=f"team-{project}")
+
+
+class _Tiers:
+ """Stands in for TierManager: the service tier has been cut to 10 a minute."""
+
+ async def key_limits(self, tier):
+ return {"rpm_limit": 10 if tier == "service" else 120, "tpm_limit": 60000}
+
+
+def _keys():
+ def litellm(call):
+ # key/generate answers with the key it made
+ return {"key": "sk-secret", "token": "tok1", **call.json}
+
+ return KeyManager(FakeTransport(litellm), _Teams(), _Tiers())
+
+
+def _sent(keys):
+ return keys._transport.calls[-1]
+
+
+def test_a_key_is_minted_with_its_tiers_limits():
+ keys = _keys()
+ created = asyncio.run(keys.create(project="airflow", purpose="ci", user="jdoe", tier="service"))
+ sent = _sent(keys)
+ assert (sent.method, sent.path) == ("POST", "key/generate")
+ assert sent.json.rpm_limit == 10, "the tier's limit as it stands now, not as config had it"
+ assert sent.json.tpm_limit == 60000
+ assert sent.json.team_id == "team-airflow"
+ assert sent.json.user_id == "jdoe"
+ assert created.secret == "sk-secret"
+ assert created.info.project == "airflow"
+
+
+def test_the_tier_asked_for_is_the_tier_recorded():
+ """Metadata naming one tier while the limits came from another is how an audit goes wrong."""
+ keys = _keys()
+ asyncio.run(keys.create(project="airflow", purpose="ci", user="jdoe", tier="project", metadata={"tier": "service"}))
+ sent = _sent(keys)
+ assert sent.json.metadata.tier == "project"
+ assert sent.json.rpm_limit == 120
diff --git a/tests/test_phase1.py b/tests/test_phase1.py
index 8a2ddb7..70b88d5 100644
--- a/tests/test_phase1.py
+++ b/tests/test_phase1.py
@@ -177,3 +177,58 @@ def test_site_admin_sees_any_project():
assert len(rows) == 1
asyncio.run(run())
+
+
+# ---------------------------------------------------------------------------
+# Which tier a key is minted on
+#
+# The tier decides a key's rate limits, and those are the only control there
+# is on self-hosted capacity. Minting used to default to the project tier,
+# so a path that did not say got project limits whatever it was for.
+# ---------------------------------------------------------------------------
+
+
+def _minted_tier(backend):
+ return backend._data["keys"][-1]["metadata"]["tier"]
+
+
+def test_a_persons_project_key_is_on_the_project_tier():
+ async def run():
+ seam, _, backend = _seam()
+ await seam.create_personal_key(Identity(uid="jdoe", projects=["airflow"], committees=[]), "airflow", "cli")
+ assert _minted_tier(backend) == "project"
+
+ asyncio.run(run())
+
+
+def test_an_automation_key_is_on_the_service_tier():
+ """No person at the keyboard is what the service tier is for."""
+
+ async def run():
+ seam, _, backend = _seam()
+ await seam.create_automation_key(Identity(uid="chair", projects=[], committees=["airflow"]), "airflow", "ci")
+ assert _minted_tier(backend) == "service"
+
+ asyncio.run(run())
+
+
+def test_a_key_request_cannot_name_a_tier_that_does_not_exist():
+ """The tier comes from a form post and ends up deciding a key's limits."""
+
+ async def run():
+ seam, _, _ = _seam()
+ ident = Identity(uid="jdoe", projects=[], committees=[])
+ with pytest.raises(AuthzError, match="no platinum tier"):
+ await seam.request_key(ident, purpose="ci", tier="platinum", reason="because")
+
+ asyncio.run(run())
+
+
+def test_an_admin_cannot_change_a_tier_that_does_not_exist():
+ async def run():
+ seam, _, _ = _seam()
+ root = Identity(uid="root", projects=[], committees=[], is_site_admin=True)
+ with pytest.raises(AuthzError, match="no platinum tier"):
+ await seam.set_tier(root, "platinum", rpm_limit=1)
+
+ asyncio.run(run())
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@tooling.apache.org
For additional commands, e-mail: commits-help@tooling.apache.org