refac

backend/open_webui/main.py

@@ -64,6 +64,7 @@ from open_webui.socket.main import (
     MODELS,
     app as socket_app,
     periodic_usage_pool_cleanup,
+    periodic_session_pool_cleanup,
     get_event_emitter,
     get_models_in_use,
 )
@@ -635,6 +636,7 @@ async def lifespan(app: FastAPI):
         limiter.total_tokens = THREAD_POOL_SIZE
 
     asyncio.create_task(periodic_usage_pool_cleanup())
+    asyncio.create_task(periodic_session_pool_cleanup())
 
     if app.state.config.ENABLE_BASE_MODELS_CACHE:
         await get_all_models(

backend/open_webui/socket/main.py

@@ -99,6 +99,7 @@ else:
 
 # Timeout duration in seconds
 TIMEOUT_DURATION = 3
+SESSION_POOL_TIMEOUT = 120  # seconds without heartbeat before session is reaped
 
 # Dictionary to maintain the user pool
 
@@ -147,6 +148,17 @@ if WEBSOCKET_MANAGER == "redis":
     aquire_func = clean_up_lock.aquire_lock
     renew_func = clean_up_lock.renew_lock
     release_func = clean_up_lock.release_lock
+
+    session_cleanup_lock = RedisLock(
+        redis_url=WEBSOCKET_REDIS_URL,
+        lock_name=f"{REDIS_KEY_PREFIX}:session_cleanup_lock",
+        timeout_secs=WEBSOCKET_REDIS_LOCK_TIMEOUT,
+        redis_sentinels=redis_sentinels,
+        redis_cluster=WEBSOCKET_REDIS_CLUSTER,
+    )
+    session_aquire_func = session_cleanup_lock.aquire_lock
+    session_renew_func = session_cleanup_lock.renew_lock
+    session_release_func = session_cleanup_lock.release_lock
 else:
     MODELS = {}
 
@@ -154,6 +166,7 @@ else:
     USAGE_POOL = {}
 
     aquire_func = release_func = renew_func = lambda: True
+    session_aquire_func = session_release_func = session_renew_func = lambda: True
 
 
 YDOC_MANAGER = YdocManager(
@@ -162,6 +175,31 @@ YDOC_MANAGER = YdocManager(
 )
 
 
+async def periodic_session_pool_cleanup():
+    """Reap orphaned SESSION_POOL entries that missed heartbeats (e.g. crashed instance)."""
+    if not session_aquire_func():
+        log.debug("Session cleanup lock held by another node. Skipping.")
+        return
+
+    try:
+        while True:
+            if not session_renew_func():
+                log.error("Unable to renew session cleanup lock. Exiting.")
+                return
+
+            now = int(time.time())
+            for sid in list(SESSION_POOL.keys()):
+                entry = SESSION_POOL.get(sid)
+                if entry and now - entry.get("last_seen_at", 0) > SESSION_POOL_TIMEOUT:
+                    log.warning(
+                        f"Reaping orphaned session {sid} (user {entry.get('id')})"
+                    )
+                    del SESSION_POOL[sid]
+            await asyncio.sleep(SESSION_POOL_TIMEOUT)
+    finally:
+        session_release_func()
+
+
 async def periodic_usage_pool_cleanup():
     max_retries = 2
     retry_delay = random.uniform(
@@ -313,15 +351,18 @@ async def connect(sid, environ, auth):
             user = Users.get_user_by_id(data["id"])
 
         if user:
-            SESSION_POOL[sid] = user.model_dump(
+            SESSION_POOL[sid] = {
-                exclude=[
+                **user.model_dump(
-                    "profile_image_url",
+                    exclude=[
-                    "profile_banner_image_url",
+                        "profile_image_url",
-                    "date_of_birth",
+                        "profile_banner_image_url",
-                    "bio",
+                        "date_of_birth",
-                    "gender",
+                        "bio",
-                ]
+                        "gender",
-            )
+                    ]
+                ),
+                "last_seen_at": int(time.time()),
+            }
             await sio.enter_room(sid, f"user:{user.id}")
 
 
@@ -340,15 +381,18 @@ async def user_join(sid, data):
     if not user:
         return
 
-    SESSION_POOL[sid] = user.model_dump(
+    SESSION_POOL[sid] = {
-        exclude=[
+        **user.model_dump(
-            "profile_image_url",
+            exclude=[
-            "profile_banner_image_url",
+                "profile_image_url",
-            "date_of_birth",
+                "profile_banner_image_url",
-            "bio",
+                "date_of_birth",
-            "gender",
+                "bio",
-        ]
+                "gender",
-    )
+            ]
+        ),
+        "last_seen_at": int(time.time()),
+    }
 
     await sio.enter_room(sid, f"user:{user.id}")
 
@@ -366,6 +410,7 @@ async def user_join(sid, data):
 async def heartbeat(sid, data):
     user = SESSION_POOL.get(sid)
     if user:
+        SESSION_POOL[sid] = {**user, "last_seen_at": int(time.time())}
         Users.update_last_active_by_id(user["id"])
 
 
@@ -709,6 +754,17 @@ async def disconnect(sid):
     if sid in SESSION_POOL:
         user = SESSION_POOL[sid]
         del SESSION_POOL[sid]
+
+        # Clean up USAGE_POOL entries for this session
+        for model_id in list(USAGE_POOL.keys()):
+            connections = USAGE_POOL.get(model_id)
+            if connections and sid in connections:
+                del connections[sid]
+                if not connections:
+                    del USAGE_POOL[model_id]
+                else:
+                    USAGE_POOL[model_id] = connections
+
         await YDOC_MANAGER.remove_user_from_all_documents(sid)
     else:
         pass

backend/open_webui/socket/utils.py

@@ -118,6 +118,8 @@ class RedisDict:
 
 
 class YdocManager:
+    COMPACTION_THRESHOLD = 500
+
     def __init__(
         self,
         redis=None,
@@ -133,10 +135,42 @@ class YdocManager:
         if self._redis:
             redis_key = f"{self._redis_key_prefix}:{document_id}:updates"
             await self._redis.rpush(redis_key, json.dumps(list(update)))
+            list_len = await self._redis.llen(redis_key)
+            if list_len >= self.COMPACTION_THRESHOLD:
+                await self._compact_updates_redis(document_id)
         else:
             if document_id not in self._updates:
                 self._updates[document_id] = []
             self._updates[document_id].append(update)
+            if len(self._updates[document_id]) >= self.COMPACTION_THRESHOLD:
+                self._compact_updates_memory(document_id)
+
+    async def _compact_updates_redis(self, document_id: str):
+        """Rolling compaction: squash oldest half into one snapshot."""
+        redis_key = f"{self._redis_key_prefix}:{document_id}:updates"
+        all_updates = await self._redis.lrange(redis_key, 0, -1)
+        if len(all_updates) <= 1:
+            return
+        mid = len(all_updates) // 2
+        ydoc = Y.Doc()
+        for raw in all_updates[:mid]:
+            ydoc.apply_update(bytes(json.loads(raw)))
+        snapshot = json.dumps(list(ydoc.get_update()))
+        pipe = self._redis.pipeline()
+        pipe.delete(redis_key)
+        pipe.rpush(redis_key, snapshot, *all_updates[mid:])
+        await pipe.execute()
+
+    def _compact_updates_memory(self, document_id: str):
+        """Rolling compaction: squash oldest half into one snapshot."""
+        updates = self._updates.get(document_id, [])
+        if len(updates) <= 1:
+            return
+        mid = len(updates) // 2
+        ydoc = Y.Doc()
+        for update in updates[:mid]:
+            ydoc.apply_update(bytes(update))
+        self._updates[document_id] = [ydoc.get_update()] + updates[mid:]
 
     async def get_updates(self, document_id: str) -> List[bytes]:
         document_id = document_id.replace(":", "_")