fix(nemo-agents): bound live Fabric sessions with an LRU cap

plugins/nemo-agents/src/nemo_agents_plugin/fabric/server.py

@@ -35,6 +35,7 @@
 from nemo_agents_plugin.fabric.session_manager import (
     DEFAULT_IDLE_SESSION_TIMEOUT_SECONDS,
     DEFAULT_MAX_CONCURRENT_INVOCATIONS,
+    DEFAULT_MAX_LIVE_SESSIONS,
     DEFAULT_SESSION_CLEANUP_INTERVAL_SECONDS,
     FabricSessionManager,
     FabricSessionStartError,
@@ -57,12 +58,15 @@ class FabricServingSettings:
     """Operational settings for the Platform-owned Fabric server."""
 
     max_concurrent_invocations: int = DEFAULT_MAX_CONCURRENT_INVOCATIONS
+    max_live_sessions: int = DEFAULT_MAX_LIVE_SESSIONS
     idle_session_timeout_seconds: float = DEFAULT_IDLE_SESSION_TIMEOUT_SECONDS
     session_cleanup_interval_seconds: float = DEFAULT_SESSION_CLEANUP_INTERVAL_SECONDS
 
     def __post_init__(self) -> None:
         if self.max_concurrent_invocations < 0:
             raise ValueError("max_concurrent_invocations must be greater than or equal to zero.")
+        if self.max_live_sessions < 0:
+            raise ValueError("max_live_sessions must be greater than or equal to zero.")
         if self.idle_session_timeout_seconds <= 0:
             raise ValueError("idle_session_timeout_seconds must be greater than zero.")
         if self.session_cleanup_interval_seconds <= 0:
@@ -257,6 +261,7 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
             base_dir=config_path.parent,
             session_registry=session_registry,
             max_concurrent_invocations=settings.max_concurrent_invocations,
+            max_live_sessions=settings.max_live_sessions,
         )
         app.state.session_manager = session_manager
         cleanup_shutdown = asyncio.Event()
@@ -390,6 +395,12 @@ def main(argv: list[str] | None = None) -> int:
         default=DEFAULT_MAX_CONCURRENT_INVOCATIONS,
         help="Maximum concurrent Fabric invocations; use 0 for unlimited.",
     )
+    parser.add_argument(
+        "--max-live-sessions",
+        type=int,
+        default=DEFAULT_MAX_LIVE_SESSIONS,
+        help="Maximum concurrently live sessions; the least recently used idle one is evicted above it. Use 0 for unlimited.",
+    )
     parser.add_argument(
         "--idle-session-timeout-seconds",
         type=float,
@@ -412,6 +423,7 @@ def main(argv: list[str] | None = None) -> int:
             args.agent_config,
             settings=FabricServingSettings(
                 max_concurrent_invocations=args.max_concurrent_invocations,
+                max_live_sessions=args.max_live_sessions,
                 idle_session_timeout_seconds=args.idle_session_timeout_seconds,
                 session_cleanup_interval_seconds=args.session_cleanup_interval_seconds,
             ),

plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_manager.py

@@ -35,6 +35,8 @@
 DEFAULT_MAX_CONCURRENT_INVOCATIONS = 8
 DEFAULT_IDLE_SESSION_TIMEOUT_SECONDS = 30 * 60
 DEFAULT_SESSION_CLEANUP_INTERVAL_SECONDS = 5 * 60
+# A caller that sends no session header cannot close what it opens, so the cap does it instead.
+DEFAULT_MAX_LIVE_SESSIONS = DEFAULT_MAX_CONCURRENT_INVOCATIONS
 
 
 class FabricSessionStartError(RuntimeError):
@@ -56,14 +58,18 @@ def __init__(
         session_registry: FabricSessionRegistry,
         fabric: Fabric | None = None,
         max_concurrent_invocations: int = DEFAULT_MAX_CONCURRENT_INVOCATIONS,
+        max_live_sessions: int = DEFAULT_MAX_LIVE_SESSIONS,
     ) -> None:
         if max_concurrent_invocations < 0:
             raise ValueError("max_concurrent_invocations must be greater than or equal to zero.")
+        if max_live_sessions < 0:
+            raise ValueError("max_live_sessions must be greater than or equal to zero.")
 
         self._agent_config = agent_config
         self._base_dir = base_dir
         self._session_registry = session_registry
         self._fabric = fabric
+        self._max_live_sessions = max_live_sessions
         self._invocation_semaphore = (
             asyncio.Semaphore(max_concurrent_invocations) if max_concurrent_invocations > 0 else None
         )
@@ -87,7 +93,7 @@ async def open_session(self) -> FabricRuntimeSession:
             raise FabricSessionStartError(f"Fabric runtime startup failed: {error}") from error
 
         try:
-            return await self._session_registry.register(runtime)
+            session = await self._session_registry.register(runtime)
         except BaseException:
             # A started runtime must not leak if registration fails or is cancelled.
             try:
@@ -96,6 +102,22 @@ async def open_session(self) -> FabricRuntimeSession:
                 logger.exception("Failed to stop Fabric runtime after session registration failed.")
             raise
 
+        await self._evict_over_capacity(keep=session.session_id)
+        return session
+
+    async def _evict_over_capacity(self, *, keep: str) -> None:
+        """Stop the least recently used idle sessions above the live-session cap."""
+        evicted = await self._session_registry.evict_over_capacity(
+            max_sessions=self._max_live_sessions,
+            keep=keep,
+        )
+        for session in evicted:
+            logger.info("Evicting idle Fabric session %s to stay within the live-session cap.", session.session_id)
+            try:
+                await self._stop_session(session)
+            except FabricSessionStopError:
+                logger.exception("Failed to stop evicted Fabric session %s.", session.session_id)
+
     async def resolve_session(self, session_id: str | None) -> FabricRuntimeSession:
         """Open a new session or resolve an existing session by its opaque ID."""
         if session_id is None:

plugins/nemo-agents/src/nemo_agents_plugin/fabric/session_registry.py

@@ -95,6 +95,27 @@ async def remove(self, session_id: str) -> FabricRuntimeSession | None:
                 session.closing = True
             return session
 
+    async def evict_over_capacity(self, *, max_sessions: int, keep: str) -> list[FabricRuntimeSession]:
+        """Remove the least recently used idle sessions above *max_sessions*, never *keep*."""
+        if max_sessions <= 0:
+            return []
+
+        evicted: list[FabricRuntimeSession] = []
+        async with self._lock:
+            while len(self._sessions) > max_sessions:
+                candidates = [
+                    session
+                    for session in self._sessions.values()
+                    if session.session_id != keep and not session.invocation_lock.locked()
+                ]
+                if not candidates:
+                    break
+                victim = min(candidates, key=lambda session: session.last_accessed_at)
+                victim.closing = True
+                evicted.append(victim)
+                del self._sessions[victim.session_id]
+        return evicted
+
     async def remove_expired(self, *, idle_timeout_seconds: float) -> list[FabricRuntimeSession]:
         """Remove inactive sessions that are not currently invoking."""
         cutoff = time.monotonic() - idle_timeout_seconds

plugins/nemo-agents/tests/unit/test_fabric_session_registry.py

@@ -117,6 +117,69 @@ async def test_remove_expired_removes_only_idle_sessions(monkeypatch: pytest.Mon
     assert await registry.count() == 2
 
 
+@pytest.mark.asyncio
+async def test_evict_over_capacity_removes_least_recently_used_first() -> None:
+    registry = FabricSessionRegistry()
+    oldest = await registry.register(cast(Any, object()), session_id="oldest")
+    middle = await registry.register(cast(Any, object()), session_id="middle")
+    newest = await registry.register(cast(Any, object()), session_id="newest")
+    oldest.last_accessed_at = 10.0
+    middle.last_accessed_at = 20.0
+    newest.last_accessed_at = 30.0
+
+    evicted = await registry.evict_over_capacity(max_sessions=1, keep="newest")
+
+    assert evicted == [oldest, middle]
+    assert oldest.closing is True
+    assert middle.closing is True
+    assert newest.closing is False
+    assert await registry.count() == 1
+
+
+@pytest.mark.asyncio
+async def test_evict_over_capacity_keeps_the_session_just_opened() -> None:
+    registry = FabricSessionRegistry()
+    established = await registry.register(cast(Any, object()), session_id="established")
+    opened = await registry.register(cast(Any, object()), session_id="opened")
+    established.last_accessed_at = 90.0
+    opened.last_accessed_at = 10.0
+
+    evicted = await registry.evict_over_capacity(max_sessions=1, keep="opened")
+
+    assert evicted == [established]
+    assert await registry.count() == 1
+
+
+@pytest.mark.asyncio
+async def test_evict_over_capacity_never_evicts_a_session_mid_invocation() -> None:
+    registry = FabricSessionRegistry()
+    busy = await registry.register(cast(Any, object()), session_id="busy")
+    opened = await registry.register(cast(Any, object()), session_id="opened")
+    busy.last_accessed_at = 10.0
+    opened.last_accessed_at = 20.0
+
+    await busy.invocation_lock.acquire()
+    try:
+        evicted = await registry.evict_over_capacity(max_sessions=1, keep="opened")
+    finally:
+        busy.invocation_lock.release()
+
+    assert evicted == []
+    assert busy.closing is False
+    assert await registry.count() == 2
+
+
+@pytest.mark.asyncio
+async def test_evict_over_capacity_is_a_no_op_under_the_cap_or_when_unlimited() -> None:
+    registry = FabricSessionRegistry()
+    await registry.register(cast(Any, object()), session_id="session-1")
+    await registry.register(cast(Any, object()), session_id="session-2")
+
+    assert await registry.evict_over_capacity(max_sessions=4, keep="session-2") == []
+    assert await registry.evict_over_capacity(max_sessions=0, keep="session-2") == []
+    assert await registry.count() == 2
+
+
 @pytest.mark.asyncio
 async def test_drain_removes_sessions_and_rejects_new_registrations() -> None:
     registry = FabricSessionRegistry()