From 498a04dc6c5baa6bf1de39b68cec4d8555bc8909 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 20 Apr 2026 14:24:33 +0000 Subject: [PATCH] feat(plugins): cross-worker tool + function cache invalidation via Redis MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bug: request.app.state.TOOLS / FUNCTIONS are per-worker Python dicts. Save on worker A does not touch worker B's cache, so worker B keeps serving the stale compiled module until the process restarts. utils/tools.py's invocation path was also completely cache-blind — a bare dict.get with no content-hash fallback, so even same-worker edits skipped the TOOL_CONTENTS check and served the first-seen module forever. Symptom: editing a tool in the Workspace UI, triggering a chat, and watching the OLD code execute until backend restart. Confirmed on a multi-worker deployment (UVICORN_WORKERS > 1) with Redis. Fix — three parts: 1. backend/open_webui/utils/plugin_cache.py (new) * publish_invalidation(app, kind, id) - drops the local TOOLS/FUNCTIONS + _CONTENTS entry - publishes {kind, id} on REDIS_PLUGIN_CACHE_CHANNEL * plugin_cache_listener(app) - subscribes on startup, drops caches on every inbound message - mirrors the existing redis_task_command_listener pattern * falls back to pure local invalidation when Redis is unconfigured 2. utils/tools.py * invocation path swapped from `TOOLS.get(tool_id)` to `await get_tool_module_from_cache(request, tool_id)`, which does content-hash invalidation from DB. Covers edits that predate or race the Redis pub/sub message. 3. routers/tools.py + routers/functions.py * create / update / delete / toggle / toggle_global all call publish_invalidation after the DB mutation commits. * toggle_function_by_id + toggle_global_by_id gained a `request` parameter they were missing. main.py lifespan starts the listener alongside the task listener and cancels it on shutdown. No new env vars — reuses REDIS_URL and REDIS_KEY_PREFIX. Single-worker deployments still work: publish_invalidation unconditionally drops the local cache before trying to publish. --- backend/open_webui/main.py | 4 + backend/open_webui/routers/functions.py | 20 ++++- backend/open_webui/routers/tools.py | 4 + backend/open_webui/utils/plugin_cache.py | 110 +++++++++++++++++++++++ backend/open_webui/utils/tools.py | 14 +-- 5 files changed, 145 insertions(+), 7 deletions(-) create mode 100644 backend/open_webui/utils/plugin_cache.py diff --git a/backend/open_webui/main.py b/backend/open_webui/main.py index ba7f74c830..e0ca6ef40f 100644 --- a/backend/open_webui/main.py +++ b/backend/open_webui/main.py @@ -583,6 +583,7 @@ from open_webui.tasks import ( stop_item_tasks, list_tasks, ) # Import from tasks.py +from open_webui.utils.plugin_cache import plugin_cache_listener from open_webui.utils.redis import get_sentinels_from_env @@ -666,6 +667,7 @@ async def lifespan(app: FastAPI): if app.state.redis is not None: app.state.redis_task_command_listener = asyncio.create_task(redis_task_command_listener(app)) + app.state.plugin_cache_listener = asyncio.create_task(plugin_cache_listener(app)) if THREAD_POOL_SIZE and THREAD_POOL_SIZE > 0: limiter = anyio.to_thread.current_default_thread_limiter() @@ -741,6 +743,8 @@ async def lifespan(app: FastAPI): if hasattr(app.state, 'redis_task_command_listener'): app.state.redis_task_command_listener.cancel() + if hasattr(app.state, 'plugin_cache_listener'): + app.state.plugin_cache_listener.cancel() app = FastAPI( diff --git a/backend/open_webui/routers/functions.py b/backend/open_webui/routers/functions.py index baec1f0870..c9ead8698d 100644 --- a/backend/open_webui/routers/functions.py +++ b/backend/open_webui/routers/functions.py @@ -23,6 +23,7 @@ from open_webui.utils.plugin import ( ) from open_webui.config import CACHE_DIR from open_webui.constants import ERROR_MESSAGES +from open_webui.utils.plugin_cache import publish_invalidation from fastapi import APIRouter, Depends, HTTPException, Request, status from open_webui.utils.auth import get_admin_user, get_verified_user from pydantic import BaseModel, HttpUrl @@ -214,6 +215,7 @@ async def create_new_function( await Functions.update_function_metadata_by_id(form_data.id, {'toggle': True}, db=db) if function: + await publish_invalidation(request.app, 'function', form_data.id) return function else: raise HTTPException( @@ -257,12 +259,18 @@ async def get_function_by_id(id: str, user=Depends(get_admin_user), db: AsyncSes @router.post('/id/{id}/toggle', response_model=Optional[FunctionModel]) -async def toggle_function_by_id(id: str, user=Depends(get_admin_user), db: AsyncSession = Depends(get_async_session)): +async def toggle_function_by_id( + request: Request, + id: str, + user=Depends(get_admin_user), + db: AsyncSession = Depends(get_async_session), +): function = await Functions.get_function_by_id(id, db=db) if function: function = await Functions.update_function_by_id(id, {'is_active': not function.is_active}, db=db) if function: + await publish_invalidation(request.app, 'function', id) return function else: raise HTTPException( @@ -282,12 +290,18 @@ async def toggle_function_by_id(id: str, user=Depends(get_admin_user), db: Async @router.post('/id/{id}/toggle/global', response_model=Optional[FunctionModel]) -async def toggle_global_by_id(id: str, user=Depends(get_admin_user), db: AsyncSession = Depends(get_async_session)): +async def toggle_global_by_id( + request: Request, + id: str, + user=Depends(get_admin_user), + db: AsyncSession = Depends(get_async_session), +): function = await Functions.get_function_by_id(id, db=db) if function: function = await Functions.update_function_by_id(id, {'is_global': not function.is_global}, db=db) if function: + await publish_invalidation(request.app, 'function', id) return function else: raise HTTPException( @@ -331,6 +345,7 @@ async def update_function_by_id( await Functions.update_function_metadata_by_id(id, {'toggle': True}, db=db) if function: + await publish_invalidation(request.app, 'function', id) return function else: raise HTTPException( @@ -363,6 +378,7 @@ async def delete_function_by_id( FUNCTIONS = request.app.state.FUNCTIONS if id in FUNCTIONS: del FUNCTIONS[id] + await publish_invalidation(request.app, 'function', id) return result diff --git a/backend/open_webui/routers/tools.py b/backend/open_webui/routers/tools.py index 4c3e77e566..c16937add8 100644 --- a/backend/open_webui/routers/tools.py +++ b/backend/open_webui/routers/tools.py @@ -29,6 +29,7 @@ from open_webui.utils.plugin import ( resolve_valves_schema_options, ) from open_webui.utils.tools import get_tool_specs +from open_webui.utils.plugin_cache import publish_invalidation from open_webui.utils.auth import get_admin_user, get_verified_user from open_webui.utils.access_control import ( has_permission, @@ -375,6 +376,7 @@ async def create_new_tools( tool_cache_dir.mkdir(parents=True, exist_ok=True) if tools: + await publish_invalidation(request.app, 'tool', form_data.id) return tools else: raise HTTPException( @@ -505,6 +507,7 @@ async def update_tools_by_id( tools = await Tools.update_tool_by_id(id, updated, db=db) if tools: + await publish_invalidation(request.app, 'tool', id) return tools else: raise HTTPException( @@ -612,6 +615,7 @@ async def delete_tools_by_id( TOOLS = request.app.state.TOOLS if id in TOOLS: del TOOLS[id] + await publish_invalidation(request.app, 'tool', id) return result diff --git a/backend/open_webui/utils/plugin_cache.py b/backend/open_webui/utils/plugin_cache.py new file mode 100644 index 0000000000..1884747251 --- /dev/null +++ b/backend/open_webui/utils/plugin_cache.py @@ -0,0 +1,110 @@ +# plugin_cache.py +# +# Cross-worker invalidation for tool + function module caches. +# +# Tool / function source code lives in the DB (shared across workers), +# but the compiled Python module is cached per-worker in +# ``request.app.state.TOOLS`` / ``FUNCTIONS``. Saving a tool on +# worker A does not touch worker B's module dict, so worker B keeps +# serving the stale module until the process restarts. +# +# We fix that by publishing an invalidation message on Redis whenever +# a plugin is created / updated / deleted / toggled. Every worker +# subscribes on startup and drops the matching entries from its local +# caches, forcing the next invocation to reload fresh from the DB. +# +# When Redis is not configured we fall back to in-process invalidation +# only — the single-worker case "just works" because save already +# updates that worker's cache directly. + +import asyncio +import json +import logging +from typing import Literal + +from redis.asyncio import Redis + +from open_webui.env import REDIS_KEY_PREFIX + +log = logging.getLogger(__name__) + +REDIS_PLUGIN_CACHE_CHANNEL = f"{REDIS_KEY_PREFIX}:plugins:invalidate" + +PluginKind = Literal["tool", "function"] + + +def _local_invalidate(app, kind: PluginKind, plugin_id: str) -> None: + """Drop the per-worker cache entries for a plugin id. + + Safe to call when the caches don't exist yet (first invocation on + this worker, startup race with the listener, etc.). + """ + try: + if kind == "tool": + store = getattr(app.state, "TOOLS", None) + content_store = getattr(app.state, "TOOL_CONTENTS", None) + else: + store = getattr(app.state, "FUNCTIONS", None) + content_store = getattr(app.state, "FUNCTION_CONTENTS", None) + if store is not None: + store.pop(plugin_id, None) + if content_store is not None: + content_store.pop(plugin_id, None) + except Exception as e: + log.exception(f"plugin-cache: local invalidate failed ({kind}/{plugin_id}): {e}") + + +async def publish_invalidation(app, kind: PluginKind, plugin_id: str) -> None: + """Invalidate locally + publish to Redis for the other workers. + + Called by the save / update / delete / toggle handlers. The local + invalidation covers the single-worker case and the worker that + received the write request; the Redis publish covers every other + worker in a multi-process deployment. + """ + _local_invalidate(app, kind, plugin_id) + redis: Redis | None = getattr(app.state, "redis", None) + if redis is None: + return + try: + payload = json.dumps({"kind": kind, "id": plugin_id}) + await redis.publish(REDIS_PLUGIN_CACHE_CHANNEL, payload) + except Exception as e: + log.exception(f"plugin-cache: redis publish failed ({kind}/{plugin_id}): {e}") + + +async def plugin_cache_listener(app) -> None: + """Subscribe to invalidations and drop local caches on receipt. + + Mirrors the existing ``redis_task_command_listener`` pattern. + Started from the lifespan hook in main.py. Silently exits if + Redis is unavailable. + """ + redis: Redis | None = getattr(app.state, "redis", None) + if redis is None: + return + pubsub = redis.pubsub() + try: + await pubsub.subscribe(REDIS_PLUGIN_CACHE_CHANNEL) + except Exception as e: + log.exception(f"plugin-cache: redis subscribe failed: {e}") + return + + async for message in pubsub.listen(): + if message.get("type") != "message": + continue + try: + data = message.get("data") + if isinstance(data, (bytes, bytearray)): + data = data.decode("utf-8", "replace") + payload = json.loads(data) + kind = payload.get("kind") + plugin_id = payload.get("id") + if kind not in ("tool", "function") or not plugin_id: + continue + _local_invalidate(app, kind, plugin_id) + log.info(f"plugin-cache: invalidated {kind}/{plugin_id}") + except asyncio.CancelledError: + raise + except Exception as e: + log.exception(f"plugin-cache: listener handler failed: {e}") diff --git a/backend/open_webui/utils/tools.py b/backend/open_webui/utils/tools.py index 3f4eac7e91..8b8c4d9897 100644 --- a/backend/open_webui/utils/tools.py +++ b/backend/open_webui/utils/tools.py @@ -40,7 +40,10 @@ from open_webui.models.tools import Tools from open_webui.models.users import UserModel from open_webui.models.groups import Groups from open_webui.models.access_grants import AccessGrants -from open_webui.utils.plugin import load_tool_module_by_id +from open_webui.utils.plugin import ( + load_tool_module_by_id, + get_tool_module_from_cache, +) from open_webui.utils.access_control import has_access, has_connection_access from open_webui.config import BYPASS_ADMIN_ACCESS_CONTROL from open_webui.env import ( @@ -189,10 +192,11 @@ async def get_tools(request: Request, tool_ids: list[str], user: UserModel, extr log.warning(f'Access denied to tool {tool_id} for user {user.id}') continue - module = request.app.state.TOOLS.get(tool_id, None) - if module is None: - module, _ = await load_tool_module_by_id(tool_id) - request.app.state.TOOLS[tool_id] = module + # Content-hash aware cache lookup — picks up edits from other + # workers even before the Redis invalidation arrives. Falls + # back to loading fresh from DB when cached content no longer + # matches DB. + module, _ = await get_tool_module_from_cache(request, tool_id) __user__ = { **extra_params['__user__'],