RomanNum3ral
/
openwebui-claude-code
mirror of https://github.com/tfriedel/openwebui-claude-code.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Fix: inline session_store/session_marker into claude_agent_pipe.py 

OpenWebUI loads Functions as a single in-memory module via exec()
(open_webui/utils/plugin.py: load_function_module_by_id). There is no
mechanism to ship sibling .py files alongside a Function, so
'from session_store import get_store' / 'from session_marker import ...'
fail with ModuleNotFoundError as soon as the function is pasted/imported
into OpenWebUI (see traceback: exec(content, module.__dict__) -> line 24
-> ModuleNotFoundError: No module named 'session_store').

Fix: inline both modules' full content directly into
claude_agent_pipe.py so it is fully self-contained, and remove the now
redundant standalone session_store.py / session_marker.py files (their
logic lives inline now; keeping both would only invite drift).

Verified by:
- python3 -m py_compile claude_agent_pipe.py
- exec()'ing the file the same way OpenWebUI's plugin loader does (with
  claude_agent_sdk stubbed out), confirming no ModuleNotFoundError/NameError
  and that SessionStore/get_store/make_marker/extract_session_id_from_messages
  and the Pipe class all load correctly.
This commit is contained in:
Assistant Patch 2026-08-10 20:45:48 +00:00
parent 3b4fa5e099
commit 86933746aa
3 changed files with 159 additions and 232 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

161
claude_agent_pipe.py
View File

@ -21,8 +21,165 @@ from typing import Any, AsyncGenerator, Callable, Dict, List, Optional, Set
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
from session_store import get_store # ---------------------------------------------------------------------------
from session_marker import extract_session_id_from_messages, make_marker # Inlined from session_store.py / session_marker.py.
#
# OpenWebUI Functions are loaded as a single in-memory module via exec() (see
# open_webui/utils/plugin.py: load_function_module_by_id) — there is no
# mechanism to ship additional .py files alongside a Function, so importing
# sibling modules (`from session_store import ...`) fails with
# ModuleNotFoundError at install time. Both modules are therefore inlined
# here verbatim. Keep the original files in the repo in sync if you edit
# this logic — they are the canonical, independently testable source; this
# inlined copy is what actually ships to OpenWebUI.
# ---------------------------------------------------------------------------
import sqlite3
import threading
_SESSION_STORE_SCHEMA = """
CREATE TABLE IF NOT EXISTS claude_code_sessions (
chat_id TEXT PRIMARY KEY,
session_id TEXT NOT NULL,
updated_at REAL NOT NULL
);
"""
class SessionStore:
"""Thread-safe, file-backed chat_id -> session_id store.
Persists the chat_id -> Claude Code session_id mapping in a small SQLite
database next to the per-chat workdir (WORKDIR_ROOT/.session_store.sqlite3).
This replaces a plain in-process dict, which loses all mappings on every
backend restart and is inconsistent across multiple worker processes.
Deliberately NOT integrated with OpenWebUI's own database: OpenWebUI's
internal SQLAlchemy engine/schema is not a stable, documented plugin API,
and writing into OpenWebUI's own chat/message tables risks racing with
OpenWebUI's own read-modify-write cycle on those rows. A separate SQLite
file has no such overlap and needs no coordination with OpenWebUI.
"""
def __init__(self, root_dir, filename: str = ".session_store.sqlite3") -> None:
self._path = Path(root_dir) / filename
self._path.parent.mkdir(parents=True, exist_ok=True)
self._lock = threading.Lock()
self._conn = sqlite3.connect(str(self._path), check_same_thread=False, timeout=10)
self._conn.execute("PRAGMA journal_mode=WAL;")
self._conn.execute("PRAGMA busy_timeout=5000;")
with self._lock:
self._conn.execute(_SESSION_STORE_SCHEMA)
self._conn.commit()
def get(self, chat_id: str) -> Optional[str]:
try:
with self._lock:
cur = self._conn.execute(
"SELECT session_id FROM claude_code_sessions WHERE chat_id = ?",
(chat_id,),
)
row = cur.fetchone()
return row[0] if row else None
except sqlite3.Error:
log.exception("SessionStore.get failed for chat_id=%s", chat_id)
return None
def set(self, chat_id: str, session_id: str) -> None:
try:
with self._lock:
self._conn.execute(
"""
INSERT INTO claude_code_sessions (chat_id, session_id, updated_at)
VALUES (?, ?, ?)
ON CONFLICT(chat_id) DO UPDATE SET
session_id = excluded.session_id,
updated_at = excluded.updated_at
""",
(chat_id, session_id, time.time()),
)
self._conn.commit()
except sqlite3.Error:
log.exception("SessionStore.set failed for chat_id=%s", chat_id)
def close(self) -> None:
with self._lock:
self._conn.close()
_session_stores: Dict[str, "SessionStore"] = {}
_session_stores_lock = threading.Lock()
def get_store(root_dir) -> "SessionStore":
"""Return a process-wide singleton SessionStore per WORKDIR_ROOT."""
key = str(root_dir)
with _session_stores_lock:
store = _session_stores.get(key)
if store is None:
store = SessionStore(root_dir)
_session_stores[key] = store
return store
_SESSION_MARKER_RE = re.compile(r"\[claude-session:([A-Za-z0-9_-]{1,128})\]:\s*#")
def make_marker(session_id: str) -> str:
"""Invisible in-text marker carrying the session_id, appended to replies.
Fallback for when the SQLite store is unreachable/wiped independently of
OpenWebUI's own chat history (e.g. different volume lifecycle across a
container redeploy). This is a Markdown link reference definition
(`[label]: #`), which CommonMark-compliant renderers register but never
render as visible output. Rides inside the ordinary assistant message
text that OpenWebUI already persists — no extra write path into
OpenWebUI's own data model, hence no extra race-condition surface.
"""
return f"\n\n[claude-session:{session_id}]: #\n"
def extract_session_id_from_messages(messages: List[Dict[str, Any]]) -> Optional[str]:
"""Scan chat history (most recent first) for the last embedded marker.
Only works if OpenWebUI resends prior assistant messages in
body["messages"] (the normal case for Chat-Completions-shaped pipes).
Fallback only, not a replacement for the SQLite store.
"""
for message in reversed(messages or []):
if message.get("role") != "assistant":
continue
content = message.get("content")
text = _flatten_marker_content(content)
if not text:
continue
match = _SESSION_MARKER_RE.search(text)
if match:
return match.group(1)
return None
def strip_marker(text: str) -> str:
"""Remove marker lines from text before showing/logging it elsewhere."""
return _SESSION_MARKER_RE.sub("", text)
def _flatten_marker_content(content: Any) -> str:
if isinstance(content, str):
return content
if isinstance(content, list):
parts = [
part.get("text", "")
for part in content
if isinstance(part, dict) and part.get("type") == "text"
]
return "\n".join(parts)
return ""
# ---------------------------------------------------------------------------
# End of inlined session_store.py / session_marker.py
# ---------------------------------------------------------------------------
_IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".gif", ".svg", ".webp"} _IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".gif", ".svg", ".webp"}
_DOWNLOAD_EXTENSIONS = { _DOWNLOAD_EXTENSIONS = {

96
session_marker.py
View File

@ -1,96 +0,0 @@
"""
Invisible in-text marker carrying the Claude Code session_id as a fallback
to the SQLite-backed SessionStore (see session_store.py).
Rationale
---------
The SQLite store (Variant A) is the primary source of truth: it survives
process restarts and worker changes without depending on what OpenWebUI
sends back in the request body. But it has one failure mode worth guarding
against: if WORKDIR_ROOT points at ephemeral/container-local storage that
gets wiped independently of OpenWebUI's own chat database (e.g. different
persistent-volume lifecycle in a container redeploy), the SQLite file can
disappear while OpenWebUI's chat history — which is persisted separately and
NOT touched by this pipe — still exists and still contains the last
assistant reply.
To cover that case, every assistant reply gets an invisible marker appended:
[claude-session:<session_id>]: #
This is a Markdown *link reference definition* (CommonMark spec, "reference
link" section): a line of the form `[label]: destination` that is never
rendered as visible output by CommonMark-compliant renderers — it just
registers a reference. Using `#` as the destination keeps it inert (no
actual link target semantics matter here; we only care about it being
absent from rendered output). This is the same technique already validated
in production by rbb-dev/Open-WebUI-OpenRouter-pipe for larger artifact
references.
Important: this marker rides inside the ordinary assistant message text
that OpenWebUI already persists as part of the normal chat flow. Nothing is
written into OpenWebUI's own chat/meta database columns from here — there is
no additional write path into OpenWebUI's data model at all, so there is no
extra conflict/race-condition surface beyond what streaming a normal
response already has.
Limitations (see also the write-up in chat): this only works if OpenWebUI
resends the previous assistant message inside `body["messages"]` on the next
turn (the normal case for OpenAI-Chat-Completions-shaped pipes). It is a
fallback, not a replacement, for the SQLite store.
"""
from __future__ import annotations
import re
from typing import Any, Dict, List, Optional
_MARKER_RE = re.compile(r"\[claude-session:([A-Za-z0-9_-]{1,128})\]:\s*#")
def make_marker(session_id: str) -> str:
"""Return the invisible marker text to append to a streamed reply."""
return f"\n\n[claude-session:{session_id}]: #\n"
def extract_session_id_from_messages(messages: List[Dict[str, Any]]) -> Optional[str]:
"""Scan chat history (most recent first) for the last embedded marker.
Looks only at assistant messages, newest to oldest, and returns the
session_id from the first marker found. Returns None if no marker is
present (e.g. first turn in a chat, or history was trimmed/edited).
"""
for message in reversed(messages or []):
if message.get("role") != "assistant":
continue
content = message.get("content")
text = _flatten_content(content)
if not text:
continue
match = _MARKER_RE.search(text)
if match:
return match.group(1)
return None
def strip_marker(text: str) -> str:
"""Remove marker lines from text before showing/logging it elsewhere.
Not needed for normal rendering (CommonMark already hides it), but
useful if the raw text is re-used somewhere that doesn't apply Markdown
rendering (e.g. plain-text export, logs).
"""
return _MARKER_RE.sub("", text)
def _flatten_content(content: Any) -> str:
if isinstance(content, str):
return content
if isinstance(content, list):
parts = [
part.get("text", "")
for part in content
if isinstance(part, dict) and part.get("type") == "text"
]
return "\n".join(parts)
return ""

134
session_store.py
View File

@ -1,134 +0,0 @@
"""
Persistent chat_id -> Claude Code session_id mapping.
Why this exists
----------------
`claude_agent_pipe.py` used to keep this mapping in a plain in-process dict
(`_chat_sessions`). That breaks in three common situations:
1. The backend process restarts (redeploy, crash, admin reloads the
function) -> the dict is empty again.
2. The backend runs with more than one worker process / replica -> each
worker has its own dict, so whichever worker handles the next turn may
simply not know about the session the previous turn created.
3. Long-lived deployments accumulate sessions for chats that are no longer
active with no way to expire them.
This module replaces the in-memory dict with a small SQLite database that
lives next to the per-chat workdir (`WORKDIR_ROOT/.session_store.sqlite3`).
It is intentionally NOT integrated with OpenWebUI's own database: OpenWebUI's
internal SQLAlchemy engine/schema is not a stable, documented plugin API, and
writing into OpenWebUI's own `chat`/`message` tables risks racing with
OpenWebUI's own read-modify-write cycle on the same rows. A separate SQLite
file has no such overlap and needs no coordination with OpenWebUI at all.
SQLite is a reasonable choice here (rather than e.g. requiring Redis) because
the write volume is tiny (one row write per chat turn) and SQLite's built-in
locking is sufficient for the "one active turn per chat_id at a time" access
pattern this pipe has. If you run many worker processes hammering the *same*
chat_id concurrently, SQLite's default locking will simply serialize those
writes rather than corrupt anything.
Usage
-----
store = SessionStore(root_dir)
resume_id = store.get(chat_id)
...
store.set(chat_id, session_id)
"""
from __future__ import annotations
import logging
import sqlite3
import threading
import time
from pathlib import Path
from typing import Optional
log = logging.getLogger(__name__)
_SCHEMA = """
CREATE TABLE IF NOT EXISTS claude_code_sessions (
chat_id TEXT PRIMARY KEY,
session_id TEXT NOT NULL,
updated_at REAL NOT NULL
);
"""
class SessionStore:
"""Thread-safe, file-backed chat_id -> session_id store.
One SQLite connection per instance, guarded by a lock. Pipes are invoked
concurrently for different chats within the same process, so the lock is
held only for the duration of a single get/set (a few milliseconds),
not for the whole agent turn.
"""
def __init__(self, root_dir: str | Path, filename: str = ".session_store.sqlite3") -> None:
self._path = Path(root_dir) / filename
self._path.parent.mkdir(parents=True, exist_ok=True)
self._lock = threading.Lock()
# check_same_thread=False: pipe() runs inside asyncio, potentially
# scheduled on different threads across calls (e.g. via run_in_executor
# elsewhere in the codebase). Access is still serialized by self._lock.
self._conn = sqlite3.connect(str(self._path), check_same_thread=False, timeout=10)
self._conn.execute("PRAGMA journal_mode=WAL;")
self._conn.execute("PRAGMA busy_timeout=5000;")
with self._lock:
self._conn.execute(_SCHEMA)
self._conn.commit()
def get(self, chat_id: str) -> Optional[str]:
try:
with self._lock:
cur = self._conn.execute(
"SELECT session_id FROM claude_code_sessions WHERE chat_id = ?",
(chat_id,),
)
row = cur.fetchone()
return row[0] if row else None
except sqlite3.Error:
log.exception("SessionStore.get failed for chat_id=%s", chat_id)
return None
def set(self, chat_id: str, session_id: str) -> None:
try:
with self._lock:
self._conn.execute(
"""
INSERT INTO claude_code_sessions (chat_id, session_id, updated_at)
VALUES (?, ?, ?)
ON CONFLICT(chat_id) DO UPDATE SET
session_id = excluded.session_id,
updated_at = excluded.updated_at
""",
(chat_id, session_id, time.time()),
)
self._conn.commit()
except sqlite3.Error:
log.exception("SessionStore.set failed for chat_id=%s", chat_id)
def close(self) -> None:
with self._lock:
self._conn.close()
_stores: dict[str, SessionStore] = {}
_stores_lock = threading.Lock()
def get_store(root_dir: str | Path) -> SessionStore:
"""Return a process-wide singleton SessionStore per WORKDIR_ROOT.
Avoids opening a new SQLite connection on every single pipe() call while
still keying correctly if the valve ever points at a different root.
"""
key = str(root_dir)
with _stores_lock:
store = _stores.get(key)
if store is None:
store = SessionStore(root_dir)
_stores[key] = store
return store