mirror of
https://github.com/open-webui/open-webui.git
synced 2026-11-08 08:27:19 +00:00
9189033c37
A note written outside the editor (the write_note tool, or any POST to /api/v1/notes/create) has no Yjs update history, so the first client to open it received an empty document and the save that followed wrote that emptiness back over the note. The content was still in the database up to that moment, then gone, with only the title left. The server now builds the Yjs document from the note's stored ProseMirror json when there is no update history yet, so a client opens the note with its real content instead of a blank document. The client id is derived from the content, which makes two workers seeding the same note a no-op rather than a duplication. Notes that only carry markdown still rely on the editor's own seeding, since building ProseMirror json from markdown server-side would mean reimplementing the editor's schema in Python. That path was also gated on being the only session joined to the note, so a second viewer disabled it entirely; it now elects the lowest joined session, which keeps a single seeder without giving up on seeding.
368 lines
14 KiB
Python
368 lines
14 KiB
Python
"""Redis-backed distributed data structures for WebSocket state management."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import uuid
|
|
|
|
import pycrdt as Y
|
|
from open_webui.env import REDIS_KEY_PREFIX
|
|
from open_webui.utils.json_codec import JSONCodec
|
|
from open_webui.utils.redis import get_redis_connection
|
|
|
|
YDOC_KEY_PREFIX = f'{REDIS_KEY_PREFIX}:ydoc:documents'
|
|
|
|
|
|
class RedisLock:
|
|
"""Distributed lock backed by a Redis SET with NX/EX semantics."""
|
|
|
|
_RENEW_SCRIPT = """
|
|
if redis.call('get', KEYS[1]) == ARGV[1] then
|
|
return redis.call('expire', KEYS[1], ARGV[2])
|
|
end
|
|
return 0
|
|
"""
|
|
_RELEASE_SCRIPT = """
|
|
if redis.call('get', KEYS[1]) == ARGV[1] then
|
|
return redis.call('del', KEYS[1])
|
|
end
|
|
return 0
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
redis_url,
|
|
lock_name,
|
|
timeout_secs,
|
|
redis_sentinels=[],
|
|
redis_cluster=False,
|
|
):
|
|
self.lock_name = lock_name
|
|
self.lock_id = str(uuid.uuid4())
|
|
self.timeout_secs = timeout_secs
|
|
self.lock_obtained = False
|
|
self.redis = get_redis_connection(
|
|
redis_url,
|
|
redis_sentinels,
|
|
redis_cluster=redis_cluster,
|
|
decode_responses=True,
|
|
)
|
|
|
|
def aquire_lock(self):
|
|
# nx=True will only set this key if it _hasn't_ already been set
|
|
self.lock_obtained = self.redis.set(self.lock_name, self.lock_id, nx=True, ex=self.timeout_secs)
|
|
return self.lock_obtained
|
|
|
|
def renew_lock(self):
|
|
return bool(self.redis.eval(self._RENEW_SCRIPT, 1, self.lock_name, self.lock_id, self.timeout_secs))
|
|
|
|
def release_lock(self):
|
|
self.redis.eval(self._RELEASE_SCRIPT, 1, self.lock_name, self.lock_id)
|
|
|
|
|
|
class RedisDict:
|
|
def __init__(self, name, redis_url, redis_sentinels=[], redis_cluster=False):
|
|
self.name = name
|
|
# Per-process cache of the last payload fingerprint written by set().
|
|
# Used to skip redundant HSET round-trips when the model list hasn't
|
|
# changed — the dominant Redis write source on busy multi-pod setups.
|
|
self._last_signature: str | None = None
|
|
self.redis = get_redis_connection(
|
|
redis_url,
|
|
redis_sentinels,
|
|
redis_cluster=redis_cluster,
|
|
decode_responses=True,
|
|
)
|
|
|
|
def __setitem__(self, key, value):
|
|
serialized_value = JSONCodec.dumps(value)
|
|
self.redis.hset(self.name, key, serialized_value)
|
|
|
|
def __getitem__(self, key):
|
|
value = self.redis.hget(self.name, key)
|
|
if value is None:
|
|
raise KeyError(key)
|
|
return JSONCodec.loads(value)
|
|
|
|
def __delitem__(self, key):
|
|
result = self.redis.hdel(self.name, key)
|
|
if result == 0:
|
|
raise KeyError(key)
|
|
|
|
def __contains__(self, key):
|
|
return self.redis.hexists(self.name, key)
|
|
|
|
def __len__(self):
|
|
return self.redis.hlen(self.name)
|
|
|
|
def keys(self):
|
|
return self.redis.hkeys(self.name)
|
|
|
|
def values(self):
|
|
return [JSONCodec.loads(v) for v in self.redis.hvals(self.name)]
|
|
|
|
def items(self):
|
|
return [(k, JSONCodec.loads(v)) for k, v in self.redis.hgetall(self.name).items()]
|
|
|
|
def set(self, mapping: dict):
|
|
if not mapping:
|
|
self.redis.delete(self.name)
|
|
self._last_signature = None
|
|
return
|
|
|
|
# Serialize values once — reused for both the fingerprint and the write.
|
|
serialized = {k: JSONCodec.dumps(v) for k, v in mapping.items()}
|
|
digest = hashlib.sha256()
|
|
for key in sorted(serialized):
|
|
digest.update(key.encode())
|
|
digest.update(b'\0')
|
|
digest.update(serialized[key].encode())
|
|
digest.update(b'\0')
|
|
signature = digest.hexdigest()
|
|
|
|
# Skip the write when the prepared mapping is identical to the last one
|
|
# this process wrote. The check is per-instance (not distributed), but
|
|
# still eliminates the majority of redundant writes because each pod
|
|
# typically produces the same model list on consecutive refreshes.
|
|
if signature == self._last_signature:
|
|
return
|
|
|
|
# Fetch existing keys before writing so we know which ones to remove.
|
|
# HKEYS is cheap — it transfers only short key strings, not large JSON values.
|
|
existing_keys = set(self.redis.hkeys(self.name))
|
|
new_keys = set(mapping.keys())
|
|
keys_to_remove = existing_keys - new_keys
|
|
|
|
# HSET first (add/update all new values), then HDEL (remove stale keys).
|
|
# We never DELETE the whole hash — this eliminates the race window
|
|
# where concurrent readers would see an empty models dict.
|
|
self.redis.hset(self.name, mapping=serialized)
|
|
if keys_to_remove:
|
|
self.redis.hdel(self.name, *keys_to_remove)
|
|
|
|
self._last_signature = signature
|
|
|
|
def get(self, key, default=None):
|
|
try:
|
|
return self[key]
|
|
except KeyError:
|
|
return default
|
|
|
|
def clear(self):
|
|
self.redis.delete(self.name)
|
|
self._last_signature = None
|
|
|
|
def update(self, other=None, **kwargs):
|
|
if other is not None:
|
|
for k, v in other.items() if hasattr(other, 'items') else other:
|
|
self[k] = v
|
|
for k, v in kwargs.items():
|
|
self[k] = v
|
|
|
|
def setdefault(self, key, default=None):
|
|
if key not in self:
|
|
self[key] = default
|
|
return self[key]
|
|
|
|
|
|
def _append_prosemirror_nodes(parent, nodes) -> None:
|
|
# Children are appended before being written to: pycrdt only accepts writes once integrated.
|
|
text_run = None # a run of text nodes shares one XmlText, as y-prosemirror lays them out
|
|
|
|
for node in nodes or []:
|
|
if not isinstance(node, dict) or not isinstance(node.get('type'), str):
|
|
continue
|
|
|
|
if node['type'] == 'text':
|
|
if not isinstance(node.get('text'), str) or not node['text']:
|
|
continue
|
|
if text_run is None:
|
|
text_run = Y.XmlText()
|
|
parent.children.append(text_run)
|
|
marks = {
|
|
mark['type']: mark.get('attrs') or {}
|
|
for mark in node.get('marks') or []
|
|
if isinstance(mark, dict) and mark.get('type')
|
|
}
|
|
text_run.insert(len(text_run), node['text'], marks)
|
|
continue
|
|
|
|
text_run = None
|
|
element = Y.XmlElement(node['type'])
|
|
parent.children.append(element)
|
|
for key, value in (node.get('attrs') or {}).items():
|
|
if value is not None:
|
|
element.attributes[key] = value
|
|
_append_prosemirror_nodes(element, node.get('content'))
|
|
|
|
|
|
def build_prosemirror_update(content_json) -> bytes | None:
|
|
"""Encode a ProseMirror document as a Yjs update in the shape y-prosemirror reads back."""
|
|
if not isinstance(content_json, dict) or content_json.get('type') != 'doc' or not content_json.get('content'):
|
|
return None
|
|
|
|
# A content-derived client id makes a concurrent seed of the same content a no-op.
|
|
canonical = json.dumps(content_json, sort_keys=True, separators=(',', ':'))
|
|
content_digest = hashlib.sha256(canonical.encode()).digest()
|
|
doc = Y.Doc(client_id=int.from_bytes(content_digest[:4], 'big'))
|
|
fragment = Y.XmlFragment()
|
|
doc['prosemirror'] = fragment
|
|
_append_prosemirror_nodes(fragment, content_json['content'])
|
|
return doc.get_update()
|
|
|
|
|
|
class YdocManager:
|
|
COMPACTION_THRESHOLD = 500
|
|
|
|
def __init__(
|
|
self,
|
|
redis=None,
|
|
redis_key_prefix: str = YDOC_KEY_PREFIX,
|
|
):
|
|
self._updates = {}
|
|
self._users = {}
|
|
self._redis = redis
|
|
self._redis_key_prefix = redis_key_prefix
|
|
|
|
async def append_to_updates(self, document_id: str, update: bytes):
|
|
document_id = document_id.replace(':', '_')
|
|
if self._redis:
|
|
redis_key = f'{self._redis_key_prefix}:{document_id}:updates'
|
|
await self._redis.rpush(redis_key, JSONCodec.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(JSONCodec.loads(raw)))
|
|
snapshot = JSONCodec.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(':', '_')
|
|
|
|
if self._redis:
|
|
redis_key = f'{self._redis_key_prefix}:{document_id}:updates'
|
|
updates = await self._redis.lrange(redis_key, 0, -1)
|
|
return [bytes(JSONCodec.loads(update)) for update in updates]
|
|
else:
|
|
return self._updates.get(document_id, [])
|
|
|
|
async def document_exists(self, document_id: str) -> bool:
|
|
document_id = document_id.replace(':', '_')
|
|
|
|
if self._redis:
|
|
redis_key = f'{self._redis_key_prefix}:{document_id}:updates'
|
|
return await self._redis.exists(redis_key) > 0
|
|
else:
|
|
return document_id in self._updates
|
|
|
|
async def get_users(self, document_id: str) -> list[str]:
|
|
document_id = document_id.replace(':', '_')
|
|
|
|
if self._redis:
|
|
redis_key = f'{self._redis_key_prefix}:{document_id}:users'
|
|
users = await self._redis.smembers(redis_key)
|
|
return list(users)
|
|
else:
|
|
return self._users.get(document_id, [])
|
|
|
|
async def add_user(self, document_id: str, user_id: str):
|
|
document_id = document_id.replace(':', '_')
|
|
|
|
if self._redis:
|
|
redis_key = f'{self._redis_key_prefix}:{document_id}:users'
|
|
await self._redis.sadd(redis_key, user_id)
|
|
# Maintain a per-session reverse index so disconnect cleanup
|
|
# can look up only the documents this session joined, instead
|
|
# of issuing a cluster-wide SCAN over the entire keyspace.
|
|
session_key = f'{self._redis_key_prefix}:session:{user_id}:documents'
|
|
await self._redis.sadd(session_key, document_id)
|
|
else:
|
|
if document_id not in self._users:
|
|
self._users[document_id] = set()
|
|
self._users[document_id].add(user_id)
|
|
|
|
async def remove_user(self, document_id: str, user_id: str):
|
|
document_id = document_id.replace(':', '_')
|
|
|
|
if self._redis:
|
|
redis_key = f'{self._redis_key_prefix}:{document_id}:users'
|
|
await self._redis.srem(redis_key, user_id)
|
|
# Keep the reverse index in sync.
|
|
session_key = f'{self._redis_key_prefix}:session:{user_id}:documents'
|
|
await self._redis.srem(session_key, document_id)
|
|
else:
|
|
if document_id in self._users and user_id in self._users[document_id]:
|
|
self._users[document_id].remove(user_id)
|
|
|
|
async def remove_user_from_all_documents(self, user_id: str):
|
|
if self._redis:
|
|
# Use the per-session reverse index instead of a cluster-wide
|
|
# SCAN. This set contains only the document IDs that this
|
|
# session actually joined, so the cost is proportional to
|
|
# the session's footprint — not the total number of documents.
|
|
session_key = f'{self._redis_key_prefix}:session:{user_id}:documents'
|
|
document_ids = await self._redis.smembers(session_key)
|
|
|
|
for document_id in document_ids:
|
|
users_key = f'{self._redis_key_prefix}:{document_id}:users'
|
|
await self._redis.srem(users_key, user_id)
|
|
|
|
if len(await self.get_users(document_id)) == 0:
|
|
await self.clear_document(document_id)
|
|
|
|
# Clean up the reverse index itself.
|
|
await self._redis.delete(session_key)
|
|
|
|
else:
|
|
for document_id in list(self._users.keys()):
|
|
if user_id in self._users[document_id]:
|
|
self._users[document_id].remove(user_id)
|
|
if not self._users[document_id]:
|
|
del self._users[document_id]
|
|
|
|
await self.clear_document(document_id)
|
|
|
|
async def clear_document(self, document_id: str):
|
|
document_id = document_id.replace(':', '_')
|
|
|
|
if self._redis:
|
|
redis_key = f'{self._redis_key_prefix}:{document_id}:updates'
|
|
await self._redis.delete(redis_key)
|
|
redis_users_key = f'{self._redis_key_prefix}:{document_id}:users'
|
|
await self._redis.delete(redis_users_key)
|
|
else:
|
|
if document_id in self._updates:
|
|
del self._updates[document_id]
|
|
if document_id in self._users:
|
|
del self._users[document_id]
|