From 1464caeab9bf385743d61e2632d612698791af59 Mon Sep 17 00:00:00 2001
From: moe-berri <moe@berri.ai>
Date: Sat, 12 Sep 2026 21:32:17 -0700
Subject: [PATCH] fix(memory): reserve pilot validation capacity for registered
 keys

---
 deploy/memory-pilot/pilot.py                  |  8 ++++---
 .../proxy/memory/test_memory_pilot.py         | 24 ++++++++++++++-----
 2 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/deploy/memory-pilot/pilot.py b/deploy/memory-pilot/pilot.py
index 40aacc4518d..27a9a040332 100644
--- a/deploy/memory-pilot/pilot.py
+++ b/deploy/memory-pilot/pilot.py
@@ -66,7 +66,8 @@ forward_credential: Final = ForwardCredential()
 class PilotGateway:
     def __init__(self, app: ASGIApp) -> None:
         self.app = app
-        self.validation_slots = asyncio.Semaphore(16)
+        self.registered_validation_slots = asyncio.Semaphore(12)
+        self.enrollment_validation_slots = asyncio.Semaphore(4)
         self.upstream = get_async_httpx_client(
             httpxSpecialProvider.PassThroughEndpoint,
             params={"timeout": 20, "client_alias": "memory-pilot-upstream"},
@@ -122,8 +123,9 @@ class PilotGateway:
                 {"error": "Upstream keys can only use inference and their own memories"}, status_code=403
             )(scope, receive, send)
             return
+        validation_slots: Final = self.registered_validation_slots if local_key else self.enrollment_validation_slots
         try:
-            await asyncio.wait_for(self.validation_slots.acquire(), timeout=0.05)
+            await asyncio.wait_for(validation_slots.acquire(), timeout=0.05)
         except TimeoutError:
             await JSONResponse(
                 {"error": "Pilot credential validation is busy; retry shortly"},
@@ -139,7 +141,7 @@ class PilotGateway:
             await JSONResponse({"error": "Upstream gateway unavailable"}, status_code=503)(scope, receive, send)
             return
         finally:
-            self.validation_slots.release()
+            validation_slots.release()
         if models.is_error:
             await JSONResponse({"error": "Upstream gateway rejected this key"}, status_code=models.status_code)(
                 scope, receive, send
diff --git a/tests/test_litellm/proxy/memory/test_memory_pilot.py b/tests/test_litellm/proxy/memory/test_memory_pilot.py
index bb2e8d29499..72075216bf8 100644
--- a/tests/test_litellm/proxy/memory/test_memory_pilot.py
+++ b/tests/test_litellm/proxy/memory/test_memory_pilot.py
@@ -1,6 +1,7 @@
 """Bound unauthenticated upstream validation without replacing gateway authentication."""
 
 import asyncio
+import hashlib
 import importlib.util
 from pathlib import Path
 from unittest.mock import AsyncMock, MagicMock, patch
@@ -11,7 +12,9 @@ from starlette.applications import Starlette
 
 
 @pytest.mark.asyncio
-async def test_upstream_validation_is_bounded_and_slots_release_after_failure(monkeypatch: pytest.MonkeyPatch) -> None:
+async def test_unknown_keys_cannot_consume_registered_validation_capacity_and_slots_release(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
     monkeypatch.setenv("UPSTREAM_LITELLM_BASE_URL", "https://upstream.example.invalid")
     filename = Path(__file__).resolve().parents[4] / "deploy" / "memory-pilot" / "pilot.py"
     spec = importlib.util.spec_from_file_location("memory_pilot_test", filename)
@@ -24,9 +27,11 @@ async def test_upstream_validation_is_bounded_and_slots_release_after_failure(mo
     count = 0
 
     async def upstream_get(*args: object, **kwargs: object) -> httpx.Response:
+        if kwargs.get("headers") == {"Authorization": "Bearer sk-established"}:
+            return httpx.Response(200, json={"data": [{"id": "model"}]})
         nonlocal count
         count += 1
-        if count == 16:
+        if count == 4:
             entered.set()
         await release.wait()
         raise httpx.ConnectError("unavailable")
@@ -34,21 +39,28 @@ async def test_upstream_validation_is_bounded_and_slots_release_after_failure(mo
     upstream = MagicMock(get=AsyncMock(side_effect=upstream_get))
     gateway.upstream = upstream
     database = MagicMock()
-    database.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None)
+
+    async def key_lookup(*, where: dict[str, str]) -> dict[str, str] | None:
+        digest = hashlib.sha256(b"sk-established").hexdigest()
+        return {"token": digest} if where == {"token": digest} else None
+
+    database.db.litellm_verificationtoken.find_unique = AsyncMock(side_effect=key_lookup)
     with patch.multiple(  # test-quality-ok: Inject external database/config; exercise real ASGI admission.
         "litellm.proxy.proxy_server", prisma_client=database, master_key="local-admin"
     ):
         async with httpx.AsyncClient(transport=httpx.ASGITransport(app=gateway), base_url="http://pilot") as client:
             pending = [
                 asyncio.create_task(client.get("/v1/models", headers={"Authorization": f"Bearer sk-invalid-{i}"}))
-                for i in range(16)
+                for i in range(4)
             ]
             await asyncio.wait_for(entered.wait(), timeout=2)
             refused = await client.get("/v1/models", headers={"Authorization": "Bearer sk-overload"})
             assert refused.status_code == 503 and refused.headers["retry-after"] == "1"
-            assert upstream.get.await_count == 16
+            assert upstream.get.await_count == 4
+            established = await client.get("/v1/models", headers={"Authorization": "Bearer sk-established"})
+            assert established.status_code == 200 and established.json() == {"data": [{"id": "model"}]}
             release.set()
             assert all(response.status_code == 503 for response in await asyncio.gather(*pending))
             again = await client.get("/v1/models", headers={"Authorization": "Bearer sk-next"})
             assert again.status_code == 503 and "unavailable" in again.text
-            assert upstream.get.await_count == 17
+            assert upstream.get.await_count == 6