Grayscale/nanobot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(session): preserve history across storage relocation

Browse Source 
Co-authored-by: lmzopq <1646888+lmzopq@users.noreply.github.com>
This commit is contained in:
Xubin Ren
2026-08-13 01:41:10 +09:00
co-authored by lmzopq
parent 45245b5e55
commit cd7480945b
18 changed files with 820 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nanobot/agent/loop.py
+6
View File
@@ -476,6 +476,12 @@ class AgentLoop:
if bus is None:
bus = MessageBus()
defaults = config.agents.defaults
if "session_manager" not in extra:
data_dir = config.runtime_data_dir
extra["session_manager"] = SessionManager(
config.workspace_path,
sessions_root=data_dir / "sessions" if data_dir is not None else None,
)
provider = extra.pop("provider", None) or make_provider(config)
resolved = config.resolve_preset()
model = extra.pop("model", None) or resolved.model
nanobot/cli/commands.py
+38
View File
@@ -440,6 +440,44 @@ app.add_typer(
app.command(name="agent")(agent)
# ============================================================================
# Session Commands
# ============================================================================
sessions_app = typer.Typer(help="Manage persisted session history")
app.add_typer(sessions_app, name="sessions")
@sessions_app.command("restore-workspace")
def sessions_restore_workspace(
config: str | None = typer.Option(None, "--config", "-c", help="Path to config file"),
workspace: str | None = typer.Option(None, "--workspace", "-w", help="Workspace directory"),
) -> None:
"""Copy sessions back into the workspace before downgrading nanobot."""
from nanobot.session.manager import SessionManager
runtime_config = _load_runtime_config(config, workspace)
data_dir = runtime_config.runtime_data_dir
manager = SessionManager(
runtime_config.workspace_path,
sessions_root=data_dir / "sessions" if data_dir is not None else None,
)
result = manager.restore_sessions_to_workspace()
console.print(
f"Restored {result.restored} session file(s) to "
f"{escape(str(runtime_config.workspace_path / 'sessions'))}; "
f"{result.unchanged} already matched."
)
if result.conflicts:
console.print(
"[red]Rollback is incomplete: existing or invalid files require manual review.[/red]"
)
for path in result.conflicts:
console.print(Text(f"- {path}", style="red"))
raise typer.Exit(1)
# ============================================================================
# Channel Commands
# ============================================================================
nanobot/config/loader.py
+2
View File
@@ -75,6 +75,7 @@ def load_config(config_path: Path | None = None) -> Config:
summary="Environment-based configuration is invalid.",
issues=validation_issues(exc),
) from exc
config.bind_source_path(path)
_apply_ssrf_whitelist(config)
return config
@@ -130,6 +131,7 @@ def load_config(config_path: Path | None = None) -> Config:
issues=issues,
) from exc
config.bind_source_path(path)
_apply_ssrf_whitelist(config)
return config
nanobot/config/schema.py
+12 -1
View File
@@ -4,7 +4,7 @@ from __future__ import annotations
from pathlib import Path
from typing import TYPE_CHECKING, Any, ClassVar, Literal, cast
from pydantic import AliasChoices, ConfigDict, Field, field_validator, model_validator
from pydantic import AliasChoices, ConfigDict, Field, PrivateAttr, field_validator, model_validator
from pydantic_settings import BaseSettings, SettingsConfigDict
from nanobot.config.timezone import detect_system_timezone
@@ -431,6 +431,8 @@ class ToolsConfig(Base):
class Config(BaseSettings):
"""Root configuration for nanobot."""
_source_path: Path | None = PrivateAttr(default=None)
agents: AgentsConfig = Field(default_factory=AgentsConfig)
channels: ChannelsConfig = Field(default_factory=ChannelsConfig)
transcription: TranscriptionConfig = Field(default_factory=TranscriptionConfig)
@@ -449,6 +451,15 @@ class Config(BaseSettings):
_resolve_tool_config_refs()
super().__init__(**values)
def bind_source_path(self, path: Path) -> None:
"""Record the config file that owns instance-level runtime data."""
self._source_path = path.expanduser().resolve(strict=False)
@property
def runtime_data_dir(self) -> Path | None:
"""Return the active instance data directory when loaded from a config path."""
return self._source_path.parent if self._source_path is not None else None
@model_validator(mode="after")
def _validate_model_preset(self) -> "Config":
if "default" in self.model_presets:
nanobot/session/manager.py
+461 -31
View File
@@ -6,7 +6,8 @@ import hashlib
import json
import os
import re
import shutil
import secrets
import stat
from collections import OrderedDict
from contextlib import suppress
from copy import deepcopy
@@ -16,9 +17,10 @@ from pathlib import Path
from typing import Any, Callable, Collection, Protocol, TypedDict, cast
from weakref import WeakValueDictionary
from filelock import FileLock
from loguru import logger
from nanobot.config.paths import get_legacy_sessions_dir
from nanobot.config.paths import get_legacy_sessions_dir, get_runtime_subdir
from nanobot.providers.base import ProviderConversationState
from nanobot.runtime_context import (
RUNTIME_CONTEXT_HISTORY_META,
@@ -59,6 +61,11 @@ _FORK_VOLATILE_METADATA_KEYS = {
"title",
"title_user_edited",
}
_WORKSPACE_STATE_DIR = ".nanobot"
_WORKSPACE_ID_FILE = "workspace-id"
_WORKSPACE_ID_RE = re.compile(r"^[0-9a-f]{32}$")
_SESSION_MIGRATION_LOCK_TIMEOUT_SECONDS = 30
_COPY_CHUNK_SIZE = 1024 * 1024
def _json_object(value: object) -> dict[str, Any]:
@@ -505,6 +512,23 @@ class SessionInfo(TypedDict):
path: str
@dataclass(frozen=True)
class _SessionFileSnapshot:
digest: str
size: int
mtime_ns: int
updated_at: float
device: int
inode: int
@dataclass(frozen=True)
class SessionRestoreResult:
restored: int
unchanged: int
conflicts: tuple[Path, ...]
class SessionStore(Protocol):
def load(self, key: str) -> Session | None: ...
@@ -522,34 +546,356 @@ class SessionStore(Protocol):
class JsonlSessionStore:
"""JSONL implementation of session persistence."""
def __init__(self, workspace: Path):
root = get_legacy_sessions_dir()
self.sessions_dir = ensure_dir(root / self._workspace_hash(workspace))
self.legacy_sessions_dir = root
self._write_workspace_marker(self.sessions_dir, workspace)
self._migrate_from_workspace(workspace)
@staticmethod
def _workspace_hash(workspace: Path) -> str:
canonical = str(Path(workspace).expanduser().resolve(strict=False))
return hashlib.sha256(canonical.encode("utf-8")).hexdigest()[:16]
@staticmethod
def _write_workspace_marker(sessions_dir: Path, workspace: Path) -> None:
marker = sessions_dir / ".workspace"
if marker.exists():
return
try:
marker.write_text(
str(Path(workspace).expanduser().resolve(strict=False)),
encoding="utf-8",
def __init__(self, workspace: Path, *, sessions_root: Path | None = None):
canonical_workspace = Path(workspace).expanduser().resolve(strict=False)
ensure_dir(canonical_workspace)
root = (
Path(sessions_root).expanduser().resolve(strict=False)
if sessions_root is not None
else get_runtime_subdir("sessions").resolve(strict=False)
)
if root == canonical_workspace or root.is_relative_to(canonical_workspace):
raise RuntimeError(
"session storage must be outside the agent workspace; "
"move --config outside --workspace or choose a nested workspace directory"
)
except OSError as exc:
logger.debug("Failed to write sessions workspace marker: {}", exc)
ensure_dir(root)
with suppress(OSError):
os.chmod(root, 0o700)
self.workspace = canonical_workspace
self._migration_lock = FileLock(
str(root / ".workspace-migration.lock"),
timeout=_SESSION_MIGRATION_LOCK_TIMEOUT_SECONDS,
)
with self._migration_lock:
workspace_id = self._load_or_create_workspace_id(canonical_workspace, root)
workspace_id = self._claim_workspace_namespace(
root,
canonical_workspace,
workspace_id,
)
self.sessions_dir = ensure_dir(root / workspace_id)
self.legacy_sessions_dir = get_legacy_sessions_dir()
self._migrate_from_workspace(canonical_workspace)
@staticmethod
def _fsync_directory(path: Path) -> None:
with suppress(PermissionError, NotImplementedError):
fd = os.open(path, os.O_RDONLY)
try:
os.fsync(fd)
except OSError as exc:
if exc.errno != errno.EINVAL:
raise
finally:
os.close(fd)
@classmethod
def _write_text_atomic(cls, path: Path, content: str, *, mode: int = 0o600) -> None:
tmp = path.with_name(f".{path.name}.{secrets.token_hex(8)}.tmp")
try:
with open(tmp, "x", encoding="utf-8") as handle:
os.chmod(tmp, mode)
handle.write(content)
handle.flush()
os.fsync(handle.fileno())
os.replace(tmp, path)
cls._fsync_directory(path.parent)
finally:
tmp.unlink(missing_ok=True)
@classmethod
def _read_workspace_id(cls, marker: Path) -> str:
if marker.is_symlink():
raise RuntimeError(f"workspace identity marker must not be a symlink: {marker}")
value = marker.read_text(encoding="utf-8").strip()
if not _WORKSPACE_ID_RE.fullmatch(value):
raise RuntimeError(
f"workspace identity marker is invalid: {marker}; "
"restore its original 32-character identifier before starting nanobot"
)
return value
@staticmethod
def _workspace_id_path(workspace: Path) -> Path:
state_dir = workspace / _WORKSPACE_STATE_DIR
if state_dir.is_symlink():
raise RuntimeError(f"workspace state directory must not be a symlink: {state_dir}")
ensure_dir(state_dir)
return state_dir / _WORKSPACE_ID_FILE
@classmethod
def _find_workspace_namespace(cls, workspace: Path, root: Path) -> str | None:
"""Recover an identity marker removed by cleanup at the same workspace path."""
matches: list[str] = []
for sessions_dir in root.iterdir():
if (
not _WORKSPACE_ID_RE.fullmatch(sessions_dir.name)
or sessions_dir.is_symlink()
or not sessions_dir.is_dir()
):
continue
marker = sessions_dir / ".workspace"
if marker.is_symlink() or not marker.is_file():
continue
try:
recorded = Path(marker.read_text(encoding="utf-8").strip()).expanduser()
recorded = recorded.resolve(strict=False)
same_workspace = recorded == workspace or (
recorded.exists() and recorded.samefile(workspace)
)
except (OSError, UnicodeError, ValueError):
continue
if same_workspace:
matches.append(sessions_dir.name)
if len(matches) > 1:
raise RuntimeError(
f"multiple session namespaces claim workspace {workspace}; "
"remove the stale namespace marker before starting nanobot"
)
return matches[0] if matches else None
@classmethod
def _load_or_create_workspace_id(cls, workspace: Path, root: Path) -> str:
marker = cls._workspace_id_path(workspace)
if marker.exists() or marker.is_symlink():
return cls._read_workspace_id(marker)
recovered = cls._find_workspace_namespace(workspace, root)
if recovered is not None:
cls._write_text_atomic(marker, f"{recovered}\n")
return recovered
workspace_id = secrets.token_hex(16)
flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL
flags |= getattr(os, "O_NOFOLLOW", 0)
try:
fd = os.open(marker, flags, 0o600)
except FileExistsError:
return cls._read_workspace_id(marker)
try:
payload = f"{workspace_id}\n".encode("ascii")
view = memoryview(payload)
while view:
written = os.write(fd, view)
view = view[written:]
os.fsync(fd)
except BaseException:
with suppress(OSError):
marker.unlink()
raise
finally:
os.close(fd)
cls._fsync_directory(marker.parent)
return workspace_id
@classmethod
def _replace_workspace_id(cls, workspace: Path, workspace_id: str) -> None:
cls._write_text_atomic(cls._workspace_id_path(workspace), f"{workspace_id}\n")
@classmethod
def _write_workspace_marker(cls, sessions_dir: Path, workspace: Path) -> None:
cls._write_text_atomic(sessions_dir / ".workspace", f"{workspace}\n")
@classmethod
def _claim_workspace_namespace(
cls,
root: Path,
workspace: Path,
workspace_id: str,
) -> str:
"""Bind a stable workspace ID, rotating copied live workspaces apart."""
for _attempt in range(3):
sessions_dir = root / workspace_id
marker = sessions_dir / ".workspace"
if sessions_dir.is_symlink():
raise RuntimeError(f"session namespace must not be a symlink: {sessions_dir}")
if not sessions_dir.exists():
ensure_dir(sessions_dir)
cls._write_workspace_marker(sessions_dir, workspace)
return workspace_id
if marker.is_symlink():
raise RuntimeError(f"session workspace marker must not be a symlink: {marker}")
if not marker.exists():
if any(sessions_dir.iterdir()):
raise RuntimeError(
f"session namespace has data but no workspace marker: {sessions_dir}"
)
cls._write_workspace_marker(sessions_dir, workspace)
return workspace_id
recorded_text = marker.read_text(encoding="utf-8").strip()
if not recorded_text:
raise RuntimeError(f"session workspace marker is empty: {marker}")
recorded = Path(recorded_text).expanduser().resolve(strict=False)
if recorded == workspace:
return workspace_id
try:
same_workspace = recorded.exists() and recorded.samefile(workspace)
except OSError:
same_workspace = False
if same_workspace:
cls._write_workspace_marker(sessions_dir, workspace)
return workspace_id
if not recorded.exists():
# The identity marker travelled with a renamed or moved workspace.
cls._write_workspace_marker(sessions_dir, workspace)
return workspace_id
# Both paths exist and are different: this is a copy, not a move.
workspace_id = secrets.token_hex(16)
cls._replace_workspace_id(workspace, workspace_id)
raise RuntimeError(f"could not allocate an isolated session namespace for {workspace}")
@staticmethod
def _session_file_snapshot(path: Path) -> _SessionFileSnapshot | None:
flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)
try:
fd = os.open(path, flags)
except OSError:
return None
try:
before = os.fstat(fd)
if not stat.S_ISREG(before.st_mode):
return None
digest = hashlib.sha256()
saw_record = False
updated_at: float | None = None
with os.fdopen(fd, "rb", closefd=False) as handle:
for raw_line in handle:
digest.update(raw_line)
if not raw_line.strip():
continue
value: object = json.loads(raw_line.decode("utf-8"))
data = _json_object(value)
saw_record = True
if data.get("_type") == "metadata":
raw_updated_at = cast(object, data.get("updated_at"))
if isinstance(raw_updated_at, str) and raw_updated_at:
updated_at = datetime.fromisoformat(raw_updated_at).timestamp()
after = os.fstat(fd)
if (
not saw_record
or before.st_dev != after.st_dev
or before.st_ino != after.st_ino
or before.st_size != after.st_size
or before.st_mtime_ns != after.st_mtime_ns
):
return None
return _SessionFileSnapshot(
digest=digest.hexdigest(),
size=after.st_size,
mtime_ns=after.st_mtime_ns,
updated_at=(updated_at if updated_at is not None else after.st_mtime_ns / 1e9),
device=after.st_dev,
inode=after.st_ino,
)
except (OSError, UnicodeDecodeError, json.JSONDecodeError, ValueError, TypeError):
return None
finally:
os.close(fd)
@classmethod
def _prepare_copy(
cls,
src: Path,
dst_dir: Path,
snapshot: _SessionFileSnapshot,
) -> Path:
tmp = dst_dir / f".{src.name}.{secrets.token_hex(8)}.tmp"
flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)
src_fd = os.open(src, flags)
try:
before = os.fstat(src_fd)
if (
before.st_dev != snapshot.device
or before.st_ino != snapshot.inode
or before.st_size != snapshot.size
or before.st_mtime_ns != snapshot.mtime_ns
):
raise OSError("session source changed before migration")
digest = hashlib.sha256()
size = 0
with os.fdopen(src_fd, "rb", closefd=False) as source, open(tmp, "xb") as target:
os.chmod(tmp, 0o600)
while chunk := source.read(_COPY_CHUNK_SIZE):
digest.update(chunk)
size += len(chunk)
target.write(chunk)
target.flush()
os.fsync(target.fileno())
after = os.fstat(src_fd)
if (
digest.hexdigest() != snapshot.digest
or size != snapshot.size
or after.st_dev != snapshot.device
or after.st_ino != snapshot.inode
or after.st_size != snapshot.size
or after.st_mtime_ns != snapshot.mtime_ns
):
raise OSError("session source changed during migration")
return tmp
except BaseException:
tmp.unlink(missing_ok=True)
raise
finally:
os.close(src_fd)
@classmethod
def _install_snapshot(
cls,
src: Path,
dst: Path,
snapshot: _SessionFileSnapshot,
) -> None:
tmp = cls._prepare_copy(src, dst.parent, snapshot)
try:
os.replace(tmp, dst)
cls._fsync_directory(dst.parent)
installed = cls._session_file_snapshot(dst)
if installed is None or installed.digest != snapshot.digest:
raise OSError(f"session migration verification failed: {dst}")
finally:
tmp.unlink(missing_ok=True)
def _archive_conflict(
self,
src: Path,
snapshot: _SessionFileSnapshot,
label: str,
) -> Path:
conflict_dir = ensure_dir(self.sessions_dir / ".migration-conflicts")
conflict = conflict_dir / (
f"{src.stem}.{label}.{snapshot.digest[:12]}.{secrets.token_hex(4)}.jsonl"
)
self._install_snapshot(src, conflict, snapshot)
return conflict
@classmethod
def _remove_migrated_source(
cls,
src: Path,
snapshot: _SessionFileSnapshot,
) -> bool:
try:
current = src.stat(follow_symlinks=False)
if (
current.st_dev != snapshot.device
or current.st_ino != snapshot.inode
or current.st_size != snapshot.size
or current.st_mtime_ns != snapshot.mtime_ns
):
return False
src.unlink()
cls._fsync_directory(src.parent)
return True
except OSError:
return False
def _migrate_from_workspace(self, workspace: Path) -> None:
"""Move legacy in-workspace session files into the out-of-workspace store."""
old_dir = Path(workspace).expanduser() / "sessions"
"""Durably copy legacy sessions out of the workspace, then remove the source."""
old_dir = workspace / "sessions"
if old_dir.is_symlink() or not old_dir.is_dir():
if old_dir.is_symlink():
logger.warning("Skipping symlinked legacy sessions directory: {}", old_dir)
@@ -559,13 +905,87 @@ class JsonlSessionStore:
logger.warning("Skipping unsafe legacy session file: {}", src)
continue
dst = self.sessions_dir / src.name
if dst.exists():
source_snapshot = self._session_file_snapshot(src)
if source_snapshot is None:
logger.warning("Skipping invalid or changing legacy session file: {}", src)
continue
try:
shutil.move(str(src), str(dst))
destination_snapshot = self._session_file_snapshot(dst) if dst.exists() else None
if dst.exists() and destination_snapshot is None:
logger.warning(
"Keeping legacy session because destination is invalid: {}",
dst,
)
continue
if destination_snapshot is None:
self._install_snapshot(src, dst, source_snapshot)
elif destination_snapshot.digest == source_snapshot.digest:
pass
elif source_snapshot.updated_at > destination_snapshot.updated_at:
archived = self._archive_conflict(dst, destination_snapshot, "destination")
self._install_snapshot(src, dst, source_snapshot)
logger.warning("Archived older session migration conflict at {}", archived)
else:
archived = self._archive_conflict(src, source_snapshot, "workspace")
logger.warning("Archived older session migration conflict at {}", archived)
installed = self._session_file_snapshot(dst)
if installed is None:
raise OSError(f"session migration destination is unreadable: {dst}")
selected_digest = (
source_snapshot.digest
if destination_snapshot is None
or source_snapshot.updated_at > destination_snapshot.updated_at
else destination_snapshot.digest
)
if installed.digest != selected_digest:
raise OSError(f"session migration selected unexpected data: {dst}")
if not self._remove_migrated_source(src, source_snapshot):
logger.warning(
"Session migrated but legacy source changed or could not be removed: {}",
src,
)
except OSError as exc:
logger.warning("Failed to migrate session {}: {}", src, exc)
def restore_to_workspace(self) -> SessionRestoreResult:
"""Copy canonical sessions back for an explicit downgrade or rollback."""
restored = 0
unchanged = 0
conflicts: list[Path] = []
old_dir = self.workspace / "sessions"
if old_dir.is_symlink():
raise RuntimeError(f"refusing to restore into symlinked sessions directory: {old_dir}")
ensure_dir(old_dir)
with self._migration_lock:
for src in self.sessions_dir.glob("*.jsonl"):
if self.session_key_from_path(src) is None:
continue
source_snapshot = self._session_file_snapshot(src)
if source_snapshot is None:
conflicts.append(src)
continue
dst = old_dir / src.name
if dst.exists():
destination_snapshot = self._session_file_snapshot(dst)
if (
destination_snapshot is not None
and destination_snapshot.digest == source_snapshot.digest
):
unchanged += 1
else:
conflicts.append(dst)
continue
self._install_snapshot(src, dst, source_snapshot)
restored += 1
return SessionRestoreResult(
restored=restored,
unchanged=unchanged,
conflicts=tuple(conflicts),
)
@staticmethod
def safe_key(key: str) -> str:
return safe_filename(key.replace(":", "_"))
@@ -1033,9 +1453,15 @@ class JsonlSessionStore:
class SessionManager:
"""Manage session identity, caching, retention, and persistence."""
def __init__(self, workspace: Path, *, store: SessionStore | None = None):
def __init__(
self,
workspace: Path,
*,
store: SessionStore | None = None,
sessions_root: Path | None = None,
):
self.workspace = workspace
self._jsonl_store = JsonlSessionStore(workspace)
self._jsonl_store = JsonlSessionStore(workspace, sessions_root=sessions_root)
self._store: SessionStore = store if store is not None else self._jsonl_store
self.sessions_dir = self._jsonl_store.sessions_dir
self.legacy_sessions_dir = self._jsonl_store.legacy_sessions_dir
@@ -1201,6 +1627,10 @@ class SessionManager:
self.invalidate(key)
return self._store.delete(key)
def restore_sessions_to_workspace(self) -> SessionRestoreResult:
"""Restore session files to the pre-relocation path for an explicit rollback."""
return self._jsonl_store.restore_to_workspace()
def fork_session_before_user_index(
self,
source_key: str,