@miller-tech/uap 1.76.3 → 1.76.4

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@miller-tech/uap",
-  "version": "1.76.3",
+  "version": "1.76.4",
   "description": "Autonomous AI agent memory system with CLAUDE.md protocol enforcement",
   "type": "module",
   "main": "dist/index.js",

package/tools/agents/scripts/anthropic_proxy.py

@@ -1851,6 +1851,50 @@ PROXY_CONCURRENCY_QUEUE_TIMEOUT = float(
 upstream_semaphore: asyncio.Semaphore | None = None
 
 
+# ---------------------------------------------------------------------------
+# Session admission control — cap the number of DISTINCT "hot" sessions
+# ---------------------------------------------------------------------------
+# The semaphore above limits CONCURRENT REQUESTS. On a multi-slot llama.cpp
+# (--parallel N) llama's native prompt-cache keeps each session on its own slot
+# across turns — but only N slots exist. When MORE than N distinct sessions are
+# active, a returning session finds its slot reassigned to another session, its
+# KV gone, forcing a from-scratch reprocess of the whole context (brutal on an
+# SSM/Mamba model, which cannot restore partial KV: every eviction is a full
+# prefill). The per-request semaphore can't prevent this — it lets session A
+# finish + release, then admits session E, which evicts A's slot.
+#
+# Session admission caps the number of DISTINCT sessions holding a slot at once.
+# A new session over the limit WAITS (queues) until an admitted session goes
+# idle (no request for IDLE_TTL) and is pruned, instead of barging in and
+# evicting a hot session. Admission is STICKY across a session's turns; it is
+# released only by idle-TTL expiry (or, under sustained over-subscription, a
+# wait-timeout graceful-degrade that force-admits, evicting the LRU). Default
+# OFF — opt in via PROXY_SESSION_ADMISSION=on (set the LIMIT to the llama slot
+# count / --parallel value).
+PROXY_SESSION_ADMISSION = os.environ.get(
+    "PROXY_SESSION_ADMISSION", "off"
+).lower() not in {"", "0", "off", "false", "no"}
+# Max distinct hot sessions. Default = the concurrency limit (= llama slots).
+PROXY_SESSION_ADMISSION_LIMIT = int(
+    os.environ.get("PROXY_SESSION_ADMISSION_LIMIT", str(PROXY_CONCURRENCY_LIMIT))
+)
+# Seconds a session may be idle (no request) before its admission is pruned,
+# freeing a slot for a waiting session. Tune ABOVE typical inter-turn gaps so an
+# actively-working session is never pruned mid-task. 0 = never prune on idle.
+PROXY_SESSION_ADMISSION_IDLE_TTL = float(
+    os.environ.get("PROXY_SESSION_ADMISSION_IDLE_TTL", "90")
+)
+# Max seconds a new session waits for admission before graceful-degrade
+# (force-admit, evicting the LRU admitted session). 0 = wait indefinitely.
+PROXY_SESSION_ADMISSION_WAIT_TIMEOUT = float(
+    os.environ.get("PROXY_SESSION_ADMISSION_WAIT_TIMEOUT", "300")
+)
+# How often a waiter re-checks for a freed slot (re-prunes idle admissions).
+PROXY_SESSION_ADMISSION_POLL = float(
+    os.environ.get("PROXY_SESSION_ADMISSION_POLL", "3")
+)
+
+
 # ---------------------------------------------------------------------------
 # Slot save/restore — cross-session KV-cache preservation
 # ---------------------------------------------------------------------------
@@ -1902,6 +1946,12 @@ _current_request_session: contextvars.ContextVar[str | None] = contextvars.Conte
     "uap_current_request_session", default=None
 )
 
+# Session admission state. _admitted_sessions maps session_id -> last-seen
+# monotonic ts; OrderedDict insertion order is the LRU (oldest = front).
+# Guarded by _admission_cond's lock (created lazily on the running event loop).
+_admitted_sessions: "OrderedDict[str, float]" = OrderedDict()
+_admission_cond: "asyncio.Condition | None" = None
+
 
 def _slot_endpoint_base() -> str:
     """Base URL for llama-server's /slots endpoint (LLAMA_CPP_BASE without /v1)."""
@@ -2052,6 +2102,93 @@ def _prepare_slot_save_dir() -> None:
         logger.warning("SLOT SAVE/RESTORE: startup dir prep failed: %s", exc)
 
 
+def _prune_idle_admissions(now: float) -> list[str]:
+    """Remove admitted sessions idle longer than the TTL; return their ids.
+    Caller must hold the admission lock. 0 TTL disables idle pruning."""
+    ttl = PROXY_SESSION_ADMISSION_IDLE_TTL
+    if ttl <= 0:
+        return []
+    stale = [sid for sid, ts in _admitted_sessions.items() if now - ts > ttl]
+    for sid in stale:
+        _admitted_sessions.pop(sid, None)
+    return stale
+
+
+def _try_admit_session(session_id: str, now: float) -> bool:
+    """Synchronous admission core (caller holds the admission lock). Returns
+    True if the session is admitted — already hot (refreshed) or newly admitted
+    into a free slot — False if the hot set is full. Pure/deterministic given
+    `now`, so it is unit-tested directly."""
+    _prune_idle_admissions(now)
+    if session_id in _admitted_sessions:
+        _admitted_sessions[session_id] = now
+        _admitted_sessions.move_to_end(session_id)
+        return True
+    if len(_admitted_sessions) < PROXY_SESSION_ADMISSION_LIMIT:
+        _admitted_sessions[session_id] = now
+        _admitted_sessions.move_to_end(session_id)
+        return True
+    return False
+
+
+async def _ensure_session_admitted(session_id: str | None) -> None:
+    """Block until `session_id` is admitted to the hot set (size <= LIMIT).
+
+    Sticky: admission persists across a session's turns and is released only by
+    idle-TTL pruning. A new session over the limit queues (waits) rather than
+    evicting a hot session. On wait-timeout it force-admits, evicting the LRU,
+    to avoid a hard stall (graceful degrade to pre-admission behaviour).
+    No-op when disabled or when there is no session id."""
+    global _admission_cond
+    if not PROXY_SESSION_ADMISSION or not session_id:
+        return
+    if _admission_cond is None:
+        _admission_cond = asyncio.Condition()
+    cond = _admission_cond
+    deadline = (
+        time.monotonic() + PROXY_SESSION_ADMISSION_WAIT_TIMEOUT
+        if PROXY_SESSION_ADMISSION_WAIT_TIMEOUT > 0
+        else None
+    )
+    waited = False
+    async with cond:
+        while True:
+            now = time.monotonic()
+            if _try_admit_session(session_id, now):
+                if waited:
+                    logger.info(
+                        "SESSION ADMISSION: admitted %s after wait (%d/%d hot)",
+                        session_id[:12], len(_admitted_sessions),
+                        PROXY_SESSION_ADMISSION_LIMIT,
+                    )
+                # A refresh/new admit may have pruned idle sessions; wake peers.
+                cond.notify_all()
+                return
+            # Hot set full and this session isn't in it.
+            if deadline is not None and now >= deadline:
+                evicted = "-"
+                if _admitted_sessions:
+                    evicted = next(iter(_admitted_sessions))  # LRU = oldest front
+                    _admitted_sessions.pop(evicted, None)
+                _admitted_sessions[session_id] = now
+                logger.warning(
+                    "SESSION ADMISSION: wait timeout (%ds), force-admitted %s "
+                    "(evicted LRU %s) — sustained over-subscription (>%d hot sessions)",
+                    int(PROXY_SESSION_ADMISSION_WAIT_TIMEOUT), session_id[:12],
+                    evicted[:12], PROXY_SESSION_ADMISSION_LIMIT,
+                )
+                cond.notify_all()
+                return
+            waited = True
+            timeout = PROXY_SESSION_ADMISSION_POLL
+            if deadline is not None:
+                timeout = max(0.0, min(timeout, deadline - now))
+            try:
+                await asyncio.wait_for(cond.wait(), timeout=timeout)
+            except asyncio.TimeoutError:
+                pass  # re-loop: re-prune idle admissions, retry
+
+
 async def _acquire_upstream_slot() -> bool:
     """Acquire a semaphore slot for an upstream request.
 
@@ -2136,6 +2273,11 @@ async def _post_with_retry(
     rather than all hammering llama.cpp at once. Slot is released in a
     finally block so it's always returned to the pool even on error.
     """
+    # Session admission: cap DISTINCT hot sessions to <= the slot count so they
+    # don't evict each other's KV (sticky, released by idle-TTL). No-op when
+    # PROXY_SESSION_ADMISSION is off. Runs BEFORE the per-request semaphore so a
+    # queued new session waits without holding a concurrency slot.
+    await _ensure_session_admitted(_current_request_session.get())
     acquired = await _acquire_upstream_slot()
     if not acquired:
         logger.warning(
@@ -2272,7 +2414,19 @@ async def lifespan(app: FastAPI):
     global http_client
     global default_context_window
     global upstream_semaphore
+    global _admission_cond
     upstream_semaphore = asyncio.Semaphore(PROXY_CONCURRENCY_LIMIT)
+    # Bind the admission condition to THIS event loop and clear stale state.
+    _admission_cond = asyncio.Condition()
+    _admitted_sessions.clear()
+    if PROXY_SESSION_ADMISSION:
+        logger.info(
+            "SESSION ADMISSION: on (limit=%d hot sessions, idle_ttl=%.0fs, "
+            "wait_timeout=%.0fs)",
+            PROXY_SESSION_ADMISSION_LIMIT,
+            PROXY_SESSION_ADMISSION_IDLE_TTL,
+            PROXY_SESSION_ADMISSION_WAIT_TIMEOUT,
+        )
     logger.info(
         "CONCURRENCY: upstream semaphore initialized limit=%d queue_timeout=%.0fs",
         PROXY_CONCURRENCY_LIMIT,
@@ -2369,6 +2523,8 @@ async def lifespan(app: FastAPI):
     http_client = None
     if upstream_semaphore is not None:
         upstream_semaphore = None
+    _admission_cond = None
+    _admitted_sessions.clear()
     logger.info("Proxy shut down")
 
 

package/tools/agents/tests/test_session_admission.py

@@ -0,0 +1,117 @@
+#!/usr/bin/env python3
+"""Session-admission control: cap the number of DISTINCT hot sessions to <= the
+slot count so they don't evict each other's KV.
+
+The per-request semaphore limits CONCURRENT REQUESTS; admission limits DISTINCT
+SESSIONS. A new session over the limit queues until an admitted session goes
+idle (TTL) and is pruned, instead of barging in and evicting a hot session
+(every eviction = a full reprocess on the SSM model). Admission is sticky across
+a session's turns; on a wait-timeout it force-admits (graceful degrade).
+
+The synchronous core `_try_admit_session` / `_prune_idle_admissions` is pure
+given `now`, so it's tested directly. The async wrapper is tested for the
+no-op-when-disabled and force-admit-on-timeout paths.
+"""
+
+import asyncio
+import importlib.util
+import unittest
+from pathlib import Path
+
+
+def _load():
+    p = Path(__file__).resolve().parents[1] / "scripts" / "anthropic_proxy.py"
+    spec = importlib.util.spec_from_file_location("anthropic_proxy", p)
+    m = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(m)
+    return m
+
+
+proxy = _load()
+
+
+class TestAdmissionCore(unittest.TestCase):
+    def setUp(self):
+        proxy._admitted_sessions.clear()
+        proxy.PROXY_SESSION_ADMISSION_LIMIT = 2
+        proxy.PROXY_SESSION_ADMISSION_IDLE_TTL = 90.0
+
+    def test_admit_up_to_limit_then_full(self):
+        self.assertTrue(proxy._try_admit_session("s1", now=0))
+        self.assertTrue(proxy._try_admit_session("s2", now=1))
+        # 3rd distinct session: hot set full -> rejected
+        self.assertFalse(proxy._try_admit_session("s3", now=2))
+        self.assertEqual(set(proxy._admitted_sessions), {"s1", "s2"})
+
+    def test_already_hot_session_refreshes_and_is_admitted(self):
+        proxy._try_admit_session("s1", now=0)
+        proxy._try_admit_session("s2", now=1)
+        # s1 returns — must stay admitted (no eviction of itself) and refresh ts
+        self.assertTrue(proxy._try_admit_session("s1", now=50))
+        self.assertEqual(proxy._admitted_sessions["s1"], 50)
+
+    def test_idle_session_pruned_after_ttl_frees_a_slot(self):
+        proxy._try_admit_session("s1", now=0)
+        proxy._try_admit_session("s2", now=1)
+        self.assertFalse(proxy._try_admit_session("s3", now=2))  # full
+        # advance past TTL: s1 (ts0) and s2 (ts1) are now idle > 90s -> pruned,
+        # s3 admits into the freed slot
+        self.assertTrue(proxy._try_admit_session("s3", now=200))
+        self.assertIn("s3", proxy._admitted_sessions)
+        self.assertNotIn("s1", proxy._admitted_sessions)
+
+    def test_zero_ttl_disables_idle_pruning(self):
+        proxy.PROXY_SESSION_ADMISSION_IDLE_TTL = 0.0
+        proxy._try_admit_session("s1", now=0)
+        proxy._try_admit_session("s2", now=1)
+        # even far in the future, no idle prune -> s3 stays blocked
+        self.assertFalse(proxy._try_admit_session("s3", now=10_000))
+
+    def test_lru_is_front_of_ordereddict(self):
+        proxy._try_admit_session("s1", now=0)
+        proxy._try_admit_session("s2", now=1)
+        proxy._try_admit_session("s1", now=5)  # refresh moves s1 to back
+        # s2 is now the LRU (front) -> first to be force-evicted
+        self.assertEqual(next(iter(proxy._admitted_sessions)), "s2")
+
+
+class TestAdmissionAsync(unittest.TestCase):
+    def setUp(self):
+        proxy._admitted_sessions.clear()
+        proxy._admission_cond = None
+
+    def test_noop_when_disabled(self):
+        proxy.PROXY_SESSION_ADMISSION = False
+        asyncio.run(proxy._ensure_session_admitted("anything"))
+        self.assertEqual(len(proxy._admitted_sessions), 0)  # never tracked
+
+    def test_noop_when_no_session_id(self):
+        proxy.PROXY_SESSION_ADMISSION = True
+        asyncio.run(proxy._ensure_session_admitted(None))
+        self.assertEqual(len(proxy._admitted_sessions), 0)
+
+    def test_admits_when_room(self):
+        proxy.PROXY_SESSION_ADMISSION = True
+        proxy.PROXY_SESSION_ADMISSION_LIMIT = 2
+        asyncio.run(proxy._ensure_session_admitted("s1"))
+        self.assertIn("s1", proxy._admitted_sessions)
+
+    def test_force_admit_on_wait_timeout(self):
+        # Over-subscribed + no idle prune + ~0 wait timeout -> force-admit,
+        # evicting the LRU (graceful degrade, never hard-stalls).
+        proxy.PROXY_SESSION_ADMISSION = True
+        proxy.PROXY_SESSION_ADMISSION_LIMIT = 1
+        proxy.PROXY_SESSION_ADMISSION_IDLE_TTL = 0.0   # never prune
+        proxy.PROXY_SESSION_ADMISSION_WAIT_TIMEOUT = 0.01
+        proxy.PROXY_SESSION_ADMISSION_POLL = 0.01
+
+        async def run():
+            await proxy._ensure_session_admitted("hot")     # fills the 1 slot
+            await proxy._ensure_session_admitted("waiter")  # times out -> force
+        asyncio.run(run())
+        self.assertIn("waiter", proxy._admitted_sessions)
+        self.assertNotIn("hot", proxy._admitted_sessions)   # LRU evicted
+
+
+if __name__ == "__main__":
+    unittest.main()