From 6f92e244fb32a804755ce3add2b67d55c4e4b634 Mon Sep 17 00:00:00 2001
From: moe-berri <moe@berri.ai>
Date: Sat, 12 Sep 2026 21:56:24 -0700
Subject: [PATCH] fix(memory): refresh pilot admission after successful
 validation

---
 deploy/memory-pilot/pilot.py                         |  6 ++++--
 tests/test_litellm/proxy/memory/test_memory_pilot.py | 12 ++++++++++--
 2 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/deploy/memory-pilot/pilot.py b/deploy/memory-pilot/pilot.py
index ef3f7a8b19f..fc006a98ca4 100644
--- a/deploy/memory-pilot/pilot.py
+++ b/deploy/memory-pilot/pilot.py
@@ -4,6 +4,7 @@ import asyncio
 import hashlib
 import os
 import secrets
+from collections.abc import Callable
 from contextvars import ContextVar
 from typing import Final
 
@@ -65,11 +66,11 @@ forward_credential: Final = ForwardCredential()
 
 
 class PilotGateway:
-    def __init__(self, app: ASGIApp) -> None:
+    def __init__(self, app: ASGIApp, clock: Callable[[], float] | None = None) -> None:
         self.app = app
         self.registered_validation_slots = asyncio.Semaphore(12)
         self.enrollment_validation_slots = asyncio.Semaphore(4)
-        self.recently_validated = InMemoryCache(max_size_in_memory=1000, default_ttl=60)
+        self.recently_validated = InMemoryCache(max_size_in_memory=1000, default_ttl=60, clock=clock)
         self.upstream = get_async_httpx_client(
             httpxSpecialProvider.PassThroughEndpoint,
             params={"timeout": 20, "client_alias": "memory-pilot-upstream"},
@@ -146,6 +147,7 @@ class PilotGateway:
             if models.status_code in (401, 403):
                 self.recently_validated.delete_cache(digest)
             elif models.is_success:
+                self.recently_validated.delete_cache(digest)
                 self.recently_validated.set_cache(digest, True)
         except httpx.HTTPError:
             await JSONResponse({"error": "Upstream gateway unavailable"}, status_code=503)(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 eb64ac01398..93b0b57889a 100644
--- a/tests/test_litellm/proxy/memory/test_memory_pilot.py
+++ b/tests/test_litellm/proxy/memory/test_memory_pilot.py
@@ -21,7 +21,8 @@ async def test_unknown_keys_cannot_consume_registered_validation_capacity_and_sl
     assert spec is not None and spec.loader is not None
     module = importlib.util.module_from_spec(spec)
     spec.loader.exec_module(module)
-    gateway = module.PilotGateway(Starlette())
+    clock = MagicMock(return_value=0.0)
+    gateway = module.PilotGateway(Starlette(), clock=clock)
     entered = asyncio.Event()
     release = asyncio.Event()
     count = 0
@@ -59,13 +60,20 @@ async def test_unknown_keys_cannot_consume_registered_validation_capacity_and_sl
             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 == 5
+            clock.return_value = 59.0
             established = await client.get("/v1/models", headers={"Authorization": "Bearer sk-established"})
             assert established.status_code == 200 and established.json() == {"data": [{"id": "model"}]}
+            clock.return_value = 61.0
+            refreshed = await client.get("/v1/models", headers={"Authorization": "Bearer sk-established"})
+            assert refreshed.status_code == 200
+            clock.return_value = 122.0
+            expired = await client.get("/v1/models", headers={"Authorization": "Bearer sk-established"})
+            assert expired.status_code == 503 and expired.headers["retry-after"] == "1"
             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 == 7
+            assert upstream.get.await_count == 8
 
 
 @pytest.mark.asyncio