zai-cli 0.1.0__py3-none-any.whl

zai/core/session.py

@@ -0,0 +1,425 @@
+import hashlib
+import json
+import re
+import uuid
+from datetime import datetime
+from pathlib import Path
+
+from .errors import FileError
+from .security import resolve_project_path
+from .storage import (
+    atomic_write_json,
+    atomic_write_text,
+    file_lock,
+    quarantine_corrupt_file,
+    read_json,
+)
+
+SESSION_DIR = Path.home() / ".zai" / "sessions"
+SESSION_SCHEMA_VERSION = 4
+MAX_NAMED_SESSIONS = 100
+
+
+def _safe_session_name(name: str) -> str:
+    return re.sub(r"[^a-zA-Z0-9_.-]+", "_", name).strip("._") or "session"
+
+
+def _valid_lookup(value: str) -> bool:
+    return not (
+        ".." in value
+        or any(character in value for character in r"/\*?[]")
+    )
+
+
+def _project_identity(project_path: str) -> tuple[str, str]:
+    resolved = str(Path(project_path).resolve())
+    slug = re.sub(
+        r"[^a-zA-Z0-9_-]+",
+        "-",
+        Path(resolved).name,
+    ).strip("-") or "project"
+    digest = hashlib.sha256(resolved.lower().encode("utf-8")).hexdigest()[:12]
+    return resolved, f"{slug}-{digest}"
+
+
+def _message_payload(message) -> dict:
+    item = {"role": message.role, "content": message.content}
+    if getattr(message, "pinned", False):
+        item["pinned"] = True
+    if getattr(message, "tool_call_id", ""):
+        item["tool_call_id"] = message.tool_call_id
+    if getattr(message, "tool_name", ""):
+        item["tool_name"] = message.tool_name
+    tool_calls = getattr(message, "tool_calls", [])
+    if tool_calls:
+        item["tool_calls"] = [
+            {
+                "id": call.id,
+                "name": call.name,
+                "arguments": call.arguments,
+            }
+            for call in tool_calls
+        ]
+    return item
+
+
+def _session_title(messages: list[dict]) -> str:
+    for message in reversed(messages):
+        if message.get("role") == "user" and message.get("content", "").strip():
+            return " ".join(message["content"].split())[:80]
+    return "Untitled session"
+
+
+def _session_payload(
+    history: list,
+    project_path: str | None = None,
+    existing: dict | None = None,
+) -> dict:
+    messages = [_message_payload(message) for message in history]
+    now = datetime.now().isoformat(timespec="seconds")
+    existing = existing or {}
+    return {
+        "version": SESSION_SCHEMA_VERSION,
+        "id": existing.get("id") or uuid.uuid4().hex,
+        "title": _session_title(messages),
+        "project_path": (
+            str(Path(project_path).resolve())
+            if project_path
+            else existing.get("project_path")
+        ),
+        "created_at": existing.get("created_at") or now,
+        "updated_at": now,
+        "messages": messages,
+    }
+
+
+def _read_session(path: Path) -> tuple[list, dict]:
+    data = read_json(path, None, expected_type=(list, dict))
+    if data is None:
+        raise ValueError("Corrupt session")
+    if isinstance(data, list):
+        return data, {
+            "version": 1,
+            "id": hashlib.sha256(str(path).encode()).hexdigest()[:32],
+            "title": _session_title(data),
+            "project_path": None,
+        }
+    version = data.get("version", 1)
+    if version > SESSION_SCHEMA_VERSION:
+        raise ValueError(f"Unsupported session version: {version}")
+    if isinstance(data.get("messages"), list):
+        data.setdefault(
+            "id",
+            hashlib.sha256(str(path).encode()).hexdigest()[:32],
+        )
+        data.setdefault("title", _session_title(data["messages"]))
+        return data["messages"], data
+    quarantine_corrupt_file(path)
+    raise ValueError("Invalid session format")
+
+
+def _session_info(path: Path) -> dict | None:
+    try:
+        messages, metadata = _read_session(path)
+        stat = path.stat()
+        return {
+            "id": metadata["id"],
+            "name": path.stem,
+            "title": metadata.get("title") or _session_title(messages),
+            "messages": len(messages),
+            "path": str(path),
+            "project_path": metadata.get("project_path"),
+            "created_at": metadata.get("created_at"),
+            "updated_at": metadata.get("updated_at"),
+            "modified_at": stat.st_mtime,
+            "auto": path.stem.startswith("auto-"),
+        }
+    except (OSError, ValueError, KeyError):
+        return None
+
+
+def _all_session_paths() -> list[Path]:
+    if not SESSION_DIR.exists():
+        return []
+    return list(SESSION_DIR.glob("*.json"))
+
+
+def _resolve_session_path(identifier: str) -> Path | None:
+    if not identifier or not _valid_lookup(identifier):
+        return None
+    safe = _safe_session_name(identifier)
+    exact = SESSION_DIR / f"{safe}.json"
+    if exact.is_file():
+        return exact
+
+    identifier_lower = identifier.lower()
+    candidates = []
+    for path in _all_session_paths():
+        info = _session_info(path)
+        if not info:
+            continue
+        if (
+            info["id"].lower() == identifier_lower
+            or info["id"].lower().startswith(identifier_lower)
+        ):
+            candidates.append(path)
+        elif identifier_lower in info["name"].lower():
+            candidates.append(path)
+    unique = list(dict.fromkeys(candidates))
+    return unique[0] if len(unique) == 1 else None
+
+
+def _messages_to_objects(messages: list):
+    from ..providers.base import Message, ToolCall
+
+    return [
+        Message(
+            role=item["role"],
+            content=item.get("content", ""),
+            tool_call_id=item.get("tool_call_id", ""),
+            tool_name=item.get("tool_name", ""),
+            pinned=item.get("pinned", False),
+            tool_calls=[
+                ToolCall(
+                    id=call["id"],
+                    name=call["name"],
+                    arguments=call.get("arguments", {}),
+                )
+                for call in item.get("tool_calls", [])
+            ],
+        )
+        for item in messages
+    ]
+
+
+def _prune_named_sessions() -> None:
+    named = [
+        path
+        for path in _all_session_paths()
+        if not path.stem.startswith("auto-")
+    ]
+    named.sort(key=lambda path: path.stat().st_mtime, reverse=True)
+    for old in named[MAX_NAMED_SESSIONS:]:
+        old.unlink(missing_ok=True)
+        old.with_name(f"{old.name}.lock").unlink(missing_ok=True)
+
+
+def save_history(
+    history: list,
+    name: str | None = None,
+    project_path: str | None = None,
+) -> str:
+    """Save or update a named conversation and return its path."""
+    SESSION_DIR.mkdir(parents=True, exist_ok=True)
+    name = _safe_session_name(name or datetime.now().strftime("%Y%m%d-%H%M%S"))
+    if name.startswith("auto-"):
+        name = f"named-{name}"
+    path = SESSION_DIR / f"{name}.json"
+    with file_lock(SESSION_DIR / ".sessions"):
+        existing: dict = {}
+        if path.exists():
+            try:
+                _, existing = _read_session(path)
+            except ValueError:
+                existing = {}
+        atomic_write_json(path, _session_payload(history, project_path, existing))
+        _prune_named_sessions()
+    return str(path)
+
+
+def save_auto_history(history: list, project_path: str) -> str:
+    """Continuously update exactly one resumable session per project."""
+    SESSION_DIR.mkdir(parents=True, exist_ok=True)
+    _, project_id = _project_identity(project_path)
+    path = SESSION_DIR / f"auto-{project_id}.json"
+    with file_lock(SESSION_DIR / ".sessions"):
+        existing: dict = {}
+        if path.exists():
+            try:
+                _, existing = _read_session(path)
+            except ValueError:
+                existing = {}
+        atomic_write_json(path, _session_payload(history, project_path, existing))
+    return str(path)
+
+
+def load_history(identifier: str):
+    """Load one unambiguous session by exact name, ID, or unique partial name."""
+    path = _resolve_session_path(identifier)
+    if not path:
+        return None
+    try:
+        messages, _ = _read_session(path)
+        return _messages_to_objects(messages)
+    except (ValueError, KeyError):
+        return None
+
+
+def load_latest_session(project_path: str | None = None):
+    """Return (messages, metadata) for the latest matching session."""
+    expected_project = (
+        str(Path(project_path).resolve())
+        if project_path
+        else None
+    )
+    sessions = list_sessions(
+        limit=None,
+        project_path=expected_project,
+        include_auto=True,
+    )
+    for info in sessions:
+        loaded = load_history(info["id"])
+        if loaded is not None:
+            return loaded, info
+    return None
+
+
+def load_latest_history(project_path: str | None = None):
+    latest = load_latest_session(project_path)
+    return latest[0] if latest else None
+
+
+def list_sessions(
+    limit: int | None = 50,
+    project_path: str | None = None,
+    include_auto: bool = True,
+) -> list[dict]:
+    """List sessions newest first, optionally scoped to one project."""
+    expected_project = (
+        str(Path(project_path).resolve())
+        if project_path
+        else None
+    )
+    sessions = []
+    for path in _all_session_paths():
+        info = _session_info(path)
+        if not info:
+            continue
+        if not include_auto and info["auto"]:
+            continue
+        if expected_project and info["project_path"] != expected_project:
+            continue
+        sessions.append(info)
+    sessions.sort(key=lambda item: item["modified_at"], reverse=True)
+    return sessions if limit is None else sessions[:limit]
+
+
+def search_sessions(
+    query: str,
+    project_path: str | None = None,
+    limit: int = 20,
+) -> list[dict]:
+    """Search session names, titles, and message content."""
+    needle = query.strip().lower()
+    if not needle:
+        return []
+    results = []
+    for info in list_sessions(
+        limit=None,
+        project_path=project_path,
+        include_auto=True,
+    ):
+        path = Path(info["path"])
+        try:
+            messages, _ = _read_session(path)
+        except ValueError:
+            continue
+        haystack = " ".join([
+            info["name"],
+            info["title"],
+            *(message.get("content", "") for message in messages),
+        ]).lower()
+        if needle in haystack:
+            results.append(info)
+            if len(results) >= limit:
+                break
+    return results
+
+
+def get_session_info(identifier: str) -> dict | None:
+    path = _resolve_session_path(identifier)
+    return _session_info(path) if path else None
+
+
+def rename_session(identifier: str, new_name: str) -> str | None:
+    """Rename a named session. Automatic project sessions cannot be renamed."""
+    source = _resolve_session_path(identifier)
+    if not source or source.stem.startswith("auto-"):
+        return None
+    safe_name = _safe_session_name(new_name)
+    destination = SESSION_DIR / f"{safe_name}.json"
+    if destination.exists():
+        return None
+    with file_lock(SESSION_DIR / ".sessions"):
+        if not source.exists() or destination.exists():
+            return None
+        source.replace(destination)
+        source.with_name(f"{source.name}.lock").unlink(missing_ok=True)
+    return str(destination)
+
+
+def delete_session(identifier: str) -> bool:
+    path = _resolve_session_path(identifier)
+    if not path:
+        return False
+    with file_lock(SESSION_DIR / ".sessions"):
+        if not path.exists():
+            return False
+        path.unlink()
+        path.with_name(f"{path.name}.lock").unlink(missing_ok=True)
+    return True
+
+
+def export_session(
+    identifier: str,
+    export_format: str = "md",
+    output: str | None = None,
+    project_path: str = ".",
+) -> str | None:
+    """Export one session as Markdown or JSON inside the selected project."""
+    path = _resolve_session_path(identifier)
+    if not path or export_format not in {"md", "json"}:
+        return None
+    try:
+        messages, metadata = _read_session(path)
+    except ValueError:
+        return None
+
+    suffix = export_format
+    try:
+        destination = resolve_project_path(
+            project_path,
+            output or f"{path.stem}.{suffix}",
+        )
+    except FileError:
+        return None
+    if export_format == "json":
+        payload = {
+            **metadata,
+            "messages": messages,
+        }
+        atomic_write_text(
+            destination,
+            json.dumps(payload, indent=2, ensure_ascii=False) + "\n",
+            mode=0o644,
+            lock=False,
+        )
+    else:
+        lines = [
+            f"# {metadata.get('title') or _session_title(messages)}",
+            "",
+            f"- Session ID: `{metadata.get('id', '')}`",
+            f"- Project: `{metadata.get('project_path') or '-'}`",
+            f"- Updated: `{metadata.get('updated_at') or '-'}`",
+            "",
+        ]
+        for message in messages:
+            role = message.get("role", "message").replace("_", " ").title()
+            lines.extend([f"## {role}", "", message.get("content", ""), ""])
+        atomic_write_text(
+            destination,
+            "\n".join(lines).rstrip() + "\n",
+            mode=0o644,
+            lock=False,
+        )
+    return str(destination)

zai/core/storage.py

@@ -0,0 +1,193 @@
+"""Crash-safe local state storage with per-file process locking."""
+from __future__ import annotations
+
+import json
+import os
+import tempfile
+import threading
+import time
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Any, Callable, Iterator
+
+_thread_locks: dict[str, threading.RLock] = {}
+_thread_locks_guard = threading.Lock()
+
+
+def _thread_lock(path: Path) -> threading.RLock:
+    key = str(path.resolve())
+    with _thread_locks_guard:
+        return _thread_locks.setdefault(key, threading.RLock())
+
+
+@contextmanager
+def file_lock(path: str | Path) -> Iterator[None]:
+    """Hold a cross-thread and cross-process lock for one state file."""
+    target = Path(path)
+    lock_path = target.with_name(f"{target.name}.lock")
+    lock_path.parent.mkdir(parents=True, exist_ok=True)
+    thread_lock = _thread_lock(lock_path)
+    with thread_lock:
+        with open(lock_path, "a+b") as handle:
+            handle.seek(0, os.SEEK_END)
+            if handle.tell() == 0:
+                handle.write(b"\0")
+                handle.flush()
+            handle.seek(0)
+            if os.name == "nt":
+                import msvcrt
+
+                msvcrt.locking(  # type: ignore[attr-defined]
+                    handle.fileno(),
+                    msvcrt.LK_LOCK,  # type: ignore[attr-defined]
+                    1,
+                )
+            else:
+                import fcntl
+
+                fcntl.flock(  # type: ignore[attr-defined]
+                    handle.fileno(),
+                    fcntl.LOCK_EX,  # type: ignore[attr-defined]
+                )
+            try:
+                yield
+            finally:
+                handle.seek(0)
+                if os.name == "nt":
+                    import msvcrt
+
+                    msvcrt.locking(  # type: ignore[attr-defined]
+                        handle.fileno(),
+                        msvcrt.LK_UNLCK,  # type: ignore[attr-defined]
+                        1,
+                    )
+                else:
+                    import fcntl
+
+                    fcntl.flock(  # type: ignore[attr-defined]
+                        handle.fileno(),
+                        fcntl.LOCK_UN,  # type: ignore[attr-defined]
+                    )
+
+
+def atomic_write_text(
+    path: str | Path,
+    content: str,
+    *,
+    encoding: str = "utf-8",
+    lock: bool = True,
+    mode: int | None = None,
+) -> None:
+    target = Path(path)
+    target.parent.mkdir(parents=True, exist_ok=True)
+
+    def write() -> None:
+        existing_mode = None
+        try:
+            existing_mode = target.stat().st_mode & 0o777
+        except OSError:
+            pass
+        fd, temporary_name = tempfile.mkstemp(
+            prefix=f".{target.name}.",
+            suffix=".tmp",
+            dir=target.parent,
+        )
+        temporary = Path(temporary_name)
+        try:
+            with os.fdopen(fd, "w", encoding=encoding, newline="") as handle:
+                handle.write(content)
+                handle.flush()
+                os.fsync(handle.fileno())
+            if existing_mode is not None:
+                os.chmod(temporary, existing_mode)
+            elif mode is not None:
+                os.chmod(temporary, mode)
+            os.replace(temporary, target)
+        finally:
+            temporary.unlink(missing_ok=True)
+
+    if lock:
+        with file_lock(target):
+            write()
+    else:
+        write()
+
+
+def atomic_write_json(
+    path: str | Path,
+    data: Any,
+    *,
+    indent: int = 2,
+    lock: bool = True,
+) -> None:
+    atomic_write_text(
+        path,
+        json.dumps(data, indent=indent, ensure_ascii=False) + "\n",
+        lock=lock,
+    )
+
+
+def quarantine_corrupt_file(path: str | Path) -> Path | None:
+    target = Path(path)
+    if not target.exists():
+        return None
+    stamp = time.strftime("%Y%m%d-%H%M%S")
+    destination = target.with_name(f"{target.name}.corrupt-{stamp}")
+    counter = 1
+    while destination.exists():
+        destination = target.with_name(
+            f"{target.name}.corrupt-{stamp}-{counter}"
+        )
+        counter += 1
+    try:
+        os.replace(target, destination)
+        return destination
+    except OSError:
+        return None
+
+
+def read_json(
+    path: str | Path,
+    default: Any,
+    *,
+    expected_type: type | tuple[type, ...] | None = None,
+    quarantine: bool = True,
+) -> Any:
+    target = Path(path)
+    if not target.exists():
+        return default
+    with file_lock(target):
+        try:
+            data = json.loads(target.read_text(encoding="utf-8"))
+            if expected_type is not None and not isinstance(data, expected_type):
+                raise ValueError("unexpected JSON root type")
+            return data
+        except (OSError, UnicodeError, json.JSONDecodeError, ValueError):
+            if quarantine:
+                quarantine_corrupt_file(target)
+            return default
+
+
+def update_json(
+    path: str | Path,
+    default: Any,
+    updater: Callable[[Any], Any],
+    *,
+    expected_type: type | tuple[type, ...] | None = None,
+) -> Any:
+    """Read-modify-write one JSON file while holding its process lock."""
+    target = Path(path)
+    with file_lock(target):
+        if target.exists():
+            try:
+                current = json.loads(target.read_text(encoding="utf-8"))
+                if expected_type is not None and not isinstance(current, expected_type):
+                    raise ValueError("unexpected JSON root type")
+            except (OSError, UnicodeError, json.JSONDecodeError, ValueError):
+                quarantine_corrupt_file(target)
+                current = default
+        else:
+            current = default
+        updated = updater(current)
+        atomic_write_json(target, updated, lock=False)
+        return updated