pythinker-code 0.8.0__py3-none-any.whl

pythinker_code/session.py

@@ -0,0 +1,319 @@
+from __future__ import annotations
+
+import asyncio
+import builtins
+import json
+import shutil
+import uuid
+from dataclasses import dataclass
+from pathlib import Path
+
+from pythinker_core.message import Message
+from pythinker_host.path import HostPath
+
+from pythinker_code.metadata import WorkDirMeta, load_metadata, save_metadata
+from pythinker_code.session_state import SessionState, load_session_state, save_session_state
+from pythinker_code.utils.logging import logger
+from pythinker_code.utils.string import shorten
+from pythinker_code.wire.file import WireFile
+from pythinker_code.wire.types import TurnBegin
+
+
+@dataclass(slots=True, kw_only=True)
+class Session:
+    """A session of a work directory."""
+
+    # static metadata
+    id: str
+    """The session ID."""
+    work_dir: HostPath
+    """The absolute path of the work directory."""
+    work_dir_meta: WorkDirMeta
+    """The metadata of the work directory."""
+    context_file: Path
+    """The absolute path to the file storing the message history."""
+    wire_file: WireFile
+    """The wire message log file wrapper."""
+
+    # session state
+    state: SessionState
+    """Persisted session state (approval settings, plan mode, workspace scope, etc.)."""
+
+    # refreshable metadata
+    title: str
+    """The title of the session."""
+    updated_at: float
+    """The timestamp of the last update to the session."""
+
+    @property
+    def dir(self) -> Path:
+        """The absolute path of the session directory."""
+        path = self.work_dir_meta.sessions_dir / self.id
+        path.mkdir(parents=True, exist_ok=True)
+        return path
+
+    @property
+    def subagents_dir(self) -> Path:
+        """The absolute path of the subagent instances directory."""
+        path = self.dir / "subagents"
+        path.mkdir(parents=True, exist_ok=True)
+        return path
+
+    def is_empty(self) -> bool:
+        """Whether the session has any context history or a custom title."""
+        if self.state.custom_title:
+            return False
+        if not self.wire_file.is_empty():
+            return False
+        try:
+            with self.context_file.open(encoding="utf-8") as f:
+                for line in f:
+                    line = line.strip()
+                    if not line:
+                        continue
+                    role = json.loads(line, strict=False).get("role")
+                    if isinstance(role, str) and not role.startswith("_"):
+                        return False
+        except FileNotFoundError:
+            return True
+        except (OSError, ValueError, TypeError):
+            logger.exception("Failed to read context file {file}:", file=self.context_file)
+            return False
+        return True
+
+    def save_state(self) -> None:
+        """Persist the session state to disk.
+
+        Reloads externally-mutable fields (title, archive) from disk first
+        to avoid overwriting concurrent changes made by the web API.
+        """
+        fresh = load_session_state(self.dir)
+        self.state.custom_title = fresh.custom_title
+        self.state.title_generated = fresh.title_generated
+        self.state.title_generate_attempts = fresh.title_generate_attempts
+        self.state.archived = fresh.archived
+        self.state.archived_at = fresh.archived_at
+        self.state.auto_archive_exempt = fresh.auto_archive_exempt
+        save_session_state(self.state, self.dir)
+
+    async def delete(self) -> None:
+        """Delete the session directory."""
+        session_dir = self.work_dir_meta.sessions_dir / self.id
+        if not session_dir.exists():
+            return
+        await asyncio.to_thread(shutil.rmtree, session_dir, True)
+
+    async def refresh(self) -> None:
+        self.title = "Untitled"
+        self.updated_at = self.context_file.stat().st_mtime if self.context_file.exists() else 0.0
+
+        if self.state.custom_title:
+            self.title = self.state.custom_title
+            return
+
+        try:
+            async for record in self.wire_file.iter_records():
+                wire_msg = record.to_wire_message()
+                if isinstance(wire_msg, TurnBegin):
+                    self.title = shorten(
+                        Message(role="user", content=wire_msg.user_input).extract_text(" "),
+                        width=50,
+                    )
+                    return
+        except Exception:
+            logger.exception(
+                "Failed to derive session title from wire file {file}:",
+                file=self.wire_file.path,
+            )
+
+    @staticmethod
+    async def create(
+        work_dir: HostPath,
+        session_id: str | None = None,
+        _context_file: Path | None = None,
+    ) -> Session:
+        """Create a new session for a work directory."""
+        work_dir = work_dir.canonical()
+        logger.debug("Creating new session for work directory: {work_dir}", work_dir=work_dir)
+
+        metadata = load_metadata()
+        work_dir_meta = metadata.get_work_dir_meta(work_dir)
+        if work_dir_meta is None:
+            work_dir_meta = metadata.new_work_dir_meta(work_dir)
+
+        if session_id is None:
+            session_id = str(uuid.uuid4())
+        session_dir = work_dir_meta.sessions_dir / session_id
+        session_dir.mkdir(parents=True, exist_ok=True)
+
+        if _context_file is None:
+            context_file = session_dir / "context.jsonl"
+        else:
+            logger.warning(
+                "Using provided context file: {context_file}", context_file=_context_file
+            )
+            _context_file.parent.mkdir(parents=True, exist_ok=True)
+            if _context_file.exists():
+                assert _context_file.is_file()
+            context_file = _context_file
+
+        if context_file.exists():
+            # truncate if exists
+            logger.warning(
+                "Context file already exists, truncating: {context_file}", context_file=context_file
+            )
+            context_file.unlink()
+        context_file.touch()
+
+        save_metadata(metadata)
+
+        session = Session(
+            id=session_id,
+            work_dir=work_dir,
+            work_dir_meta=work_dir_meta,
+            context_file=context_file,
+            wire_file=WireFile(path=session_dir / "wire.jsonl"),
+            state=SessionState(),
+            title="",
+            updated_at=0.0,
+        )
+        await session.refresh()
+        return session
+
+    @staticmethod
+    async def find(work_dir: HostPath, session_id: str) -> Session | None:
+        """Find a session by work directory and session ID."""
+        work_dir = work_dir.canonical()
+        logger.debug(
+            "Finding session for work directory: {work_dir}, session ID: {session_id}",
+            work_dir=work_dir,
+            session_id=session_id,
+        )
+
+        metadata = load_metadata()
+        work_dir_meta = metadata.get_work_dir_meta(work_dir)
+        if work_dir_meta is None:
+            logger.debug("Work directory never been used")
+            return None
+
+        _migrate_session_context_file(work_dir_meta, session_id)
+
+        session_dir = work_dir_meta.sessions_dir / session_id
+        if not session_dir.is_dir():
+            logger.debug("Session directory not found: {session_dir}", session_dir=session_dir)
+            return None
+
+        context_file = session_dir / "context.jsonl"
+        if not context_file.exists():
+            logger.debug(
+                "Session context file not found: {context_file}", context_file=context_file
+            )
+            return None
+
+        session = Session(
+            id=session_id,
+            work_dir=work_dir,
+            work_dir_meta=work_dir_meta,
+            context_file=context_file,
+            wire_file=WireFile(path=session_dir / "wire.jsonl"),
+            state=load_session_state(session_dir),
+            title="",
+            updated_at=0.0,
+        )
+        await session.refresh()
+        return session
+
+    @staticmethod
+    async def list(work_dir: HostPath) -> builtins.list[Session]:
+        """List all sessions for a work directory."""
+        work_dir = work_dir.canonical()
+        logger.debug("Listing sessions for work directory: {work_dir}", work_dir=work_dir)
+
+        metadata = load_metadata()
+        work_dir_meta = metadata.get_work_dir_meta(work_dir)
+        if work_dir_meta is None:
+            logger.debug("Work directory never been used")
+            return []
+
+        session_ids = {
+            path.name if path.is_dir() else path.stem
+            for path in work_dir_meta.sessions_dir.iterdir()
+            if path.is_dir() or path.suffix == ".jsonl"
+        }
+
+        sessions: list[Session] = []
+        for session_id in session_ids:
+            _migrate_session_context_file(work_dir_meta, session_id)
+            session_dir = work_dir_meta.sessions_dir / session_id
+            if not session_dir.is_dir():
+                logger.debug("Session directory not found: {session_dir}", session_dir=session_dir)
+                continue
+            context_file = session_dir / "context.jsonl"
+            if not context_file.exists():
+                logger.debug(
+                    "Session context file not found: {context_file}", context_file=context_file
+                )
+                continue
+            session = Session(
+                id=session_id,
+                work_dir=work_dir,
+                work_dir_meta=work_dir_meta,
+                context_file=context_file,
+                wire_file=WireFile(path=session_dir / "wire.jsonl"),
+                state=load_session_state(session_dir),
+                title="",
+                updated_at=0.0,
+            )
+            if session.is_empty():
+                logger.debug(
+                    "Session context file is empty: {context_file}", context_file=context_file
+                )
+                continue
+            await session.refresh()
+            sessions.append(session)
+        sessions.sort(key=lambda session: session.updated_at, reverse=True)
+        return sessions
+
+    @classmethod
+    async def list_all(cls) -> builtins.list[Session]:
+        """List sessions across all known work directories."""
+        all_sessions: list[Session] = []
+        for wd in load_metadata().work_dirs:
+            sessions = await cls.list(HostPath.unsafe_from_local_path(Path(wd.path)))
+            all_sessions.extend(sessions)
+        all_sessions.sort(key=lambda s: s.updated_at, reverse=True)
+        return all_sessions
+
+    @staticmethod
+    async def continue_(work_dir: HostPath) -> Session | None:
+        """Get the last session for a work directory."""
+        work_dir = work_dir.canonical()
+        logger.debug("Continuing session for work directory: {work_dir}", work_dir=work_dir)
+
+        metadata = load_metadata()
+        work_dir_meta = metadata.get_work_dir_meta(work_dir)
+        if work_dir_meta is None:
+            logger.debug("Work directory never been used")
+            return None
+        if work_dir_meta.last_session_id is None:
+            logger.debug("Work directory never had a session")
+            return None
+
+        logger.debug(
+            "Found last session for work directory: {session_id}",
+            session_id=work_dir_meta.last_session_id,
+        )
+        return await Session.find(work_dir, work_dir_meta.last_session_id)
+
+
+def _migrate_session_context_file(work_dir_meta: WorkDirMeta, session_id: str) -> None:
+    old_context_file = work_dir_meta.sessions_dir / f"{session_id}.jsonl"
+    new_context_file = work_dir_meta.sessions_dir / session_id / "context.jsonl"
+    if old_context_file.exists() and not new_context_file.exists():
+        new_context_file.parent.mkdir(parents=True, exist_ok=True)
+        old_context_file.rename(new_context_file)
+        logger.info(
+            "Migrated session context file from {old} to {new}",
+            old=old_context_file,
+            new=new_context_file,
+        )

pythinker_code/session_fork.py

@@ -0,0 +1,325 @@
+"""Session fork utilities.
+
+Provides turn enumeration, wire/context truncation, and session forking
+for both the Web API and CLI slash commands (/undo, /fork).
+"""
+
+from __future__ import annotations
+
+import json
+import mimetypes
+import re
+import shutil
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, cast
+
+from pythinker_code.session_state import load_session_state, save_session_state
+
+CHECKPOINT_USER_PATTERN = re.compile(r"^<system>CHECKPOINT \d+</system>$")
+
+
+@dataclass(frozen=True, slots=True)
+class TurnInfo:
+    """Summary of a single turn for display in the selector."""
+
+    index: int
+    """0-based turn index."""
+    user_text: str
+    """First-line text of the user message."""
+
+
+# ---------------------------------------------------------------------------
+# Turn enumeration
+# ---------------------------------------------------------------------------
+
+
+def enumerate_turns(wire_path: Path) -> list[TurnInfo]:
+    """Scan wire.jsonl and return a list of all turns with user message text.
+
+    Each turn is identified by a ``TurnBegin`` record, whose ``user_input``
+    field provides the user message text.
+    """
+    if not wire_path.exists():
+        return []
+
+    turns: list[TurnInfo] = []
+    current_turn = -1
+
+    with open(wire_path, encoding="utf-8") as f:
+        for line in f:
+            stripped = line.strip()
+            if not stripped:
+                continue
+
+            try:
+                record: dict[str, Any] = json.loads(stripped)
+            except json.JSONDecodeError:
+                continue
+
+            if record.get("type") == "metadata":
+                continue
+
+            message: dict[str, Any] = record.get("message", {})
+            msg_type: str | None = message.get("type")
+
+            if msg_type == "TurnBegin":
+                current_turn += 1
+                user_input = message.get("payload", {}).get("user_input", "")
+                text = _extract_user_text(user_input)
+                turns.append(TurnInfo(index=current_turn, user_text=text))
+
+    return turns
+
+
+def _extract_user_text(user_input: str | list[Any]) -> str:
+    """Extract plain text from a TurnBegin user_input (string or content parts)."""
+    if isinstance(user_input, str):
+        return user_input
+
+    parts: list[str] = []
+    for part in user_input:
+        if isinstance(part, dict):
+            part_dict = cast(dict[str, Any], part)
+            text = part_dict.get("text")
+            if isinstance(text, str):
+                parts.append(text)
+        elif isinstance(part, str):
+            parts.append(part)
+    return " ".join(parts)
+
+
+# ---------------------------------------------------------------------------
+# Wire / context truncation (extracted from web/api/sessions.py)
+# ---------------------------------------------------------------------------
+
+
+def truncate_wire_at_turn(wire_path: Path, turn_index: int) -> list[str]:
+    """Read wire.jsonl and return all lines up to and including the given turn.
+
+    Args:
+        wire_path: Path to the wire.jsonl file
+        turn_index: 0-based turn index. Returns turns 0..turn_index inclusive.
+
+    Returns:
+        List of raw JSON lines (including the metadata header)
+
+    Raises:
+        ValueError: If turn_index is out of range
+    """
+    if not wire_path.exists():
+        raise ValueError("wire.jsonl not found")
+
+    lines: list[str] = []
+    current_turn = -1  # Will become 0 on first TurnBegin
+
+    with open(wire_path, encoding="utf-8") as f:
+        for line in f:
+            stripped = line.strip()
+            if not stripped:
+                continue
+
+            try:
+                record: dict[str, Any] = json.loads(stripped)
+            except json.JSONDecodeError:
+                continue
+
+            # Always keep metadata header
+            if record.get("type") == "metadata":
+                lines.append(stripped)
+                continue
+
+            message: dict[str, Any] = record.get("message", {})
+            msg_type: str | None = message.get("type")
+
+            if msg_type == "TurnBegin":
+                current_turn += 1
+                if current_turn > turn_index:
+                    break
+
+            if current_turn <= turn_index:
+                lines.append(stripped)
+
+            # Stop after the TurnEnd of the target turn
+            if msg_type == "TurnEnd" and current_turn == turn_index:
+                break
+
+    if current_turn < turn_index:
+        raise ValueError(f"turn_index {turn_index} out of range (max turn: {current_turn})")
+
+    return lines
+
+
+def _is_checkpoint_user_message(record: dict[str, Any]) -> bool:
+    """Whether a context line is the synthetic user checkpoint marker."""
+    if record.get("role") != "user":
+        return False
+
+    content = record.get("content")
+    if isinstance(content, str):
+        return CHECKPOINT_USER_PATTERN.fullmatch(content.strip()) is not None
+
+    parts = cast(list[Any], content) if isinstance(content, list) else []
+    if len(parts) == 1 and isinstance(parts[0], dict):
+        first_part = cast(dict[str, Any], parts[0])
+        text = first_part.get("text")
+        if isinstance(text, str):
+            return CHECKPOINT_USER_PATTERN.fullmatch(text.strip()) is not None
+
+    return False
+
+
+def truncate_context_at_turn(context_path: Path, turn_index: int) -> list[str]:
+    """Read context.jsonl and return all lines up to and including the given turn.
+
+    Turn detection is based on real user messages, excluding synthetic checkpoint
+    user entries like ``<system>CHECKPOINT N</system>``.
+
+    Unlike wire truncation, this is best-effort: if context has fewer user turns
+    than ``turn_index`` (e.g. slash-command turns that did not mutate context),
+    return all available context lines instead of failing.
+    """
+    if not context_path.exists():
+        return []
+
+    lines: list[str] = []
+    current_turn = -1  # Will become 0 on first real user message
+
+    with open(context_path, encoding="utf-8") as f:
+        for line in f:
+            stripped = line.strip()
+            if not stripped:
+                continue
+
+            try:
+                record: dict[str, Any] = json.loads(stripped)
+            except json.JSONDecodeError:
+                continue
+
+            if record.get("role") == "user" and not _is_checkpoint_user_message(record):
+                current_turn += 1
+                if current_turn > turn_index:
+                    break
+
+            if current_turn <= turn_index:
+                lines.append(stripped)
+
+    return lines
+
+
+# ---------------------------------------------------------------------------
+# Full fork operation
+# ---------------------------------------------------------------------------
+
+
+async def fork_session(
+    source_session_dir: Path,
+    work_dir: Any,  # HostPath
+    turn_index: int | None = None,
+    title_prefix: str = "Fork",
+    source_title: str | None = None,
+) -> str:
+    """Fork a session, creating a new session with history up to the given turn.
+
+    Args:
+        source_session_dir: Path to the source session directory.
+        work_dir: The HostPath work directory.
+        turn_index: 0-based turn index (inclusive). If None, copy all turns.
+        title_prefix: Prefix for the forked session title.
+        source_title: Title of the source session (used for the fork title).
+
+    Returns:
+        The new session ID.
+
+    Raises:
+        ValueError: If turn_index is out of range.
+    """
+    from pythinker_code.session import Session as PythinkerCLISession
+
+    wire_path = source_session_dir / "wire.jsonl"
+    context_path = source_session_dir / "context.jsonl"
+
+    if turn_index is not None:
+        truncated_wire_lines = truncate_wire_at_turn(wire_path, turn_index)
+        truncated_context_lines = truncate_context_at_turn(context_path, turn_index)
+    else:
+        # Copy all content
+        truncated_wire_lines = _read_all_lines(wire_path)
+        truncated_context_lines = _read_all_lines(context_path)
+
+    new_session = await PythinkerCLISession.create(work_dir=work_dir)
+    new_session_dir = new_session.dir
+
+    # Copy referenced video files
+    _copy_referenced_videos(source_session_dir, new_session_dir, truncated_wire_lines)
+
+    # Write truncated wire.jsonl
+    new_wire_path = new_session_dir / "wire.jsonl"
+    with open(new_wire_path, "w", encoding="utf-8") as f:
+        for line in truncated_wire_lines:
+            f.write(line + "\n")
+
+    # Write truncated context.jsonl (overwrites the empty file from create())
+    new_context_path = new_session_dir / "context.jsonl"
+    with open(new_context_path, "w", encoding="utf-8") as f:
+        for line in truncated_context_lines:
+            f.write(line + "\n")
+
+    # Set title
+    if source_title is None:
+        src_state = load_session_state(source_session_dir)
+        source_title = src_state.custom_title or "Untitled"
+
+    fork_title = f"{title_prefix}: {source_title}"
+
+    new_state = load_session_state(new_session_dir)
+    new_state.custom_title = fork_title
+    new_state.title_generated = True
+    new_state.wire_mtime = new_wire_path.stat().st_mtime
+    save_session_state(new_state, new_session_dir)
+
+    return new_session.id
+
+
+def _read_all_lines(path: Path) -> list[str]:
+    """Read all non-empty lines from a file."""
+    if not path.exists():
+        return []
+    lines: list[str] = []
+    with open(path, encoding="utf-8") as f:
+        for line in f:
+            stripped = line.strip()
+            if stripped:
+                lines.append(stripped)
+    return lines
+
+
+def _copy_referenced_videos(
+    source_dir: Path,
+    new_session_dir: Path,
+    wire_lines: list[str],
+) -> None:
+    """Copy video files referenced in the wire lines to the new session."""
+    source_uploads = source_dir / "uploads"
+    if not source_uploads.is_dir():
+        return
+
+    referenced_videos: set[str] = set()
+    for line in wire_lines:
+        for match in re.finditer(r"uploads/([^\"\\<>\s]+)", line):
+            fname = match.group(1)
+            mime, _ = mimetypes.guess_type(fname)
+            if mime and mime.startswith("video/"):
+                referenced_videos.add(fname)
+
+    files_to_copy = [
+        source_uploads / name for name in referenced_videos if (source_uploads / name).is_file()
+    ]
+    if files_to_copy:
+        new_uploads = new_session_dir / "uploads"
+        new_uploads.mkdir(parents=True, exist_ok=True)
+        copied_names: list[str] = []
+        for vf in files_to_copy:
+            shutil.copy2(vf, new_uploads / vf.name)
+            copied_names.append(vf.name)
+        (new_uploads / ".sent").write_text(json.dumps(copied_names), encoding="utf-8")