feat(cloud): harden multi-tenant runtime resources

src/langbot/pkg/provider/session/sessionmgr.py

@@ -2,6 +2,8 @@ from __future__ import annotations
 
 import asyncio
 import dataclasses
+import heapq
+import time
 
 from langbot_plugin.api.entities.builtin.provider import message as provider_message, prompt as provider_prompt
 import langbot_plugin.api.entities.builtin.provider.session as provider_session
@@ -24,6 +26,10 @@ SessionKey = tuple[
     str,
     int | str,
 ]
+SessionExpiryEntry = tuple[float, int, SessionKey]
+
+_SESSION_EXPIRY_HEAP_MIN_LIMIT = 64
+_SESSION_EXPIRY_HEAP_ACTIVE_MULTIPLIER = 4
 
 
 def _query_session_key(query: pipeline_query.Query) -> tuple[SessionKey, ExecutionContext]:
@@ -49,11 +55,215 @@ class SessionManager:
 
     ap: app.Application
 
-    session_list: list[provider_session.Session]
-
     def __init__(self, ap: app.Application):
         self.ap = ap
-        self.session_list = []
+        self._legacy_sessions: list[provider_session.Session] = []
+        self._session_index: dict[SessionKey, provider_session.Session] = {}
+        self._session_keys_by_workspace: dict[str, set[SessionKey]] = {}
+        self._session_expiry_heap: list[SessionExpiryEntry] = []
+        self._next_access_revision = 0
+
+    @property
+    def session_list(self) -> list[provider_session.Session]:
+        """Compatibility view for API services that enumerate sessions."""
+
+        return [
+            *self._legacy_sessions,
+            *self._session_index.values(),
+        ]
+
+    @session_list.setter
+    def session_list(self, sessions: list[provider_session.Session]) -> None:
+        """Replace the cache while keeping the O(1) index consistent."""
+
+        session_values = list(sessions)
+        self._legacy_sessions = []
+        self._session_index = {}
+        self._session_keys_by_workspace = {}
+        self._session_expiry_heap = []
+        self._next_access_revision = 0
+        now = time.monotonic()
+        for session in session_values:
+            key = getattr(session, '_langbot_session_key', None)
+            if isinstance(key, tuple) and len(key) == 6:
+                self._session_index[key] = session
+                self._session_keys_by_workspace.setdefault(key[1], set()).add(key)
+                last_accessed = getattr(session, '_langbot_last_accessed', None)
+                if last_accessed is None:
+                    last_accessed = now
+                self._touch_session(
+                    session,
+                    key,
+                    float(last_accessed),
+                    compact=False,
+                )
+            else:
+                self._legacy_sessions.append(session)
+        self._compact_session_expiry_heap(force=True)
+
+    def _retention_config(self) -> dict:
+        instance_config = getattr(getattr(self.ap, 'instance_config', None), 'data', {})
+        if not isinstance(instance_config, dict):
+            return {}
+        config = instance_config.get('system', {}).get('session_retention', {})
+        return config if isinstance(config, dict) else {}
+
+    def _positive_config_int(self, name: str, default: int) -> int:
+        try:
+            value = int(self._retention_config().get(name, default))
+        except (TypeError, ValueError):
+            value = default
+        return max(value, 1)
+
+    @staticmethod
+    def _session_is_idle(session: provider_session.Session) -> bool:
+        semaphore = getattr(session, '_semaphore', None)
+        concurrency = getattr(session, '_langbot_session_concurrency', None)
+        if semaphore is None or not isinstance(concurrency, int):
+            return True
+        return getattr(semaphore, '_value', -1) == concurrency
+
+    def _remove_session(self, session: provider_session.Session) -> None:
+        key = getattr(session, '_langbot_session_key', None)
+        if isinstance(key, tuple) and len(key) == 6 and self._session_index.get(key) is session:
+            self._session_index.pop(key, None)
+            workspace_keys = self._session_keys_by_workspace.get(key[1])
+            if workspace_keys is not None:
+                workspace_keys.discard(key)
+                if not workspace_keys:
+                    self._session_keys_by_workspace.pop(key[1], None)
+        else:
+            try:
+                self._legacy_sessions.remove(session)
+            except ValueError:
+                pass
+
+    def _touch_session(
+        self,
+        session: provider_session.Session,
+        key: SessionKey,
+        now: float,
+        *,
+        compact: bool = True,
+    ) -> None:
+        self._next_access_revision += 1
+        revision = self._next_access_revision
+        object.__setattr__(session, '_langbot_last_accessed', now)
+        object.__setattr__(session, '_langbot_access_revision', revision)
+        heapq.heappush(
+            self._session_expiry_heap,
+            (now, revision, key),
+        )
+        if compact:
+            self._compact_session_expiry_heap()
+
+    def _compact_session_expiry_heap(self, *, force: bool = False) -> None:
+        limit = max(
+            len(self._session_index) * _SESSION_EXPIRY_HEAP_ACTIVE_MULTIPLIER,
+            _SESSION_EXPIRY_HEAP_MIN_LIMIT,
+        )
+        if not force and len(self._session_expiry_heap) <= limit:
+            return
+        self._session_expiry_heap = [
+            (
+                float(getattr(session, '_langbot_last_accessed', 0.0)),
+                int(getattr(session, '_langbot_access_revision', 0)),
+                key,
+            )
+            for key, session in self._session_index.items()
+        ]
+        heapq.heapify(self._session_expiry_heap)
+
+    def _pop_current_expiry_entry(
+        self,
+    ) -> tuple[float, int, SessionKey, provider_session.Session] | None:
+        while self._session_expiry_heap:
+            last_accessed, revision, key = heapq.heappop(self._session_expiry_heap)
+            session = self._session_index.get(key)
+            if session is None:
+                continue
+            if getattr(session, '_langbot_access_revision', None) != revision:
+                continue
+            return last_accessed, revision, key, session
+        return None
+
+    def _prune_expired_sessions(self, now: float) -> None:
+        idle_ttl = self._positive_config_int('idle_ttl_seconds', 86400)
+        cutoff = now - idle_ttl
+        while self._session_expiry_heap:
+            last_accessed, _, _ = self._session_expiry_heap[0]
+            if last_accessed > cutoff:
+                break
+            current = self._pop_current_expiry_entry()
+            if current is None:
+                break
+            last_accessed, revision, key, session = current
+            if last_accessed > cutoff:
+                heapq.heappush(
+                    self._session_expiry_heap,
+                    (last_accessed, revision, key),
+                )
+                break
+            if self._session_is_idle(session):
+                self._remove_session(session)
+                continue
+            # The session became active without another cache lookup. Give it
+            # a fresh TTL instead of repeatedly examining the same expired
+            # entry or losing its future expiry record.
+            self._touch_session(session, key, now)
+
+    def _prune_workspace_capacity(
+        self,
+        workspace_uuid: str,
+        max_entries_per_workspace: int,
+    ) -> None:
+        workspace_keys = self._session_keys_by_workspace.get(workspace_uuid, set())
+        overflow = len(workspace_keys) - max_entries_per_workspace + 1
+        if overflow <= 0:
+            return
+        idle_workspace_sessions = sorted(
+            (
+                session
+                for key in tuple(workspace_keys)
+                if (session := self._session_index.get(key)) is not None and self._session_is_idle(session)
+            ),
+            key=lambda session: float(getattr(session, '_langbot_last_accessed', 0.0)),
+        )
+        for session in idle_workspace_sessions[:overflow]:
+            self._remove_session(session)
+
+    def _evict_oldest_idle_session(self, now: float) -> bool:
+        # At most one current entry per active session is examined. Stale heap
+        # revisions do not count and are discarded in O(log N).
+        current_probes = 0
+        max_probes = len(self._session_index)
+        while current_probes < max_probes:
+            current = self._pop_current_expiry_entry()
+            if current is None:
+                return False
+            _, _, key, session = current
+            current_probes += 1
+            if self._session_is_idle(session):
+                self._remove_session(session)
+                return True
+            self._touch_session(session, key, now)
+        return False
+
+    def _prune_sessions(self, now: float, workspace_uuid: str) -> None:
+        self._prune_expired_sessions(now)
+        max_entries_per_workspace = self._positive_config_int('max_entries_per_workspace', 200)
+        self._prune_workspace_capacity(
+            workspace_uuid,
+            max_entries_per_workspace,
+        )
+
+        max_entries = self._positive_config_int('max_entries', 2000)
+        overflow = len(self._session_index) - max_entries + 1
+        if overflow <= 0:
+            return
+        for _ in range(overflow):
+            if not self._evict_oldest_idle_session(now):
+                break
 
     async def initialize(self):
         pass
@@ -61,9 +271,25 @@ class SessionManager:
     async def get_session(self, query: pipeline_query.Query) -> provider_session.Session:
         """获取会话"""
         session_key, execution_context = _query_session_key(query)
-        for session in self.session_list:
-            if getattr(session, '_langbot_session_key', None) == session_key:
-                return session
+        now = time.monotonic()
+        session = self._session_index.get(session_key)
+        if session is not None:
+            self._touch_session(session, session_key, now)
+            return session
+
+        self._prune_sessions(now, execution_context.workspace_uuid)
+        max_entries_per_workspace = self._positive_config_int('max_entries_per_workspace', 200)
+        workspace_entries = len(
+            self._session_keys_by_workspace.get(
+                execution_context.workspace_uuid,
+                (),
+            )
+        )
+        if workspace_entries >= max_entries_per_workspace:
+            raise RuntimeError(f'Workspace session cache capacity reached ({max_entries_per_workspace})')
+        max_entries = self._positive_config_int('max_entries', 2000)
+        if len(self._session_index) >= max_entries:
+            raise RuntimeError(f'Session cache capacity reached ({max_entries})')
 
         session_concurrency = self.ap.instance_config.data['concurrency']['session']
 
@@ -93,8 +319,14 @@ class SessionManager:
         object.__setattr__(session, 'bot_uuid', query.bot_uuid)
         object.__setattr__(session, '_execution_context', session_context)
         object.__setattr__(session, '_langbot_session_key', session_key)
+        object.__setattr__(session, '_langbot_session_concurrency', session_concurrency)
         session._semaphore = asyncio.Semaphore(session_concurrency)
-        self.session_list.append(session)
+        self._session_index[session_key] = session
+        self._session_keys_by_workspace.setdefault(
+            execution_context.workspace_uuid,
+            set(),
+        ).add(session_key)
+        self._touch_session(session, session_key, now)
         return session
 
     async def get_conversation(
@@ -144,6 +376,47 @@ class SessionManager:
                 bot_uuid=bot_uuid,
             )
             session.conversations.append(conversation)
+            max_conversations = self._positive_config_int('max_conversations_per_session', 20)
+            if len(session.conversations) > max_conversations:
+                del session.conversations[:-max_conversations]
             session.using_conversation = conversation
 
         return session.using_conversation
+
+    def trim_conversation_messages(
+        self,
+        conversation: provider_session.Conversation,
+        *,
+        max_rounds: int,
+    ) -> None:
+        """Bound retained process-local history after a completed turn."""
+
+        try:
+            max_rounds = int(max_rounds)
+        except (TypeError, ValueError):
+            max_rounds = 10
+        max_rounds = max(max_rounds, 1)
+        max_messages = self._positive_config_int('max_messages_per_conversation', 100)
+
+        kept_reversed = []
+        user_rounds = 0
+        for message in reversed(conversation.messages):
+            if user_rounds >= max_rounds:
+                break
+            kept_reversed.append(message)
+            if getattr(message, 'role', None) == 'user':
+                user_rounds += 1
+        retained = list(reversed(kept_reversed))[-max_messages:]
+        # Binary payloads are needed for the current model call, but retaining
+        # them in process-local history makes a few image/file turns consume
+        # hundreds of MB. Historical URL and text references remain intact.
+        for message in retained:
+            content = getattr(message, 'content', None)
+            if not isinstance(content, list):
+                continue
+            for element in content:
+                if getattr(element, 'image_base64', None) is not None:
+                    element.image_base64 = None
+                if getattr(element, 'file_base64', None) is not None:
+                    element.file_base64 = None
+        conversation.messages = retained