diff --git a/claude_agent_pipe.py b/claude_agent_pipe.py index 34ab0ec..0fd949e 100644 --- a/claude_agent_pipe.py +++ b/claude_agent_pipe.py @@ -21,8 +21,165 @@ from typing import Any, AsyncGenerator, Callable, Dict, List, Optional, Set 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"} _DOWNLOAD_EXTENSIONS = { diff --git a/session_marker.py b/session_marker.py deleted file mode 100644 index 714c435..0000000 --- a/session_marker.py +++ /dev/null @@ -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:]: # - -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 "" diff --git a/session_store.py b/session_store.py deleted file mode 100644 index 32de925..0000000 --- a/session_store.py +++ /dev/null @@ -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