Coverage for src/lilbee/sessions/store.py: 100%
268 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""Append-only JSONL store for chat sessions.
3Each session is one ``<id>.jsonl`` file under ``cfg.data_dir/sessions``. The file
4is a strictly append-only event log: one JSON object per line, appended and
5fsynced, never rewritten. Event types are ``meta`` (first line), ``title``
6(newest wins, so rename appends rather than rewrites), ``message``, and
7``summary`` (newest wins; compaction's condensed view of the turns that no
8longer fit the prompt). The only corruption an append log can suffer is a torn
9final line, which the reader skips. A fork's file is written whole once, then
10only appended to.
11"""
13from __future__ import annotations
15import json
16import os
17from collections.abc import Callable, Iterator
18from dataclasses import dataclass
19from datetime import UTC, datetime
20from enum import StrEnum
21from pathlib import Path
22from typing import Any
23from uuid import uuid4
25from filelock import FileLock
27from lilbee.core.config import cfg
28from lilbee.core.security import (
29 OWNER_ONLY_MODE,
30 ensure_private_dir,
31 private_opener,
32 write_private_text,
33)
35SESSIONS_DIRNAME = "sessions"
36SESSIONS_DISABLED_HINT = (
37 "Sessions are off. Turn them on with /set sessions_enabled true in the TUI, "
38 "settings_set via MCP, or sessions_enabled = true in config.toml."
39)
40AGENT_SESSIONS_DISABLED_HINT = (
41 "Agent sessions are off. Turn them on with settings_set mcp_sessions_enabled "
42 "true, /set mcp_sessions_enabled true in the TUI, or mcp_sessions_enabled = "
43 "true in config.toml."
44)
45# Bounds a wedged lock holder; a healthy append holds the lock for milliseconds.
46_APPEND_LOCK_TIMEOUT_S = 10
47UNTITLED_SESSION_TITLE = "Untitled chat"
48TITLE_MAX_LEN = 60
49TITLE_ELLIPSIS = "…"
50FORK_TITLE_SUFFIX = " (fork {number})"
53def sessions_enabled() -> bool:
54 """True when session persistence is on for the human surfaces (default on)."""
55 return cfg.sessions_enabled
58def agent_sessions_enabled() -> bool:
59 """True when agent (MCP) sessions are on (default off).
61 Independent of ``sessions_enabled``: the two govern separate domains, so an
62 agent's working state can be off while a human's conversations are saved.
63 """
64 return cfg.mcp_sessions_enabled
67class SessionEventType(StrEnum):
68 """The tag on every line of a session log."""
70 META = "meta"
71 TITLE = "title"
72 MESSAGE = "message"
73 SUMMARY = "summary"
74 ORIGIN = "origin"
77class SessionOrigin(StrEnum):
78 """The surface a session belongs to: whoever created it, or was last
79 explicitly transferred to. Appends from any other surface are refused, so
80 an agent cannot splice its turns into a conversation a human owns."""
82 TUI = "tui"
83 MCP = "mcp"
84 HTTP = "http"
85 CLI = "cli"
88# The surfaces a human drives directly. Their sessions are one conversation
89# space (start in Obsidian, continue in the TUI); agent sessions are working
90# state and stay out of it unless asked for.
91HUMAN_ORIGINS: frozenset[SessionOrigin] = frozenset(
92 {SessionOrigin.TUI, SessionOrigin.HTTP, SessionOrigin.CLI}
93)
96class MessageRole(StrEnum):
97 """Author of a chat message."""
99 USER = "user"
100 ASSISTANT = "assistant"
103class TitleSource(StrEnum):
104 """Where a session title came from."""
106 AUTO = "auto"
107 CUSTOM = "custom"
110@dataclass(frozen=True)
111class StoredSourceRef:
112 """One entry of a turn's Sources list: the source, the ``[n]`` the answer cites
113 it by (None when the number is unknown) and the plain-text locator of its page
114 or line span ('' when it has none). The label is not stored: a reader computes
115 it from the current stored title."""
117 source: str
118 citation_number: int | None = None
119 locator: str = ""
122@dataclass(frozen=True)
123class SessionMessage:
124 """One turn in a session. ``ts`` is stamped by the store on write.
126 ``source_refs`` is None for a turn saved before the store kept it; such a turn
127 has only ``sources``, or a Sources list at the end of ``content``.
128 """
130 role: MessageRole
131 content: str
132 sources: tuple[str, ...] = ()
133 ts: str = ""
134 source_refs: tuple[StoredSourceRef, ...] | None = None
137def message_with_refs(
138 role: MessageRole, content: str, refs: tuple[StoredSourceRef, ...]
139) -> SessionMessage:
140 """A turn with the Sources list *refs*, its ``sources`` kept for older readers."""
141 return SessionMessage(
142 role=role, content=content, sources=tuple(r.source for r in refs), source_refs=refs
143 )
146@dataclass(frozen=True)
147class SessionMeta:
148 """Session metadata, reconstructed from the log without its message bodies."""
150 id: str
151 title: str
152 created_at: str
153 updated_at: str
154 model_ref: str
155 scope: str
156 message_count: int
157 origin: SessionOrigin = SessionOrigin.TUI
158 """Owning surface. Files written before ownership existed carry no origin;
159 the only writer then was the TUI, so that is the fallback."""
160 forked_from: str = ""
161 """Id of the session this one was forked from; empty when it is not a fork."""
164@dataclass(frozen=True)
165class Session:
166 """A session's metadata plus its full transcript."""
168 meta: SessionMeta
169 messages: tuple[SessionMessage, ...]
170 # Rolling summary of the turns compaction has folded away, empty until the
171 # conversation first outgrows the prompt budget. It lives here rather than on
172 # the meta because only replaying a session needs it: listing does not, and
173 # carrying a paragraph per session would bloat the drawer's hot path and
174 # every HTTP/MCP list payload.
175 summary: str = ""
178class SessionNotFoundError(Exception):
179 """Raised when a session id has no backing file."""
181 def __init__(self, session_id: str) -> None:
182 super().__init__(f"No session with id {session_id!r}")
183 self.session_id = session_id
186def _may_append(surface: SessionOrigin, owner: SessionOrigin) -> bool:
187 """Whether *surface* may append to a session owned by *owner*.
189 The human surfaces are one conversation space (the same person in the TUI,
190 Obsidian, or the shell), so they append to each other's sessions freely.
191 Agent sessions are working state: only the agent surface appends to them,
192 and it appends to nothing else without an explicit claim.
193 """
194 if surface is owner:
195 return True
196 return surface in HUMAN_ORIGINS and owner in HUMAN_ORIGINS
199class SessionOwnershipError(Exception):
200 """Raised when a surface appends to a session another surface owns."""
202 def __init__(self, session_id: str, owner: SessionOrigin, surface: SessionOrigin) -> None:
203 super().__init__(
204 f"Session {session_id!r} belongs to the {owner.value} surface; "
205 f"claim it before appending from {surface.value}."
206 )
207 self.session_id = session_id
208 self.owner = owner
209 self.surface = surface
212class SessionForkRangeError(ValueError):
213 """Raised when a fork point lies outside the source session's transcript."""
215 def __init__(self, session_id: str, message_count: int, available: int) -> None:
216 super().__init__(
217 f"Cannot fork session {session_id!r} after {message_count} messages: "
218 f"choose a count from 0 to {available}."
219 )
220 self.session_id = session_id
221 self.message_count = message_count
222 self.available = available
225def derive_title(text: str) -> str:
226 """Title a session from its first user message: first line, truncated."""
227 stripped = text.strip()
228 if not stripped:
229 return UNTITLED_SESSION_TITLE
230 first = stripped.splitlines()[0]
231 if len(first) > TITLE_MAX_LEN:
232 return first[:TITLE_MAX_LEN] + TITLE_ELLIPSIS
233 return first
236def _fork_title(source_title: str, number: int) -> str:
237 """``<source title> (fork N)``, the source part clipped so the whole fits TITLE_MAX_LEN."""
238 suffix = FORK_TITLE_SUFFIX.format(number=number)
239 room = TITLE_MAX_LEN - len(suffix)
240 if len(source_title) <= room:
241 return source_title + suffix
242 return source_title[: room - len(TITLE_ELLIPSIS)] + TITLE_ELLIPSIS + suffix
245def _fork_count(source: Session, message_count: int | None) -> int:
246 """The number of leading messages to copy, checked against *source*'s transcript."""
247 available = len(source.messages)
248 if message_count is None:
249 return available
250 if not 0 <= message_count <= available:
251 raise SessionForkRangeError(source.meta.id, message_count, available)
252 return message_count
255def _message_event(message: SessionMessage, ts: str) -> dict[str, Any]:
256 """The ``message`` event line for *message*, stamped *ts*."""
257 event: dict[str, Any] = {
258 "type": SessionEventType.MESSAGE,
259 "role": message.role,
260 "content": message.content,
261 "sources": list(message.sources),
262 "ts": ts,
263 }
264 if message.source_refs is not None:
265 event["source_refs"] = [_ref_event(r) for r in message.source_refs]
266 return event
269def _ref_event(ref: StoredSourceRef) -> dict[str, Any]:
270 """*ref* as its entry in a ``message`` event, its locator only when it has one."""
271 entry: dict[str, Any] = {"source": ref.source, "citation_number": ref.citation_number}
272 if ref.locator:
273 entry["locator"] = ref.locator
274 return entry
277def _message_from_event(event: dict[str, Any], ts: str) -> SessionMessage:
278 """Reconstruct one message from its ``message`` event line."""
279 refs = event.get("source_refs")
280 return SessionMessage(
281 role=MessageRole(event["role"]),
282 content=event["content"],
283 sources=tuple(event.get("sources", [])),
284 ts=ts,
285 source_refs=None
286 if refs is None
287 else tuple(
288 StoredSourceRef(r["source"], r.get("citation_number"), r.get("locator", ""))
289 for r in refs
290 ),
291 )
294class SessionStore:
295 """Reads and appends session logs under ``cfg.data_dir/sessions``.
297 The directory is resolved late-bound from ``cfg`` on every call, so the store
298 follows a reconfigured data dir (and test isolation) without reconstruction.
299 ``clock`` is injectable for deterministic tests.
300 """
302 def __init__(self, clock: Callable[[], datetime] | None = None) -> None:
303 self._clock = clock or (lambda: datetime.now(UTC))
304 # path -> (size, mtime, meta) from the last fold of that file; see _meta_for.
305 self._meta_cache: dict[Path, tuple[int, float, SessionMeta]] = {}
307 @property
308 def _dir(self) -> Path:
309 return cfg.data_dir / SESSIONS_DIRNAME
311 def _path(self, session_id: str) -> Path:
312 return self._dir / f"{session_id}.jsonl"
314 def _now(self) -> str:
315 return self._clock().isoformat()
317 def _require(self, session_id: str) -> Path:
318 path = self._path(session_id)
319 if not path.exists():
320 raise SessionNotFoundError(session_id)
321 return path
323 @staticmethod
324 def _write_event(path: Path, event: dict[str, Any]) -> None:
325 # Per-session lock: two writers on one id (a second process, another
326 # surface) serialize instead of interleaving lines. Appends take
327 # milliseconds, so a blocked writer waits, never fails, under any
328 # realistic contention; the timeout only bounds a wedged holder.
329 with (
330 FileLock(str(path) + ".lock", timeout=_APPEND_LOCK_TIMEOUT_S, mode=OWNER_ONLY_MODE),
331 open(path, "a", encoding="utf-8", opener=private_opener) as fh,
332 ):
333 fh.write(json.dumps(event) + "\n")
334 fh.flush()
335 os.fsync(fh.fileno())
337 @staticmethod
338 def _iter_events(path: Path) -> Iterator[dict[str, Any]]:
339 # A torn multi-byte character decodes here, outside the try that skips
340 # torn lines; replacing makes it a JSON failure that try can catch.
341 with path.open(encoding="utf-8", errors="replace") as fh:
342 for raw in fh:
343 line = raw.strip()
344 if not line:
345 continue
346 try:
347 yield json.loads(line)
348 except json.JSONDecodeError:
349 continue # torn final line; skip it
351 @staticmethod
352 def _meta_event(
353 session_id: str,
354 model_ref: str,
355 scope: str,
356 origin: SessionOrigin,
357 now: str,
358 forked_from: str,
359 ) -> dict[str, Any]:
360 return {
361 "type": SessionEventType.META,
362 "id": session_id,
363 "created_at": now,
364 "model_ref": model_ref,
365 "scope": scope,
366 "origin": origin,
367 "forked_from": forked_from,
368 "ts": now,
369 }
371 def create(self, model_ref: str, scope: str, origin: SessionOrigin = SessionOrigin.TUI) -> str:
372 """Start a new session owned by *origin* and return its id."""
373 session_id = uuid4().hex
374 ensure_private_dir(self._dir)
375 meta = self._meta_event(session_id, model_ref, scope, origin, self._now(), forked_from="")
376 self._write_event(self._path(session_id), meta)
377 return session_id
379 def fork(
380 self,
381 session_id: str,
382 *,
383 message_count: int | None = None,
384 origin: SessionOrigin = SessionOrigin.TUI,
385 ) -> str:
386 """Copy the first *message_count* messages (all when None) into a new session.
388 The fork is owned by *origin*, which must be allowed to append to the
389 source. The source is read once and never written.
390 """
391 source = self.get(session_id)
392 if not _may_append(origin, source.meta.origin):
393 raise SessionOwnershipError(session_id, source.meta.origin, origin)
394 count = _fork_count(source, message_count)
395 fork_id = uuid4().hex
396 events = self._fork_events(fork_id, source, count, origin)
397 write_private_text(self._path(fork_id), "".join(json.dumps(e) + "\n" for e in events))
398 return fork_id
400 def _fork_events(
401 self, fork_id: str, source: Session, count: int, origin: SessionOrigin
402 ) -> list[dict[str, Any]]:
403 """A fork's whole log; the title event is last so ``updated_at`` is the fork time."""
404 now = self._now()
405 meta = source.meta
406 events = [self._meta_event(fork_id, meta.model_ref, meta.scope, origin, now, meta.id)]
407 events += [_message_event(message, message.ts) for message in source.messages[:count]]
408 if count == len(source.messages) and source.summary:
409 events.append({"type": SessionEventType.SUMMARY, "summary": source.summary, "ts": now})
410 title = _fork_title(meta.title, self._fork_number(meta.id))
411 events.append(
412 {"type": SessionEventType.TITLE, "title": title, "source": TitleSource.AUTO, "ts": now}
413 )
414 return events
416 def _fork_number(self, source_id: str) -> int:
417 """1 plus the number of existing forks of *source_id*."""
418 return 1 + sum(1 for meta in self.list() if meta.forked_from == source_id)
420 def add_message(
421 self, session_id: str, message: SessionMessage, *, surface: SessionOrigin | None = None
422 ) -> None:
423 """Append one message event to an existing session.
425 With *surface* given, the append is refused unless that surface owns the
426 session (see ``transfer``); without it the caller is a library embedder
427 that manages its own store and ownership does not apply.
428 """
429 path = self._require(session_id)
430 if surface is not None:
431 meta = self._meta_for(path)
432 if meta is not None and not _may_append(surface, meta.origin):
433 raise SessionOwnershipError(session_id, meta.origin, surface)
434 self._write_event(path, _message_event(message, self._now()))
436 def transfer(self, session_id: str, origin: SessionOrigin) -> None:
437 """Append an origin event handing the session to *origin*; newest wins.
439 This is the explicit bridge between the human and agent domains: an
440 agent claims a session whose id the user handed it, and POST /claim
441 brings one back. Never implicit in an append.
442 """
443 self._write_event(
444 self._require(session_id),
445 {"type": SessionEventType.ORIGIN, "origin": origin, "ts": self._now()},
446 )
448 def set_title(self, session_id: str, title: str, source: TitleSource) -> None:
449 """Append a title event; the newest title wins on read."""
450 self._write_event(
451 self._require(session_id),
452 {"type": SessionEventType.TITLE, "title": title, "source": source, "ts": self._now()},
453 )
455 def set_summary(self, session_id: str, summary: str) -> None:
456 """Append a summary event; the newest summary wins on read.
458 Compaction folds the oldest turns into a summary once they no longer fit
459 the prompt. The messages themselves stay in the log untouched: the
460 transcript the user scrolls is always complete, and only what is fed to
461 the model is condensed.
462 """
463 self._write_event(
464 self._require(session_id),
465 {"type": SessionEventType.SUMMARY, "summary": summary, "ts": self._now()},
466 )
468 def delete(self, session_id: str) -> None:
469 """Remove a session's file."""
470 self._require(session_id).unlink()
472 def get(self, session_id: str) -> Session:
473 """Replay a session's log into its reconstructed view."""
474 meta, messages, summary = self._replay(
475 session_id, self._require(session_id), collect_messages=True
476 )
477 return Session(meta=meta, messages=messages, summary=summary)
479 def list(self, origins: frozenset[SessionOrigin] | None = None) -> list[SessionMeta]:
480 """Sessions' metadata, newest first; *origins* narrows to those surfaces.
482 Listing replays every event of every session, so it is the one hot path
483 here: the drawer runs it on open. Messages are not materialised (only
484 counted), and each file's meta is memoised against its size and mtime so
485 reopening a vault that has not changed costs one stat() per session.
486 """
487 if not self._dir.exists():
488 return []
489 ensure_private_dir(self._dir)
490 paths = list(self._dir.glob("*.jsonl"))
491 metas = [meta for meta in (self._meta_for(path) for path in paths) if meta is not None]
492 if origins is not None:
493 metas = [meta for meta in metas if meta.origin in origins]
494 # Drop cache entries for sessions that no longer exist, so a long-lived
495 # store does not pin the meta of every session ever deleted.
496 live = {path for path in paths}
497 self._meta_cache = {p: v for p, v in self._meta_cache.items() if p in live}
498 return sorted(metas, key=lambda meta: (meta.updated_at, meta.id), reverse=True)
500 def _meta_for(self, path: Path) -> SessionMeta | None:
501 """Meta for one session file, reusing the last fold when it is unchanged.
503 The log is append-only, so any new event grows the file: size plus mtime
504 is enough to notice a change. A file that grows between the stat and the
505 read is simply re-folded on the next list(), never served stale.
507 Returns None when the file goes away underneath us, which is routine: the
508 CLI or another surface can delete a session while the drawer is listing.
509 Reading it instead would raise straight out of list().
510 """
511 try:
512 stat = path.stat()
513 except OSError:
514 return None
515 cached = self._meta_cache.get(path)
516 if cached is not None and cached[0] == stat.st_size and cached[1] == stat.st_mtime:
517 return cached[2]
518 try:
519 meta = self._replay(path.stem, path, collect_messages=False)[0]
520 except OSError:
521 return None
522 self._meta_cache[path] = (stat.st_size, stat.st_mtime, meta)
523 return meta
525 def _replay(
526 self, session_id: str, path: Path, *, collect_messages: bool
527 ) -> tuple[SessionMeta, tuple[SessionMessage, ...], str]:
528 """Fold a session's event log into its meta, messages and summary.
530 ``collect_messages=False`` is for listing, which needs only the count:
531 building a SessionMessage per message across a whole vault is pure waste.
532 """
533 created_at = ""
534 model_ref = ""
535 scope = ""
536 title = UNTITLED_SESSION_TITLE
537 updated_at = ""
538 summary = ""
539 origin = SessionOrigin.TUI
540 forked_from = ""
541 message_count = 0
542 messages: list[SessionMessage] = []
543 for event in self._iter_events(path):
544 ts = event.get("ts", "")
545 updated_at = ts
546 event_type = event.get("type")
547 if event_type == SessionEventType.META:
548 created_at = event["created_at"]
549 model_ref = event["model_ref"]
550 scope = event["scope"]
551 origin = SessionOrigin(event.get("origin", SessionOrigin.TUI))
552 forked_from = event.get("forked_from", "")
553 elif event_type == SessionEventType.ORIGIN:
554 origin = SessionOrigin(event["origin"])
555 elif event_type == SessionEventType.TITLE:
556 title = event["title"]
557 elif event_type == SessionEventType.SUMMARY:
558 summary = event["summary"]
559 elif event_type == SessionEventType.MESSAGE:
560 message_count += 1
561 if collect_messages:
562 messages.append(_message_from_event(event, ts))
563 meta = SessionMeta(
564 id=session_id,
565 title=title,
566 created_at=created_at,
567 updated_at=updated_at,
568 model_ref=model_ref,
569 scope=scope,
570 message_count=message_count,
571 origin=origin,
572 forked_from=forked_from,
573 )
574 return meta, tuple(messages), summary