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: persist Claude Code session id across process restarts and workers 

- Add session_store.py: SQLite-backed, WAL-mode session store keyed by
  chat_id, robust to multi-worker access and process restarts.
- Add session_marker.py: fallback mechanism that embeds an invisible
  markdown reference marker in assistant responses so the session id
  survives even without SQLite access, by parsing it back out of
  body messages history.
- Update claude_agent_pipe.py: resolution order is in-memory cache ->
  SQLite store -> marker fallback. Session id is persisted to both the
  in-memory dict and SQLite on every init SystemMessage, and the
  marker is appended to the final visible response text.
- Warn when chat_id is missing/None instead of silently starting
  a fresh session (related to open-webui/open-webui#20563).
This commit is contained in:
Assistant Patch 2026-08-10 13:21:36 +00:00
parent 8f0d99100b
commit 3c2217115b
3 changed files with 293 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

65
claude_agent_pipe.py
View File

@ -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:

96
session_marker.py Normal file
View File

@ -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:<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 Normal file
View File

@ -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