diff --git a/claude_agent_pipe.py b/claude_agent_pipe.py index 426466c..34ab0ec 100644 --- a/claude_agent_pipe.py +++ b/claude_agent_pipe.py @@ -21,6 +21,9 @@ 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 + _IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".gif", ".svg", ".webp"} _DOWNLOAD_EXTENSIONS = { ".pdf", @@ -61,8 +64,15 @@ from claude_agent_sdk import ( log = logging.getLogger(__name__) -# OpenWebUI calls pipe() fresh for each chat turn. We keep a chat_id -> session_id -# map in-process so follow-up turns resume the same Claude Code session. +# OpenWebUI calls pipe() fresh for each chat turn. This in-process dict is a +# fast path so a hot worker doesn't have to hit SQLite on every turn. It is +# NOT the source of truth anymore: it starts empty on every process restart +# and is per-worker (stale/absent on any other worker), which is exactly why +# Claude Code used to "forget" chat history across restarts / multi-worker +# deployments. The source of truth is session_store.SessionStore (SQLite, +# survives restarts, shared across workers via the shared filesystem), with +# session_marker as a secondary fallback embedded in the chat text itself. +# See session_store.py / session_marker.py for the full rationale. _chat_sessions: Dict[str, str] = {} @@ -1495,13 +1505,49 @@ class Pipe: prompt = _strip_mode_prefix(prompt) chat_id = __chat_id__ or "default" + if not __chat_id__: + log.warning( + "pipe() called without __chat_id__ — falling back to a shared " + "'default' workdir/session. Session resume across turns will " + "not work correctly for this call (see OpenWebUI issue about " + "metadata.chat_id being unset for certain internal calls)." + ) workdir = Path(self.valves.WORKDIR_ROOT) / chat_id workdir.mkdir(parents=True, exist_ok=True) allowed_tools = [ t.strip() for t in self.valves.ALLOWED_TOOLS.split(",") if t.strip() ] + + # Resolve the session to resume, in order of preference: + # 1. In-process cache (_chat_sessions) — fastest, but empty after a + # restart or on any worker that didn't handle the previous turn. + # 2. SQLite-backed SessionStore — survives restarts and is shared + # across worker processes via the shared filesystem under + # WORKDIR_ROOT. This is the source of truth. + # 3. Marker embedded in the previous assistant message, extracted + # from OpenWebUI's own chat history in `body`. Fallback only, + # for the case where the SQLite file itself became unavailable + # (e.g. WORKDIR_ROOT lives on ephemeral storage that was wiped + # independently of OpenWebUI's own chat database). + session_store = get_store(self.valves.WORKDIR_ROOT) resume_id = _chat_sessions.get(chat_id) + resume_source = "memory" if resume_id else None + if not resume_id: + resume_id = session_store.get(chat_id) + if resume_id: + resume_source = "sqlite" + if not resume_id: + resume_id = extract_session_id_from_messages(body.get("messages") or []) + if resume_id: + resume_source = "marker" + if resume_id: + log.debug( + "Resuming Claude Code session %s for chat_id=%s (source=%s)", + resume_id, + chat_id, + resume_source, + ) # Knowledge base attached via Workspace Model → expose as an MCP tool # Claude can call agentically. OpenWebUI's middleware already added one @@ -1621,6 +1667,11 @@ class Pipe: session_id = message.data.get("session_id") if session_id: _chat_sessions[chat_id] = session_id + # Persist beyond this process's lifetime/worker. + # Cheap (single upsert) and done on every init + # message, i.e. once per turn — not hot-path + # sensitive. + session_store.set(chat_id, session_id) continue if isinstance(message, StreamEvent): @@ -1739,6 +1790,16 @@ class Pipe: yield f"\n\n**API-Fehler:** `{message.api_error_status}`\n" if message.total_cost_usd is not None: yield f"\n\n_Cost: ${message.total_cost_usd:.4f} · {message.duration_ms}ms_\n" + # Fallback-path marker (see session_marker.py): embed + # the session_id invisibly in the reply itself so it + # can be recovered from OpenWebUI's own chat history + # even if both the in-memory cache and the SQLite + # store are unavailable on a later turn. This does + # NOT write into any OpenWebUI-owned data — it's part + # of the ordinary streamed assistant text. + current_session_id = _chat_sessions.get(chat_id) + if current_session_id: + yield make_marker(current_session_id) return except Exception as exc: diff --git a/session_marker.py b/session_marker.py new file mode 100644 index 0000000..714c435 --- /dev/null +++ b/session_marker.py @@ -0,0 +1,96 @@ +""" +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 new file mode 100644 index 0000000..32de925 --- /dev/null +++ b/session_store.py @@ -0,0 +1,134 @@ +""" +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