jupyterlab-codex-sidebar 0.1.4 → 0.1.6

package/jupyterlab_codex/sessions.py

@@ -1,15 +1,24 @@
 import json
 import os
 import re
+import threading
 from datetime import datetime, timedelta, timezone
 from pathlib import Path
-from typing import Dict, List, Tuple
+from typing import Any, Dict, List, Tuple
+from uuid import uuid4
 
 
 _TRUE_VALUES = {"1", "true", "y", "yes", "on"}
 _FALSE_VALUES = {"0", "false", "n", "no", "off"}
 _DEFAULT_SESSION_RETENTION_DAYS = 30
+_DEFAULT_SESSION_MAX_MESSAGES = 100
+_DEFAULT_SESSION_MAX_BYTES = 2_000_000
+_DEFAULT_SESSION_PRUNE_INTERVAL_MINUTES = 15
 _DEFAULT_MAX_MESSAGE_CHARS = 12000
+_DEFAULT_UI_LABEL_MAX_CHARS = 80
+_DEFAULT_UI_PREVIEW_MAX_CHARS = 500
+_DEFAULT_UI_PREVIEW_MAX_ITEMS_PER_SESSION = 10
+_SESSION_FILE_VERSION = 1
 
 _SENSITIVE_PATTERNS = [
     (
@@ -25,6 +34,10 @@ _SENSITIVE_PATTERNS = [
 
 
 class SessionStore:
+    """Session persistence with bounded growth and recoverable file handling."""
+
+    _file_lock = threading.RLock()
+
     def __init__(self, base_dir: str | None = None):
         root = base_dir or os.path.join(os.path.expanduser("~"), ".jupyter", "codex-sessions")
         self._base = Path(root)
@@ -32,25 +45,46 @@ class SessionStore:
             os.environ.get("JUPYTERLAB_CODEX_SESSION_LOGGING"), default=True
         )
         self._retention_days = _as_non_negative_int(
-            os.environ.get("JUPYTERLAB_CODEX_SESSION_RETENTION_DAYS"), _DEFAULT_SESSION_RETENTION_DAYS
+            os.environ.get("JUPYTERLAB_CODEX_SESSION_RETENTION_DAYS"),
+            _DEFAULT_SESSION_RETENTION_DAYS,
         )
         self._max_message_chars = _as_positive_int(
             os.environ.get("JUPYTERLAB_CODEX_SESSION_MAX_MESSAGE_CHARS"), _DEFAULT_MAX_MESSAGE_CHARS
         )
+        self._max_messages_per_session = _as_positive_int(
+            os.environ.get("JUPYTERLAB_CODEX_SESSION_MAX_MESSAGES"),
+            _DEFAULT_SESSION_MAX_MESSAGES,
+        )
+        self._max_session_bytes = _as_non_negative_int(
+            os.environ.get("JUPYTERLAB_CODEX_SESSION_MAX_BYTES"),
+            _DEFAULT_SESSION_MAX_BYTES,
+        )
+        self._prune_interval = timedelta(
+            minutes=_as_positive_int(
+                os.environ.get("JUPYTERLAB_CODEX_SESSION_PRUNE_INTERVAL_MINUTES"),
+                _DEFAULT_SESSION_PRUNE_INTERVAL_MINUTES,
+            )
+        )
+        self._last_global_prune = datetime.min.replace(tzinfo=timezone.utc)
+
         if self._logging_enabled:
             self._base.mkdir(parents=True, exist_ok=True)
-            self.prune_expired_sessions()
+            with self._file_lock:
+                self._prune_expired_sessions_locked()
 
     def ensure_session(self, session_id: str, notebook_path: str, notebook_os_path: str = "") -> None:
         if not self._logging_enabled:
             return
+        if not session_id:
+            return
 
-        meta_path = self._meta_path(session_id)
-        if meta_path.exists():
+        existing_meta = self._load_meta(session_id)
+        if existing_meta:
             return
 
         paired_path, paired_os_path = _derive_paired_paths(notebook_path, notebook_os_path)
         meta = {
+            "schema_version": _SESSION_FILE_VERSION,
             "session_id": session_id,
             "notebook_path": notebook_path,
             "notebook_os_path": notebook_os_path,
@@ -58,58 +92,97 @@ class SessionStore:
             "paired_os_path": paired_os_path,
             "created_at": _now_iso(),
             "updated_at": _now_iso(),
+            "retention_days": self._retention_days,
+            "max_messages_per_session": self._max_messages_per_session,
         }
-        meta_path.write_text(json.dumps(meta, indent=2), encoding="utf-8")
+        with self._file_lock:
+            self._write_meta_atomic(self._meta_path(session_id), meta)
 
-    def append_message(self, session_id: str, role: str, content: str) -> None:
+    def append_message(
+        self, session_id: str, role: str, content: str, ui: Dict[str, Any] | None = None
+    ) -> None:
         if not self._logging_enabled:
             return
+        if not session_id:
+            return
 
+        normalized_role = role if role in {"system", "user", "assistant"} else "system"
         record = {
-            "role": role,
+            "role": normalized_role,
             "content": _sanitize_message(content, self._max_message_chars),
             "timestamp": _now_iso(),
         }
-        with self._jsonl_path(session_id).open("a", encoding="utf-8") as handle:
-            handle.write(json.dumps(record))
-            handle.write("\n")
-        self._touch_meta(session_id)
-        self.prune_expired_sessions()
+        ui_payload = _sanitize_ui_payload(ui)
+        if ui_payload:
+            record["ui"] = ui_payload
 
-    def load_messages(self, session_id: str) -> List[Dict[str, str]]:
+        with self._file_lock:
+            path = self._jsonl_path(session_id)
+            try:
+                with path.open("a", encoding="utf-8") as handle:
+                    handle.write(json.dumps(record))
+                    handle.write("\n")
+            except OSError:
+                return
+
+            self._touch_meta_locked(session_id)
+            self._enforce_session_limits_locked(session_id)
+            if self._is_global_prune_due():
+                self._prune_expired_sessions_locked()
+
+    def load_messages(self, session_id: str) -> List[Dict[str, Any]]:
         if not self._logging_enabled:
             return []
+        if not session_id:
+            return []
 
         path = self._jsonl_path(session_id)
         if not path.exists():
             return []
 
-        messages = []
-        with path.open("r", encoding="utf-8") as handle:
-            for line in handle:
-                line = line.strip()
-                if not line:
-                    continue
-                try:
-                    messages.append(json.loads(line))
-                except json.JSONDecodeError:
-                    continue
+        messages: List[Dict[str, Any]] = []
+        with self._file_lock:
+            try:
+                with path.open("r", encoding="utf-8") as handle:
+                    for line in handle:
+                        line = line.strip()
+                        if not line:
+                            continue
+                        try:
+                            payload = json.loads(line)
+                        except json.JSONDecodeError:
+                            continue
+                        if isinstance(payload, dict):
+                            messages.append(payload)
+            except OSError:
+                return []
         return messages
 
+    def _prune_user_ui_previews(self, session_id: str, keep_latest: int) -> None:
+        if keep_latest <= 0:
+            return
+        with self._file_lock:
+            self._trim_session_records_locked(session_id, keep_ui_previews=keep_latest)
+
     def build_prompt(
         self,
         session_id: str,
         user_content: str,
         selection: str,
         cell_output: str,
+        selection_truncated: bool = False,
+        cell_output_truncated: bool = False,
         cwd: str | None = None,
+        notebook_mode: str = "",
+        include_history: bool = True,
     ) -> str:
-        messages = self.load_messages(session_id)
+        messages = self.load_messages(session_id) if include_history else []
         meta = self._load_meta(session_id)
         notebook_path = meta.get("notebook_path", "")
         notebook_os_path = meta.get("notebook_os_path", "")
         paired_path = meta.get("paired_path", "")
         paired_os_path = meta.get("paired_os_path", "")
+        mode = _normalize_notebook_mode(notebook_mode, notebook_path, notebook_os_path)
 
         if not paired_path and not paired_os_path:
             paired_path, paired_os_path = _derive_paired_paths(notebook_path, notebook_os_path)
@@ -126,7 +199,7 @@ class SessionStore:
         if cwd:
             parts.append(f"System: Current working directory: {cwd}")
 
-        if paired_os_path:
+        if mode == "ipynb" and paired_os_path:
             parts.extend(
                 [
                     f"System: Jupytext paired file (absolute path): {paired_os_path}",
@@ -134,7 +207,7 @@ class SessionStore:
                     "System: The notebook will prompt reload when the paired file changes on disk.",
                 ]
             )
-        elif paired_path:
+        elif mode == "ipynb" and paired_path:
             parts.extend(
                 [
                     f"System: Jupytext paired file (Jupyter path): {paired_path}",
@@ -142,20 +215,55 @@ class SessionStore:
                     "System: The notebook will prompt reload when the paired file changes on disk.",
                 ]
             )
+        elif mode == "jupytext_py":
+            target = notebook_os_path or notebook_path or "<notebook>.py"
+            parts.extend(
+                [
+                    "System: Current file mode: Jupytext Python notebook script (.py).",
+                    f"System: IMPORTANT - Edit this file directly: {target}",
+                ]
+            )
+        elif mode == "plain_py":
+            target = notebook_os_path or notebook_path or "<script>.py"
+            parts.extend(
+                [
+                    "System: Current file mode: Plain Python script (.py).",
+                    f"System: IMPORTANT - Edit this file directly: {target}",
+                ]
+            )
 
-        parts.extend(
-            [
-                "",
-                "System: Instructions:",
+        if mode == "ipynb":
+            instructions = [
                 "System: 1. For code changes, modify the paired file directly using file editing tools.",
                 "System: 2. Keep edits minimal and aligned with the user request.",
                 "System: 3. The 'Current Cell Content' shows what the user is currently viewing/editing.",
                 "System: 4. If you cannot proceed due to sandbox/permission restrictions, say so explicitly and ask the user to switch Permission (shield icon) to 'Full access' and retry. If authentication is required, tell them to run `codex login` in a terminal first.",
-                "",
             ]
-        )
+        elif mode == "jupytext_py":
+            instructions = [
+                "System: 1. For code changes, modify the current .py file directly using file editing tools.",
+                "System: 2. Preserve existing Jupytext structure and metadata (YAML header and # %% cell markers) unless the user asks to change them.",
+                "System: 3. The 'Current Cell Content' is a notebook cell snippet from the .py file.",
+                "System: 4. If you cannot proceed due to sandbox/permission restrictions, say so explicitly and ask the user to switch Permission (shield icon) to 'Full access' and retry. If authentication is required, tell them to run `codex login` in a terminal first.",
+            ]
+        elif mode == "plain_py":
+            instructions = [
+                "System: 1. For code changes, modify the current .py file directly using file editing tools.",
+                "System: 2. Do not introduce Jupytext YAML headers or notebook cell markers (for example, # %%) unless the user explicitly requests it.",
+                "System: 3. If no context snippet is provided, inspect files directly before making edits.",
+                "System: 4. If you cannot proceed due to sandbox/permission restrictions, say so explicitly and ask the user to switch Permission (shield icon) to 'Full access' and retry. If authentication is required, tell them to run `codex login` in a terminal first.",
+            ]
+        else:
+            instructions = [
+                "System: 1. For code changes, inspect files directly and edit the correct target file.",
+                "System: 2. Keep edits minimal and aligned with the user request.",
+                "System: 3. The provided context snippet, if any, may be partial.",
+                "System: 4. If you cannot proceed due to sandbox/permission restrictions, say so explicitly and ask the user to switch Permission (shield icon) to 'Full access' and retry. If authentication is required, tell them to run `codex login` in a terminal first.",
+            ]
+
+        parts.extend(["", "System: Instructions:", *instructions, ""])
 
-        if messages:
+        if include_history and messages:
             parts.append("Conversation:")
             for msg in messages:
                 role = msg.get("role", "user")
@@ -163,15 +271,28 @@ class SessionStore:
                 parts.append(f"{role.title()}: {content}")
             parts.append("")
 
-        if selection:
+        include_selection = mode in {"ipynb", "jupytext_py", "plain_py"}
+        include_cell_output = mode in {"ipynb", "jupytext_py"}
+
+        if include_selection and selection:
             parts.append("Current Cell Content:")
             parts.append(selection)
             parts.append("")
+        if include_selection and selection_truncated:
+            parts.append(
+                "System: Current Cell Content was truncated before sending due size limits. If full context is needed, inspect the source file directly."
+            )
+            parts.append("")
 
-        if cell_output:
+        if include_cell_output and cell_output:
             parts.append("Current Cell Output:")
             parts.append(cell_output)
             parts.append("")
+        if include_cell_output and cell_output_truncated:
+            parts.append(
+                "System: Current Cell Output was truncated before sending due size limits."
+            )
+            parts.append("")
 
         parts.append("User:")
         parts.append(user_content)
@@ -182,6 +303,44 @@ class SessionStore:
         meta = self._load_meta(session_id)
         return meta.get("notebook_path", "")
 
+    def has_session(self, session_id: str) -> bool:
+        normalized_session_id = (session_id or "").strip()
+        if not normalized_session_id:
+            return False
+        if not self._logging_enabled:
+            return False
+        meta = self._load_meta(normalized_session_id)
+        return isinstance(meta, dict) and bool(meta)
+
+    def session_matches_notebook(
+        self, session_id: str, notebook_path: str, notebook_os_path: str = ""
+    ) -> bool:
+        """
+        Validate whether an existing session id belongs to the current notebook.
+        """
+        normalized_session_id = (session_id or "").strip()
+        if not normalized_session_id:
+            return False
+        if not self._logging_enabled:
+            return True
+
+        meta = self._load_meta(normalized_session_id)
+        if not isinstance(meta, dict) or not meta:
+            return False
+
+        normalized_notebook_path = (notebook_path or "").strip()
+        normalized_notebook_os_path = (notebook_os_path or "").strip()
+        stored_notebook_path = (meta.get("notebook_path") or "").strip()
+        stored_notebook_os_path = (meta.get("notebook_os_path") or "").strip()
+
+        if normalized_notebook_path and stored_notebook_path == normalized_notebook_path:
+            return True
+        if normalized_notebook_os_path and stored_notebook_os_path == normalized_notebook_os_path:
+            return True
+        if not normalized_notebook_path and not normalized_notebook_os_path:
+            return True
+        return False
+
     def resolve_session_for_notebook(self, notebook_path: str, notebook_os_path: str = "") -> str:
         if not self._logging_enabled:
             return ""
@@ -194,57 +353,61 @@ class SessionStore:
         latest_session_id = ""
         latest_updated_at = None
 
-        for path in self._base.glob("*.meta.json"):
-            session_id = path.stem.removesuffix(".meta")
-            try:
-                meta = json.loads(path.read_text(encoding="utf-8"))
-            except (json.JSONDecodeError, IOError):
-                continue
-            if not isinstance(meta, dict):
-                continue
+        with self._file_lock:
+            for path in self._base.glob("*.meta.json"):
+                session_id = path.stem.removesuffix(".meta")
+                try:
+                    meta = json.loads(path.read_text(encoding="utf-8"))
+                except (json.JSONDecodeError, IOError):
+                    continue
+                if not isinstance(meta, dict):
+                    continue
 
-            matched = False
-            path_match = (meta.get("notebook_path") or "").strip()
-            os_path_match = (meta.get("notebook_os_path") or "").strip()
-            if normalized_notebook_path and path_match == normalized_notebook_path:
-                matched = True
-            if not matched and normalized_notebook_os_path and os_path_match == normalized_notebook_os_path:
-                matched = True
-            if not matched:
-                continue
+                matched = False
+                path_match = (meta.get("notebook_path") or "").strip()
+                os_path_match = (meta.get("notebook_os_path") or "").strip()
+                if normalized_notebook_path and path_match == normalized_notebook_path:
+                    matched = True
+                if not matched and normalized_notebook_os_path and os_path_match == normalized_notebook_os_path:
+                    matched = True
+                if not matched:
+                    continue
 
-            updated_at = meta.get("updated_at") or meta.get("created_at")
-            parsed_updated_at = _parse_iso_datetime(updated_at)
-            if parsed_updated_at is None:
-                continue
-            if latest_updated_at is None or parsed_updated_at > latest_updated_at:
-                latest_updated_at = parsed_updated_at
-                latest_session_id = session_id
+                updated_at = meta.get("updated_at") or meta.get("created_at")
+                parsed_updated_at = _parse_iso_datetime(updated_at)
+                if parsed_updated_at is None:
+                    continue
+                if latest_updated_at is None or parsed_updated_at > latest_updated_at:
+                    latest_updated_at = parsed_updated_at
+                    latest_session_id = session_id
 
         return latest_session_id
 
     def delete_session(self, session_id: str) -> None:
         if not self._logging_enabled:
             return
-
         normalized_session_id = (session_id or "").strip()
         if not normalized_session_id:
             return
 
-        self._delete_session_files(normalized_session_id)
+        with self._file_lock:
+            self._delete_session_files(normalized_session_id)
 
     def delete_all_sessions(self) -> tuple[int, int]:
         if not self._logging_enabled:
             return (0, 0)
+        if not self._base.exists():
+            return (0, 0)
 
         deleted_count = 0
         failed_count = 0
-        for path in self._base.glob("*.meta.json"):
-            session_id = path.stem.removesuffix(".meta")
-            if self._delete_session_files(session_id):
-                deleted_count += 1
-            else:
-                failed_count += 1
+        with self._file_lock:
+            for path in self._base.glob("*.meta.json"):
+                session_id = path.stem.removesuffix(".meta")
+                if self._delete_session_files(session_id):
+                    deleted_count += 1
+                else:
+                    failed_count += 1
 
         return (deleted_count, failed_count)
 
@@ -253,26 +416,29 @@ class SessionStore:
     ) -> None:
         if not self._logging_enabled:
             return
+        if not session_id:
+            return
 
-        meta_path = self._meta_path(session_id)
-        if meta_path.exists():
-            try:
-                meta = json.loads(meta_path.read_text(encoding="utf-8"))
-            except (json.JSONDecodeError, IOError):
-                meta = {}
-        else:
-            meta = {
-                "session_id": session_id,
-                "created_at": _now_iso(),
-            }
+        with self._file_lock:
+            meta = self._load_meta(session_id)
+            if not meta:
+                meta = {
+                    "session_id": session_id,
+                    "created_at": _now_iso(),
+                    "schema_version": _SESSION_FILE_VERSION,
+                }
 
-        paired_path, paired_os_path = _derive_paired_paths(notebook_path, notebook_os_path)
-        meta["notebook_path"] = notebook_path
-        meta["notebook_os_path"] = notebook_os_path
-        meta["paired_path"] = paired_path
-        meta["paired_os_path"] = paired_os_path
-        meta["updated_at"] = _now_iso()
-        meta_path.write_text(json.dumps(meta, indent=2), encoding="utf-8")
+            paired_path, paired_os_path = _derive_paired_paths(notebook_path, notebook_os_path)
+            meta["session_id"] = session_id
+            meta["notebook_path"] = notebook_path
+            meta["notebook_os_path"] = notebook_os_path
+            meta["paired_path"] = paired_path
+            meta["paired_os_path"] = paired_os_path
+            meta["schema_version"] = _SESSION_FILE_VERSION
+            meta["updated_at"] = _now_iso()
+            meta["retention_days"] = self._retention_days
+            meta["max_messages_per_session"] = self._max_messages_per_session
+            self._write_meta_atomic(self._meta_path(session_id), meta)
 
     def close_session(self, session_id: str) -> None:
         if not self._logging_enabled:
@@ -280,10 +446,88 @@ class SessionStore:
 
         self._touch_meta(session_id)
 
+    def rename_session(self, old_session_id: str, new_session_id: str) -> str:
+        """
+        Move session files to a new id (for example, when Codex returns a real
+        `thread_id` after the first run).
+        """
+        old_id = (old_session_id or "").strip()
+        new_id = (new_session_id or "").strip()
+        if not old_id or not new_id or old_id == new_id:
+            return new_id or old_id
+        if not self._logging_enabled:
+            return new_id
+
+        with self._file_lock:
+            old_jsonl = self._jsonl_path(old_id)
+            new_jsonl = self._jsonl_path(new_id)
+            if old_jsonl.exists():
+                if new_jsonl.exists():
+                    try:
+                        with old_jsonl.open("r", encoding="utf-8") as source, new_jsonl.open(
+                            "a", encoding="utf-8"
+                        ) as target:
+                            for line in source:
+                                if not line:
+                                    continue
+                                if line.endswith("\n"):
+                                    target.write(line)
+                                else:
+                                    target.write(f"{line}\n")
+                        old_jsonl.unlink()
+                    except OSError:
+                        pass
+                else:
+                    try:
+                        old_jsonl.rename(new_jsonl)
+                    except OSError:
+                        pass
+
+            merged_meta: Dict[str, Any] = {}
+            for candidate in (self._meta_path(new_id), self._meta_path(old_id)):
+                if not candidate.exists():
+                    continue
+                try:
+                    loaded = json.loads(candidate.read_text(encoding="utf-8"))
+                except (json.JSONDecodeError, IOError):
+                    continue
+                if isinstance(loaded, dict):
+                    merged_meta.update(loaded)
+
+            if merged_meta:
+                merged_meta["session_id"] = new_id
+                merged_meta["updated_at"] = _now_iso()
+                merged_meta["schema_version"] = _SESSION_FILE_VERSION
+                merged_meta["retention_days"] = self._retention_days
+                merged_meta["max_messages_per_session"] = self._max_messages_per_session
+                try:
+                    self._write_meta_atomic(self._meta_path(new_id), merged_meta)
+                except OSError:
+                    pass
+
+            old_meta = self._meta_path(old_id)
+            new_meta = self._meta_path(new_id)
+            if old_meta != new_meta and old_meta.exists():
+                try:
+                    old_meta.unlink()
+                except OSError:
+                    pass
+
+            # Ensure the policy metadata is present even for files created before this change.
+            if new_meta.exists():
+                self._enforce_session_limits_locked(new_id, skip_size_limit=True)
+
+        return new_id
+
     def _touch_meta(self, session_id: str) -> None:
         if not self._logging_enabled:
             return
+        if not session_id:
+            return
+        with self._file_lock:
+            self._touch_meta_locked(session_id)
 
+    def _touch_meta_locked(self, session_id: str) -> None:
         meta_path = self._meta_path(session_id)
         if not meta_path.exists():
             return
@@ -291,25 +535,43 @@ class SessionStore:
         try:
             meta = json.loads(meta_path.read_text(encoding="utf-8"))
         except (json.JSONDecodeError, IOError):
-            return
+            meta = {
+                "session_id": session_id,
+                "created_at": _now_iso(),
+                "schema_version": _SESSION_FILE_VERSION,
+            }
+        if not isinstance(meta, dict):
+            meta = {
+                "session_id": session_id,
+                "created_at": _now_iso(),
+                "schema_version": _SESSION_FILE_VERSION,
+            }
 
         meta["updated_at"] = _now_iso()
-        meta_path.write_text(json.dumps(meta, indent=2), encoding="utf-8")
+        meta["schema_version"] = _SESSION_FILE_VERSION
+        meta["retention_days"] = self._retention_days
+        self._write_meta_atomic(meta_path, meta)
 
     def _load_meta(self, session_id: str) -> Dict[str, str]:
         if not self._logging_enabled:
             return {}
+        if not session_id:
+            return {}
 
         meta_path = self._meta_path(session_id)
         if not meta_path.exists():
             return {}
 
-        try:
-            meta = json.loads(meta_path.read_text(encoding="utf-8"))
-        except (json.JSONDecodeError, IOError):
-            return {}
+        with self._file_lock:
+            try:
+                meta = json.loads(meta_path.read_text(encoding="utf-8"))
+            except (json.JSONDecodeError, IOError):
+                return {}
 
-        return meta if isinstance(meta, dict) else {}
+            if not isinstance(meta, dict):
+                return {}
+
+            return meta
 
     def _jsonl_path(self, session_id: str) -> Path:
         return self._base / f"{session_id}.jsonl"
@@ -321,23 +583,163 @@ class SessionStore:
         if not self._logging_enabled or self._retention_days <= 0:
             return
 
+        with self._file_lock:
+            self._prune_expired_sessions_locked()
+            self._last_global_prune = datetime.now(timezone.utc)
+
+    def _prune_expired_sessions_locked(self) -> None:
+        if self._retention_days <= 0:
+            return
+
         now = datetime.now(timezone.utc)
         cutoff = now - timedelta(days=self._retention_days)
+        active_session_ids = set()
+
         for path in self._base.glob("*.meta.json"):
             session_id = path.stem.removesuffix(".meta")
             try:
                 meta = json.loads(path.read_text(encoding="utf-8"))
             except (json.JSONDecodeError, IOError):
+                self._delete_session_files(session_id)
                 continue
             if not isinstance(meta, dict):
+                self._delete_session_files(session_id)
                 continue
 
             updated_at = meta.get("updated_at") or meta.get("created_at")
             when = _parse_iso_datetime(updated_at)
             if not when:
-                continue
-            if when < cutoff:
                 self._delete_session_files(session_id)
+                continue
+            if when >= cutoff:
+                active_session_ids.add(session_id)
+                continue
+
+            self._delete_session_files(session_id)
+
+        for path in self._base.glob("*.jsonl"):
+            if path.stem not in active_session_ids:
+                self._delete_session_files(path.stem)
+
+        self._last_global_prune = datetime.now(timezone.utc)
+
+    def _is_global_prune_due(self) -> bool:
+        now = datetime.now(timezone.utc)
+        if now - self._last_global_prune < self._prune_interval:
+            return False
+        return True
+
+    def _enforce_session_limits_locked(self, session_id: str, skip_size_limit: bool = False) -> None:
+        path = self._jsonl_path(session_id)
+        if not path.exists():
+            return
+
+        if not skip_size_limit:
+            should_check_size = self._max_session_bytes > 0 and path.stat().st_size > self._max_session_bytes
+        else:
+            should_check_size = False
+
+        records, invalid_count = _read_jsonl_records(path)
+        if not records and invalid_count == 0:
+            return
+
+        original_records = list(records)
+
+        if self._max_messages_per_session > 0:
+            records = records[-self._max_messages_per_session :]
+
+        self._trim_user_ui_previews(records, _DEFAULT_UI_PREVIEW_MAX_ITEMS_PER_SESSION)
+        if self._max_session_bytes > 0:
+            self._trim_records_to_byte_budget(records, self._max_session_bytes)
+
+        changed = (len(records) != len(original_records)) or (invalid_count > 0) or should_check_size
+        if not changed:
+            return
+
+        if not self._write_jsonl_records(path, records):
+            return
+
+    def _trim_session_records_locked(self, session_id: str, keep_ui_previews: int) -> None:
+        if keep_ui_previews <= 0:
+            return
+        path = self._jsonl_path(session_id)
+        if not path.exists():
+            return
+
+        records, invalid_count = _read_jsonl_records(path)
+        if not records and invalid_count == 0:
+            return
+
+        original_records = list(records)
+        self._trim_user_ui_previews(records, keep_ui_previews)
+        if self._max_messages_per_session > 0:
+            records = records[-self._max_messages_per_session :]
+
+        if self._max_session_bytes > 0:
+            self._trim_records_to_byte_budget(records, self._max_session_bytes)
+
+        if records == original_records and invalid_count == 0:
+            return
+
+        self._write_jsonl_records(path, records)
+
+    def _trim_records_to_byte_budget(self, records: List[Dict[str, Any]], max_bytes: int) -> None:
+        if max_bytes <= 0:
+            return
+
+        if not records:
+            return
+
+        serialized_sizes = [len(json.dumps(record)) for record in records]
+        total_bytes = sum(length + 1 for length in serialized_sizes)
+
+        if total_bytes <= max_bytes:
+            return
+
+        while total_bytes > max_bytes and records:
+            removed = serialized_sizes.pop(0) + 1
+            records.pop(0)
+            total_bytes -= removed
+
+    def _trim_user_ui_previews(self, records: List[Dict[str, Any]], keep_latest: int) -> None:
+        if keep_latest <= 0:
+            return
+
+        ui_indices: List[int] = []
+        for idx, record in enumerate(records):
+            if record.get("role") != "user":
+                continue
+            ui = record.get("ui")
+            if isinstance(ui, dict) and (
+                isinstance(ui.get("selectionPreview"), dict)
+                or isinstance(ui.get("cellOutputPreview"), dict)
+            ):
+                ui_indices.append(idx)
+
+        if len(ui_indices) <= keep_latest:
+            return
+
+        for idx in ui_indices[: len(ui_indices) - keep_latest]:
+            if "ui" in records[idx]:
+                records[idx].pop("ui", None)
+
+    def _write_jsonl_records(self, path: Path, records: List[Dict[str, Any]]) -> bool:
+        tmp_path = path.with_name(f".{path.name}.{uuid4().hex}.tmp")
+        try:
+            with tmp_path.open("w", encoding="utf-8") as handle:
+                for record in records:
+                    handle.write(json.dumps(record))
+                    handle.write("\n")
+            tmp_path.replace(path)
+            return True
+        except OSError:
+            return False
+
+    def _write_meta_atomic(self, path: Path, meta: Dict[str, Any]) -> None:
+        path.parent.mkdir(parents=True, exist_ok=True)
+        tmp_path = path.with_name(f".{path.name}.{uuid4().hex}.tmp")
+        tmp_path.write_text(json.dumps(meta, indent=2), encoding="utf-8")
+        tmp_path.replace(path)
 
     def _delete_session_files(self, session_id: str) -> bool:
         deleted_any = False
@@ -357,14 +759,35 @@ def _derive_paired_paths(notebook_path: str, notebook_os_path: str) -> Tuple[str
     paired_path = ""
     paired_os_path = ""
 
-    if notebook_path.endswith(".ipynb"):
+    notebook_path_lower = (notebook_path or "").lower()
+    notebook_os_path_lower = (notebook_os_path or "").lower()
+
+    if notebook_path_lower.endswith(".ipynb"):
         paired_path = notebook_path[:-6] + ".py"
-    if notebook_os_path.endswith(".ipynb"):
+    elif notebook_path_lower.endswith(".py"):
+        paired_path = notebook_path[:-3] + ".ipynb"
+    if notebook_os_path_lower.endswith(".ipynb"):
         paired_os_path = notebook_os_path[:-6] + ".py"
+    elif notebook_os_path_lower.endswith(".py"):
+        paired_os_path = notebook_os_path[:-3] + ".ipynb"
 
     return paired_path, paired_os_path
 
 
+def _normalize_notebook_mode(raw_mode: str, notebook_path: str, notebook_os_path: str) -> str:
+    mode = (raw_mode or "").strip().lower()
+    if mode in {"ipynb", "jupytext_py", "plain_py"}:
+        return mode
+
+    path = (notebook_path or "").strip().lower()
+    os_path = (notebook_os_path or "").strip().lower()
+    if path.endswith(".ipynb") or os_path.endswith(".ipynb"):
+        return "ipynb"
+    if path.endswith(".py") or os_path.endswith(".py"):
+        return "plain_py"
+    return "unsupported"
+
+
 def _now_iso() -> str:
     return datetime.now(timezone.utc).isoformat()
 
@@ -416,6 +839,43 @@ def _sanitize_message(content: str, max_chars: int) -> str:
     return _truncate_text(_sanitize_sensitive_values(content), max_chars)
 
 
+def _sanitize_ui_payload(raw: Dict[str, Any] | None) -> Dict[str, Any]:
+    if not isinstance(raw, dict):
+        return {}
+
+    payload: Dict[str, Any] = {}
+
+    selection_preview = _sanitize_ui_preview(raw.get("selectionPreview"))
+    if selection_preview:
+        payload["selectionPreview"] = selection_preview
+
+    cell_output_preview = _sanitize_ui_preview(raw.get("cellOutputPreview"))
+    if cell_output_preview:
+        payload["cellOutputPreview"] = cell_output_preview
+
+    return payload
+
+
+def _sanitize_ui_preview(raw: Any) -> Dict[str, str]:
+    if not isinstance(raw, dict):
+        return {}
+
+    location_raw = raw.get("locationLabel")
+    preview_text_raw = raw.get("previewText")
+    if not isinstance(location_raw, str) or not isinstance(preview_text_raw, str):
+        return {}
+
+    location = _sanitize_message(location_raw.strip(), _DEFAULT_UI_LABEL_MAX_CHARS)
+    preview_text = _sanitize_message(
+        preview_text_raw.replace("\r\n", "\n").replace("\r", "\n").strip(),
+        _DEFAULT_UI_PREVIEW_MAX_CHARS,
+    )
+    if not location or not preview_text:
+        return {}
+
+    return {"locationLabel": location, "previewText": preview_text}
+
+
 def _sanitize_sensitive_values(raw: str) -> str:
     sanitized = raw
     for pattern, replacement in _SENSITIVE_PATTERNS:
@@ -433,3 +893,28 @@ def _truncate_text(raw: str, max_chars: int) -> str:
     if max_chars <= 3:
         return raw[:max_chars]
     return f"{raw[: max_chars - 3]}..."
+
+
+def _read_jsonl_records(path: Path) -> Tuple[List[Dict[str, Any]], int]:
+    records: List[Dict[str, Any]] = []
+    removed_invalid_count = 0
+
+    try:
+        with path.open("r", encoding="utf-8") as handle:
+            for line in handle:
+                line = line.strip()
+                if not line:
+                    continue
+                try:
+                    payload = json.loads(line)
+                except json.JSONDecodeError:
+                    removed_invalid_count += 1
+                    continue
+                if isinstance(payload, dict):
+                    records.append(payload)
+                else:
+                    removed_invalid_count += 1
+    except OSError:
+        return [], 1
+
+    return records, removed_invalid_count