test(e2e): jwt auto_register map-existing-key repro

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

tests/e2e/coverage_registry/other.yaml

@@ -54,3 +54,6 @@
 
 - {id: other.auth.jwt.wrong_issuer_denied, module: other, tier: P0, area: auth, assertions: [wrong_issuer_denied], source: "auth/handle_jwt.py", rationale: "A signed token with the correct audience and an unexpected issuer is rejected"}
 - {id: other.auth.jwt.wrong_audience_denied, module: other, tier: P0, area: auth, assertions: [wrong_audience_denied], source: "auth/handle_jwt.py", rationale: "A signed token from the trusted issuer intended for another app is rejected"}
+- {id: other.auth.jwt.auto_register_maps_existing_key, module: other, tier: P0, area: auth, assertions: [maps_existing_key], source: "user_api_key_auth.py _auto_register_jwt_mapping", rationale: "With auto_register_map_existing_key: true, the first JWT call of a user who already owns a key writes the sub-claim mapping to that existing key hash and mints nothing; the spend row lands on the pre-existing key (LIT-5378)", fail_before_fix: proven}
+- {id: other.auth.jwt.auto_register_mints_when_keyless, module: other, tier: P0, area: auth, assertions: [mints_when_keyless], source: "user_api_key_auth.py _auto_register_jwt_mapping", rationale: "With auto_register_map_existing_key: true, a user with no keys still gets exactly one minted key and a sub-claim mapping on their first JWT call (LIT-5378)", fail_before_fix: proven}
+- {id: other.auth.jwt.auto_register_default_mints, module: other, tier: P0, area: auth, assertions: [default_mints], source: "user_api_key_auth.py _auto_register_jwt_mapping", rationale: "Without auto_register_map_existing_key, auto_register keeps the current behavior: it mints a second key for a user who already has one and bills the minted key (LIT-5378)", fail_before_fix: proven}

tests/e2e/e2e_config.py

@@ -7,6 +7,7 @@ environment so the same tests run against localhost or a deployed proxy.
 from __future__ import annotations
 
 import os
+import socket
 import time
 import uuid
 from pathlib import Path
@@ -15,6 +16,7 @@ from typing import Final
 from dotenv import load_dotenv
 from fixture_mode import deterministic_marker, parse_fixture_mode, registration_owner
 from provider_edge import provider_edge_api_base
+from pydantic import TypeAdapter
 
 # Local runs keep provider / DataDog keys in tests/e2e/.env (see CONTRIBUTING.md).
 # Compose injects them into the proxy container, but pytest on the host does not
@@ -233,6 +235,15 @@ def unique_marker() -> str:
     return uuid.uuid4().hex[:12]
 
 
+INHERITED_ENV_PREFIXES: Final = ("REDIS_", "MICROSOFT_", "GOOGLE_", "GENERIC_", "PROXY_")
+
+
+def available_port() -> int:
+    with socket.socket() as listener:
+        listener.bind(("127.0.0.1", 0))
+        return TypeAdapter(tuple[str, int]).validate_python(listener.getsockname())[1]
+
+
 def settle_propagation(written_at: float) -> None:
     """Block until PROPAGATION_TIMEOUT has elapsed since `written_at`, a
     `time.monotonic()` stamp taken the moment a control-plane write returned.

tests/e2e/mcp/oauth_gateway.py

@@ -8,7 +8,6 @@ The optional live edge measures headers without recording credentials or bodies.
 from __future__ import annotations
 
 import os
-import socket
 import subprocess
 import sys
 import threading
@@ -20,13 +19,12 @@ from pathlib import Path
 from typing import Final
 
 import psycopg
+from e2e_config import INHERITED_ENV_PREFIXES, available_port
 from e2e_http import NoBody
 from idp import Keycloak, stop_process_group
 from proxy_client import ProxyClient, build_proxy_client
 from psycopg.rows import class_row
-from pydantic import BaseModel, SecretStr, TypeAdapter, ValidationError
-
-INHERITED_ENV_PREFIXES: Final = ("REDIS_", "MICROSOFT_", "GOOGLE_", "GENERIC_", "PROXY_")
+from pydantic import BaseModel, SecretStr, ValidationError
 
 
 class StoredOAuth(BaseModel):
@@ -101,12 +99,6 @@ class OAuthObservation:
         assert all(not item[2] for item in snapshot), "gateway bearer leaked to the upstream"
 
 
-def available_port() -> int:
-    with socket.socket() as listener:
-        listener.bind(("127.0.0.1", 0))
-        return TypeAdapter(tuple[str, int]).validate_python(listener.getsockname())[1]
-
-
 @dataclass(slots=True)
 class OAuthGateway:
     base_url: str

tests/e2e/models.py

@@ -1278,6 +1278,7 @@ class UserNewBody(BaseModel):
 
 class UserNewResponse(BaseModel):
     user_id: str
+    key: str | None = None
 
 
 class UserUpdateBody(BaseModel):
@@ -1321,6 +1322,40 @@ class UserListResponse(BaseModel):
     total: int
 
 
+class UserKeyRow(BaseModel):
+    token: str
+    key_alias: str | None = None
+
+
+class UserInfoWithKeysResponse(BaseModel):
+    user_id: str | None = None
+    keys: list[UserKeyRow] = []
+
+
+class JwtKeyMappingRow(BaseModel):
+    id: str
+    jwt_claim_name: str
+    jwt_claim_value: str
+    created_by: str | None = None
+
+
+class JwtKeyMappingListParams(BaseModel):
+    size: int = 100
+
+
+class JwtKeyMappingListResponse(BaseModel):
+    mappings: list[JwtKeyMappingRow]
+    total_count: int
+
+
+class JwtKeyMappingDeleteBody(BaseModel):
+    id: str
+
+
+class JwtKeyMappingDeleteResponse(BaseModel):
+    status: str
+
+
 class OrgNewBody(BaseModel):
     organization_alias: str
     models: list[str] = []

tests/e2e/other/other_client.py

@@ -18,10 +18,18 @@ from dataclasses import dataclass
 from e2e_http import NoBody, ProbeResult, Result
 from idp import Keycloak, keycloak_from_env
 from models import (
+    JwtKeyMappingDeleteBody,
+    JwtKeyMappingDeleteResponse,
+    JwtKeyMappingListParams,
+    JwtKeyMappingListResponse,
     ReadinessDetailsResponse,
     ReadinessResponse,
+    UserInfoParams,
+    UserInfoWithKeysResponse,
     UserListParams,
     UserListResponse,
+    UserNewBody,
+    UserNewResponse,
 )
 from proxy_client import ProxyClient
 
@@ -66,6 +74,44 @@ class OtherClient:
             response_type=ReadinessDetailsResponse,
         )
 
+    def user_new(self, body: UserNewBody) -> Result[UserNewResponse]:
+        """POST /user/new under the master key: seed the litellm user a JWT
+        `sub` claim resolves to, before that token ever reaches the proxy."""
+        return self.proxy.transport.post(
+            "/user/new",
+            headers=self.proxy.transport.master,
+            json=body,
+            response_type=UserNewResponse,
+        )
+
+    def user_info(self, user_id: str) -> Result[UserInfoWithKeysResponse]:
+        """GET /user/info under the master key. Only the user's key rows are
+        modelled: `token` is the stored key hash, never the plaintext key."""
+        return self.proxy.transport.get(
+            "/user/info",
+            headers=self.proxy.transport.master,
+            params=UserInfoParams(user_id=user_id),
+            response_type=UserInfoWithKeysResponse,
+        )
+
+    def jwt_mapping_list(self) -> Result[JwtKeyMappingListResponse]:
+        """GET /jwt/key/mapping/list under the master key."""
+        return self.proxy.transport.get(
+            "/jwt/key/mapping/list",
+            headers=self.proxy.transport.master,
+            params=JwtKeyMappingListParams(size=100),
+            response_type=JwtKeyMappingListResponse,
+        )
+
+    def jwt_mapping_delete(self, mapping_id: str) -> Result[JwtKeyMappingDeleteResponse]:
+        """POST /jwt/key/mapping/delete under the master key."""
+        return self.proxy.transport.post(
+            "/jwt/key/mapping/delete",
+            headers=self.proxy.transport.master,
+            json=JwtKeyMappingDeleteBody(id=mapping_id),
+            response_type=JwtKeyMappingDeleteResponse,
+        )
+
     def list_users_as(self, key: str) -> Result[UserListResponse]:
         """GET /user/list under `key`. Admin-only, so it doubles as the master
         key's authorization proof: the master key (proxy admin) reads it, a

tests/e2e/other/owned_jwt_gateway.py

@@ -0,0 +1,108 @@
+"""An owned, source-built proxy whose `litellm_jwtauth` block a test controls.
+
+The shared proxy on :4000 runs the CONTRIBUTING.md JWT block, so a test that
+needs a different `litellm_jwtauth` config boots its own gateway on a free port
+against the same database and the same Keycloak realm. The caller supplies the
+`litellm_jwtauth` mapping verbatim, which is exactly what makes a config an
+unfixed proxy rejects observable as a boot failure in this gateway's own log.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+import sys
+import time
+from collections.abc import Mapping
+from contextlib import ExitStack
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Final
+
+from e2e_config import INHERITED_ENV_PREFIXES, available_port
+from e2e_http import NoBody
+from idp import Keycloak, stop_process_group
+from proxy_client import ProxyClient, build_proxy_client
+
+MODEL_NAME: Final = "gemini-3.8-flash"
+
+
+@dataclass(slots=True)
+class OwnedJwtGateway:
+    base_url: str
+    proxy: ProxyClient
+    _environment: Mapping[str, str] = field(repr=False)
+    _command: tuple[str, ...] = field(repr=False)
+    _log_path: Path
+    _child: subprocess.Popen[bytes] | None = field(default=None, init=False, repr=False)
+
+    def start(self) -> None:
+        with self._log_path.open("ab") as log:
+            self._child = subprocess.Popen(
+                self._command,
+                env=self._environment,
+                stdout=log,
+                stderr=log,
+                start_new_session=True,
+            )
+        deadline: Final = time.monotonic() + 120
+        while time.monotonic() < deadline:
+            assert self._child.poll() is None, "owned JWT gateway exited; inspect its private log"
+            result = self.proxy.transport.probe("/health/liveliness", params=NoBody())
+            if result.status_code == 200:
+                return
+            time.sleep(0.5)
+        raise AssertionError("owned JWT gateway did not become ready")
+
+    def stop(self) -> None:
+        if self._child is not None:
+            stop_process_group(self._child)
+            assert self._child.poll() is not None, "old gateway process is still alive"
+
+
+def owned_jwt_gateway(
+    idp: Keycloak, directory: Path, cleanup: ExitStack, *, litellm_jwtauth: str, name: str
+) -> OwnedJwtGateway:
+    for env_name in ("DATABASE_URL", "LITELLM_LICENSE", "LITELLM_SALT_KEY", "LITELLM_MASTER_KEY"):
+        assert os.environ.get(env_name), f"{env_name} is required for the owned JWT gateway"
+    port: Final = available_port()
+    base_url: Final = f"http://127.0.0.1:{port}"
+    config: Final = directory / f"{name}.yaml"
+    config.write_text(
+        "model_list:\n"
+        f"  - model_name: {MODEL_NAME}\n"
+        "    litellm_params:\n"
+        f"      model: gemini/{MODEL_NAME}\n"
+        "      api_key: os.environ/GEMINI_API_KEY\n"
+        "general_settings:\n"
+        "  master_key: os.environ/LITELLM_MASTER_KEY\n"
+        "  database_url: os.environ/DATABASE_URL\n"
+        "  proxy_batch_write_at: 5\n"
+        "  enable_jwt_auth: true\n"
+        "  litellm_jwtauth:\n" + "".join(f"    {line}\n" for line in litellm_jwtauth.strip().splitlines())
+    )
+    environment: Final = {
+        **{key: value for key, value in os.environ.items() if not key.startswith(INHERITED_ENV_PREFIXES)},
+        "JWT_PUBLIC_KEY_URL": idp.jwks_url,
+        "JWT_ISSUER": idp.issuer,
+        "JWT_AUDIENCE": "litellm-e2e",
+        "LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY": "true",
+        "DISABLE_SCHEMA_UPDATE": "true",
+        "STORE_MODEL_IN_DB": "True",
+        "PYTHONPATH": str(Path(__file__).resolve().parents[3]),
+    }
+    gateway: Final = OwnedJwtGateway(
+        base_url=base_url,
+        proxy=build_proxy_client(
+            base_url=base_url,
+            control_plane_base_url=base_url,
+            replica_urls=(base_url,),
+            master_key=os.environ["LITELLM_MASTER_KEY"],
+        ),
+        _environment=environment,
+        _command=(sys.executable, "-m", "litellm.proxy.proxy_cli", "--config", str(config), "--port", str(port)),
+        _log_path=directory / f"{name}.log",
+    )
+    cleanup.callback(gateway.stop)
+    gateway.start()
+    return gateway

tests/e2e/other/test_jwt_auto_register_e2e.py

@@ -0,0 +1,179 @@
+"""auto_register with auto_register_map_existing_key binds the JWT claim to the user's existing key.
+
+`unregistered_jwt_client_behavior: auto_register` on `virtual_key_claim_field: sub` mints a fresh
+virtual key on the user's first JWT call. With `auto_register_map_existing_key: true` the proxy must
+instead point the new JWT mapping at a key the resolved user already owns, and mint only when the
+user has none. Each behavior gets its own gateway because the flag lives in `litellm_jwtauth`, so
+this file boots two owned proxies against the shared database and Keycloak realm.
+"""
+
+from __future__ import annotations
+
+import hashlib
+from collections.abc import Iterator
+from contextlib import ExitStack
+from typing import Final
+
+import pytest
+from e2e_config import unique_marker
+from e2e_http import unwrap
+from idp import Identity, Keycloak
+from lifecycle import ResourceManager
+from models import ChatBody, ChatMessage, JwtKeyMappingRow, KeyGenerateBody, TeamNewBody, UserNewBody
+from other_client import OtherClient
+from owned_jwt_gateway import MODEL_NAME, OwnedJwtGateway, owned_jwt_gateway
+
+pytestmark = pytest.mark.e2e
+
+_JWT_COMMON: Final = (
+    "user_id_jwt_field: sub\n"
+    "user_email_jwt_field: email\n"
+    "team_ids_jwt_field: groups\n"
+    "user_id_upsert: true\n"
+    "virtual_key_claim_field: sub\n"
+    "unregistered_jwt_client_behavior: auto_register"
+)
+
+
+def _key_hash(key: str) -> str:
+    return hashlib.sha256(key.encode()).hexdigest()
+
+
+def _ping() -> ChatBody:
+    return ChatBody(
+        model=MODEL_NAME,
+        messages=[ChatMessage(role="user", content=f"Reply with the single word ok. {unique_marker()}")],
+        max_tokens=5,
+    )
+
+
+def _identity_with_user(idp: Keycloak, client: OtherClient, resources: ResourceManager) -> Identity:
+    """An IdP identity plus the litellm user and team its claims resolve to, with
+    teardown that also sweeps the user's keys and JWT mapping rows the proxy
+    wrote, since those outlive the user row itself."""
+    marker: Final = unique_marker()
+    identity: Final = idp.provision(marker=marker, group=f"e2e-jwt-team-{marker}", defer=resources.defer)
+    resources.defer(lambda: client.proxy.delete_user(identity.user_id))
+    team_id: Final = client.proxy.create_team(TeamNewBody(team_alias=f"e2e-jwt-{marker}", team_id=identity.group))
+    resources.defer(lambda: client.proxy.delete_team(team_id))
+    unwrap(
+        client.user_new(
+            UserNewBody(
+                user_id=identity.user_id,
+                user_email=f"{identity.username}@example.com",
+                user_role="internal_user",
+                auto_create_key=False,
+            )
+        )
+    )
+
+    def delete_user_keys() -> None:
+        for row in unwrap(client.user_info(identity.user_id)).keys:
+            client.proxy.delete_key(row.token)
+
+    def delete_user_mappings() -> None:
+        for mapping in unwrap(client.jwt_mapping_list()).mappings:
+            if mapping.jwt_claim_value == identity.user_id:
+                _ = client.jwt_mapping_delete(mapping.id)
+
+    resources.defer(delete_user_keys)
+    resources.defer(delete_user_mappings)
+    return identity
+
+
+def _mapping_for(client: OtherClient, claim_value: str) -> JwtKeyMappingRow | None:
+    return next(
+        (row for row in unwrap(client.jwt_mapping_list()).mappings if row.jwt_claim_value == claim_value),
+        None,
+    )
+
+
+@pytest.fixture(scope="module")
+def mapping_gateway(idp: Keycloak, tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedJwtGateway]:
+    with ExitStack() as cleanup:
+        yield owned_jwt_gateway(
+            idp,
+            tmp_path_factory.mktemp("jwt-mapping"),
+            cleanup,
+            litellm_jwtauth=f"{_JWT_COMMON}\nauto_register_map_existing_key: true",
+            name="jwt-mapping-gateway",
+        )
+
+
+@pytest.fixture(scope="module")
+def minting_gateway(idp: Keycloak, tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedJwtGateway]:
+    with ExitStack() as cleanup:
+        yield owned_jwt_gateway(
+            idp,
+            tmp_path_factory.mktemp("jwt-minting"),
+            cleanup,
+            litellm_jwtauth=_JWT_COMMON,
+            name="jwt-minting-gateway",
+        )
+
+
+class TestJwtAutoRegisterMapExistingKey:
+    @pytest.mark.covers("other.auth.jwt.auto_register_maps_existing_key")
+    def test_first_jwt_call_maps_to_the_users_existing_key_and_mints_none(
+        self, client: OtherClient, idp: Keycloak, resources: ResourceManager, mapping_gateway: OwnedJwtGateway
+    ) -> None:
+        identity: Final = _identity_with_user(idp, client, resources)
+        existing_key: Final = client.proxy.generate_key(
+            KeyGenerateBody(user_id=identity.user_id, key_alias=f"e2e-jwt-existing-{unique_marker()}")
+        )
+        resources.defer(lambda: client.proxy.delete_key(existing_key))
+
+        response: Final = unwrap(mapping_gateway.proxy.chat(idp.access_token(identity), _ping()))
+
+        keys: Final = unwrap(client.user_info(identity.user_id)).keys
+        assert [row.token for row in keys] == [_key_hash(existing_key)], (
+            f"map_existing_key must leave the user with only their pre-existing key, got {keys}"
+        )
+        mapping: Final = _mapping_for(client, identity.user_id)
+        assert mapping is not None, (
+            f"no JWT mapping row for sub={identity.user_id}: {unwrap(client.jwt_mapping_list())}"
+        )
+        assert mapping.jwt_claim_name == "sub", f"mapping must bind the sub claim, got {mapping}"
+        assert mapping.created_by == "auto_register", f"mapping must be written by auto_register, got {mapping}"
+        rows: Final = client.proxy.poll_logs_for_key(existing_key)
+        assert any(row.request_id == response.id for row in rows), (
+            f"the JWT chat must be billed to the user's existing key, spend rows for it: {rows}"
+        )
+
+    @pytest.mark.covers("other.auth.jwt.auto_register_mints_when_keyless")
+    def test_first_jwt_call_mints_a_key_when_the_user_has_none(
+        self, client: OtherClient, idp: Keycloak, resources: ResourceManager, mapping_gateway: OwnedJwtGateway
+    ) -> None:
+        identity: Final = _identity_with_user(idp, client, resources)
+
+        response: Final = unwrap(mapping_gateway.proxy.chat(idp.access_token(identity), _ping()))
+        assert response.choices, f"JWT chat returned no completion: {response}"
+
+        keys: Final = unwrap(client.user_info(identity.user_id)).keys
+        assert len(keys) == 1, f"a keyless user must get exactly one minted key, got {keys}"
+        mapping: Final = _mapping_for(client, identity.user_id)
+        assert mapping is not None and mapping.jwt_claim_name == "sub", (
+            f"the minted key must be recorded as a sub-claim mapping, mappings: {unwrap(client.jwt_mapping_list())}"
+        )
+
+    @pytest.mark.covers("other.auth.jwt.auto_register_default_mints")
+    def test_default_behavior_still_mints_when_the_user_already_has_a_key(
+        self, client: OtherClient, idp: Keycloak, resources: ResourceManager, minting_gateway: OwnedJwtGateway
+    ) -> None:
+        identity: Final = _identity_with_user(idp, client, resources)
+        existing_key: Final = client.proxy.generate_key(
+            KeyGenerateBody(user_id=identity.user_id, key_alias=f"e2e-jwt-existing-{unique_marker()}")
+        )
+        resources.defer(lambda: client.proxy.delete_key(existing_key))
+
+        response: Final = unwrap(minting_gateway.proxy.chat(idp.access_token(identity), _ping()))
+        assert response.id is not None, f"JWT chat returned no response id: {response}"
+
+        keys: Final = unwrap(client.user_info(identity.user_id)).keys
+        assert len(keys) == 2, (
+            f"default auto_register must mint a second key for a user who already has one, got {keys}"
+        )
+        rows: Final = client.proxy.poll_logs_for_request_id(response.id)
+        assert rows and all(row.api_key != _key_hash(existing_key) for row in rows), (
+            f"the default path must bill the minted key, not the user's existing one: {rows}"
+        )