scriptoriumadmin/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui.git synced 2026-10-18 04:43:55 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

refac

Browse source
This commit is contained in:
Timothy Jaeryang Baek 2026-10-16 14:48:08 +04:00
parent c3c3087645
commit 362b01d0b2
2 changed files with 80 additions and 19 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
.env.example
View file

@ -6,7 +6,7 @@ OPENAI_API_BASE_URL=''
OPENAI_API_KEY=''
# With WEBSOCKET_MANAGER=redis, room channels require PUBLISH, SUBSCRIBE and
# PSUBSCRIBE permissions plus channel ACLs &socketio and &socketio#*.
# UNSUBSCRIBE/PING permissions plus channel ACLs &socketio and &socketio#*.
# Set false for shared-channel delivery. This must match on every instance;
# changing it requires a full fleet restart, not a rolling restart.
# WEBSOCKET_REDIS_ROOM_CHANNELS=true

97
backend/open_webui/socket/redis_room_channels.py
View file

@ -11,7 +11,7 @@ class AsyncRedisRoomChannelManager(AsyncRedisManager):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self._local_room_channels = set()
self._subscription_queue = asyncio.Queue()
# collision-free while namespaces contain no '#' (socket.io default '/'); rooms may contain '#'
def _room_channel(self, namespace, room):
@ -20,12 +20,29 @@ class AsyncRedisRoomChannelManager(AsyncRedisManager):
def basic_enter_room(self, sid, namespace, room, eio_sid=None):
super().basic_enter_room(sid, namespace, room, eio_sid=eio_sid)
if room is not None:
self._local_room_channels.add(self._room_channel(namespace, room))
self._subscription_queue.put_nowait(asyncio.Event())
def basic_leave_room(self, sid, namespace, room):
super().basic_leave_room(sid, namespace, room)
if room is not None and room not in self.rooms.get(namespace, {}):
self._local_room_channels.discard(self._room_channel(namespace, room))
self._subscription_queue.put_nowait(asyncio.Event())
async def _wait_for_rooms(self):
ready = asyncio.Event()
self._subscription_queue.put_nowait(ready)
await ready.wait()
async def connect(self, eio_sid, namespace):
sid = await super().connect(eio_sid, namespace)
if sid is not None:
await self._wait_for_rooms()
return sid
async def enter_room(self, sid, namespace, room, eio_sid=None):
local = self.is_connected(sid, namespace)
await super().enter_room(sid, namespace, room, eio_sid=eio_sid)
if local and room is not None:
await self._wait_for_rooms()
async def _publish(self, data):
if data.get('method') == 'emit' and isinstance(data.get('room'), str):
@ -57,23 +74,73 @@ class AsyncRedisRoomChannelManager(AsyncRedisManager):
)
break
async def _sync_rooms(self, pubsub):
wanted = {self.channel.encode()} | {
self._room_channel(namespace, room)
for namespace, rooms in self.rooms.items()
for room in rooms
if room is not None
}
if added := wanted - (pubsub.channels.keys() - pubsub.pending_unsubscribe_channels):
await pubsub.subscribe(*added)
if removed := pubsub.channels.keys() - wanted - pubsub.pending_unsubscribe_channels:
await pubsub.unsubscribe(*removed)
async def _listen_to_rooms(self, pubsub):
pubsub.ignore_subscribe_messages = False
self._subscription_queue.put_nowait(asyncio.Event())
changed = asyncio.create_task(self._subscription_queue.get())
message = asyncio.create_task(pubsub.get_message(timeout=None))
waiters = []
serial = 0
try:
while True:
await asyncio.wait([message] if waiters else [message, changed], return_when=asyncio.FIRST_COMPLETED)
if not waiters and changed.done():
waiters.append(changed.result())
while not self._subscription_queue.empty():
waiters.append(self._subscription_queue.get_nowait())
changed = asyncio.create_task(self._subscription_queue.get())
await self._sync_rooms(pubsub)
# PONG confirms Redis processed the subscriptions before a join returns.
serial += 1
await pubsub.ping(str(serial).encode())
if message.done():
data = message.result() or {}
message = asyncio.create_task(pubsub.get_message(timeout=None))
if data.get('type') == 'subscribe':
# redis-py can reconnect internally and restore a pending removal.
self._subscription_queue.put_nowait(asyncio.Event())
elif data.get('type') == 'pong' and data['data'] == str(serial).encode():
for ready in waiters:
ready.set()
waiters.clear()
else:
yield data
finally:
if changed.done() and not changed.cancelled():
waiters.append(changed.result())
for ready in waiters:
self._subscription_queue.put_nowait(ready)
message.cancel()
changed.cancel()
await asyncio.gather(message, changed, return_exceptions=True)
await pubsub.aclose()
async def _redis_listen_with_retries(self):
_, error = self._get_redis_module_and_error()
retry_sleep = 1
subscribed = False
while True:
try:
if not subscribed:
self._redis_connect()
await self.pubsub.subscribe(self.channel)
await self.pubsub.psubscribe(f'{self.channel}#*')
retry_sleep = 1
async for message in self.pubsub.listen():
self._redis_connect()
await self.pubsub.subscribe(self.channel)
retry_sleep = 1
async for message in self._listen_to_rooms(self.pubsub):
yield message
except error as exc:
if isinstance(exc, NoPermissionError):
self._get_logger().error(
'Redis denied subscribing: %s. Check SUBSCRIBE/PSUBSCRIBE permissions and channel ACLs '
'Redis denied subscribing: %s. Check SUBSCRIBE/UNSUBSCRIBE/PING permissions and channel ACLs '
'(&%s and &%s#*). To disable room channels, set '
'WEBSOCKET_REDIS_ROOM_CHANNELS=False on every instance and fully restart the fleet.',
exc,
@ -84,18 +151,12 @@ class AsyncRedisRoomChannelManager(AsyncRedisManager):
f'Cannot receive from redis... retrying in {retry_sleep} secs',
extra={'redis_exception': str(exc)},
)
subscribed = False
await asyncio.sleep(retry_sleep)
retry_sleep *= 2
if retry_sleep > 60:
retry_sleep = 60
async def _listen(self):
main_channel = self.channel.encode()
async for message in self._redis_listen_with_retries():
if 'data' not in message:
continue
if (message['type'] == 'message' and message['channel'] == main_channel) or (
message['type'] == 'pmessage' and message['channel'] in self._local_room_channels
):
if message.get('type') == 'message':
yield message['data']