dataface 0.1.2__py3-none-any.whl

dataface/agent_api/_session_store.py

@@ -0,0 +1,486 @@
+"""Session auto-save and resume for dft chat.
+
+Storage layout (mirrors Codex CLI's date-partitioned shape):
+
+    ~/.dft/sessions/
+      2026/04/28/
+        rollout-20260428142133-7e3b9f.jsonl
+      index.json   # cwd → list of SessionMeta; rebuild-able from JSONL files
+
+JSONL schema (one JSON object per line):
+
+    {"type":"meta","session_id":"...","cwd":"...","provider":"...","model":"...","started_at":"..."}
+
+    After each user prompt:
+    {"type":"turn","messages":[
+        {"role":"user","content":"..."},
+        {"role":"assistant","content":"...","tool_calls":[...]},
+        {"role":"tool","tool_call_id":"...","name":"...","content":"..."},
+        ...
+    ]}
+
+The ``turn`` event holds the exact message dicts that run_agent appended to the
+conversation for one prompt/response cycle.  Replay is trivial: concatenate all
+turn["messages"] arrays in order.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import uuid
+from collections.abc import Iterator
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+from dataface.agent_api._state import dft_home, dft_state_path
+
+logger = logging.getLogger(__name__)
+
+
+def _now_iso() -> str:
+    return datetime.now(tz=timezone.utc).astimezone().isoformat()
+
+
+def _session_filename(dt: datetime, short_id: str) -> str:
+    """Build the filename from a parsed datetime + short id."""
+    return f"rollout-{dt.strftime('%Y%m%d%H%M%S')}-{short_id}.jsonl"
+
+
+class SessionWriter:
+    """Append-only JSONL session writer.
+
+    Writes a ``meta`` line first, then one ``turn`` JSON line per prompt/response
+    cycle.  Every write flushes so a kill -9 loses at most the current in-flight
+    turn, not the whole transcript.
+    """
+
+    def __init__(
+        self,
+        session_id: str,
+        cwd: Path,
+        *,
+        provider: str,
+        model: str = "",
+        started_at: str | None = None,
+    ) -> None:
+        self.session_id = session_id
+        self.cwd = cwd
+        self.provider = provider
+        self.model = model
+        self.started_at = started_at or _now_iso()
+
+        # Derive storage path from the started_at timestamp
+        dt = datetime.fromisoformat(self.started_at)
+        yyyy = f"{dt.year:04d}"
+        mm = f"{dt.month:02d}"
+        dd = f"{dt.day:02d}"
+        fname = _session_filename(dt, session_id[:6])
+        self.path: Path = dft_state_path("sessions", yyyy, mm, dd, fname)
+
+        # Open for append; write meta line immediately.
+        self._fh = self.path.open("a")
+        meta: dict[str, Any] = {
+            "type": "meta",
+            "session_id": session_id,
+            "cwd": str(cwd),
+            "provider": provider,
+            "model": model,
+            "started_at": self.started_at,
+        }
+        self._fh.write(json.dumps(meta) + "\n")
+        self._fh.flush()
+
+    def write_turn(self, new_messages: list[dict[str, Any]]) -> None:
+        """Persist the messages appended during one prompt/response cycle."""
+        self._fh.write(
+            json.dumps({"type": "turn", "messages": new_messages}, default=str) + "\n"
+        )
+        self._fh.flush()
+
+    def close(self) -> None:
+        self._fh.close()
+
+
+class SessionReader:
+    """Iterate events from a JSONL session file."""
+
+    def __init__(self, path: Path) -> None:
+        self.path = path
+
+    def __iter__(self) -> Iterator[dict[str, Any]]:
+        with self.path.open() as f:
+            for line in f:
+                line = line.strip()
+                if not line:
+                    continue
+                try:
+                    yield json.loads(line)
+                except json.JSONDecodeError:
+                    logger.warning("Skipping malformed JSONL line in %s", self.path)
+
+
+# ---------------------------------------------------------------------------
+# Index
+# ---------------------------------------------------------------------------
+
+SessionMeta = dict[str, Any]
+
+
+class SessionIndex:
+    """Persistent index mapping cwd → list of SessionMeta.
+
+    The index is stored as ~/.dft/sessions/index.json.
+    If the file is missing or malformed the index is rebuilt by scanning
+    all JSONL files under ~/.dft/sessions/ — it is a cache, not a source of truth.
+    """
+
+    def __init__(self) -> None:
+        self._path = dft_home() / "sessions" / "index.json"
+        self._data: dict[str, list[SessionMeta]] = self._load()
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    def record(
+        self,
+        session_id: str,
+        cwd: Path,
+        started_at: str,
+        path: Path,
+    ) -> None:
+        """Add or update a session entry and persist the index."""
+        cwd_str = str(cwd)
+        entry: SessionMeta = {
+            "session_id": session_id,
+            "cwd": cwd_str,
+            "started_at": started_at,
+            "path": str(path),
+        }
+        bucket = self._data.setdefault(cwd_str, [])
+        # Replace if already present (same id), else append.
+        for i, existing in enumerate(bucket):
+            if existing["session_id"] == session_id:
+                bucket[i] = entry
+                self._save()
+                return
+        bucket.append(entry)
+        self._save()
+
+    def list_for_cwd(self, cwd: Path) -> list[SessionMeta]:
+        """Return all session metadata entries for *cwd*, newest first."""
+        entries = list(self._data.get(str(cwd), []))
+        entries.sort(key=lambda e: e.get("started_at", ""), reverse=True)
+        return entries
+
+    def list_all(self) -> list[SessionMeta]:
+        """Return all session metadata entries across all cwds, newest first."""
+        all_entries: list[SessionMeta] = []
+        for entries in self._data.values():
+            all_entries.extend(entries)
+        all_entries.sort(key=lambda e: e.get("started_at", ""), reverse=True)
+        return all_entries
+
+    def find(self, session_id: str) -> SessionMeta | None:
+        """Find a session by id, searching across all cwds."""
+        for entries in self._data.values():
+            for entry in entries:
+                if entry["session_id"] == session_id:
+                    return entry
+        return None
+
+    # ------------------------------------------------------------------
+    # Persistence helpers
+    # ------------------------------------------------------------------
+
+    def _load(self) -> dict[str, list[SessionMeta]]:
+        if self._path.exists():
+            try:
+                raw = json.loads(self._path.read_text())
+                if isinstance(raw, dict):
+                    return raw
+                logger.warning("sessions/index.json has unexpected shape; rebuilding")
+            except (json.JSONDecodeError, OSError) as exc:
+                logger.warning("sessions/index.json unreadable (%s); rebuilding", exc)
+        return self._rebuild()
+
+    def _save(self) -> None:
+        self._path.parent.mkdir(parents=True, exist_ok=True)
+        self._path.write_text(json.dumps(self._data, indent=2))
+
+    def _rebuild(self) -> dict[str, list[SessionMeta]]:
+        """Scan all JSONL files under ~/.dft/sessions/ and reconstruct the index."""
+        sessions_root = dft_home() / "sessions"
+        data: dict[str, list[SessionMeta]] = {}
+        if not sessions_root.exists():
+            return data
+        for jsonl_file in sessions_root.rglob("*.jsonl"):
+            meta = _read_meta_line(jsonl_file)
+            if meta is None:
+                continue
+            cwd = meta.get("cwd", "")
+            session_id = meta.get("session_id", "")
+            started_at = meta.get("started_at", "")
+            if not (cwd and session_id):
+                continue
+            entry: SessionMeta = {
+                "session_id": session_id,
+                "cwd": cwd,
+                "started_at": started_at,
+                "path": str(jsonl_file.relative_to(dft_home())),
+            }
+            data.setdefault(cwd, []).append(entry)
+        # Persist the rebuilt index
+        self._path.parent.mkdir(parents=True, exist_ok=True)
+        self._path.write_text(json.dumps(data, indent=2))
+        return data
+
+
+def _read_meta_line(path: Path) -> dict[str, Any] | None:
+    """Read and return the first line of a JSONL file as a dict, or None."""
+    try:
+        with path.open() as f:
+            first = f.readline().strip()
+        if not first:
+            return None
+        obj = json.loads(first)
+        return obj if isinstance(obj, dict) and obj.get("type") == "meta" else None
+    except (OSError, json.JSONDecodeError):
+        return None
+
+
+# ---------------------------------------------------------------------------
+# Replay
+# ---------------------------------------------------------------------------
+
+
+def _estimate_tokens(messages: list[dict[str, Any]]) -> int:
+    """Rough token estimate via JSON-bytes / 4. No tokenizer dep.
+
+    The 4-bytes-per-token heuristic is pessimistic enough for English-heavy
+    LLM messages; the safety margin baked into trim callers covers the noise.
+    """
+    if not messages:
+        return 0
+    return sum(len(json.dumps(m, default=str)) for m in messages) // 4
+
+
+def _split_into_units(messages: list[dict[str, Any]]) -> list[list[dict[str, Any]]]:
+    """Group messages into atomic units that must drop or stay together.
+
+    A unit is one of:
+    - A user turn (one message).
+    - An assistant turn without tool_calls (one message).
+    - An assistant turn WITH tool_calls + all following tool responses (1 + N).
+
+    Tool responses (`role == "tool"`) are never standalone — they belong to
+    the immediately preceding assistant turn that issued the tool_call.
+    Splitting on assistant boundaries preserves tool_call/tool_result pairs.
+    """
+    units: list[list[dict[str, Any]]] = []
+    current: list[dict[str, Any]] = []
+    for msg in messages:
+        role = msg.get("role")
+        if role == "tool":
+            # A tool response must follow the assistant turn that issued
+            # the tool_call — otherwise the JSONL session is corrupt.
+            if not current:
+                raise ValueError(
+                    "Session corrupt: tool response without preceding "
+                    "assistant turn. The session JSONL file is malformed."
+                )
+            current.append(msg)
+            continue
+        # New user / assistant message starts a new unit.
+        if current:
+            units.append(current)
+        current = [msg]
+    if current:
+        units.append(current)
+    return units
+
+
+def trim_messages_to_context(
+    messages: list[dict[str, Any]],
+    *,
+    max_tokens: int,
+) -> tuple[list[dict[str, Any]], int]:
+    """Drop oldest message-units until the estimated token count fits.
+
+    Tool_call/tool_result pairs drop together so the LLM never sees an
+    orphan tool response (which would error or hallucinate).
+
+    Args:
+        messages: The replayed message list.
+        max_tokens: Target token ceiling (model context minus safety margin).
+
+    Returns:
+        ``(trimmed_messages, dropped_message_count)``.
+
+    Raises:
+        ValueError: If the most-recent unit alone exceeds *max_tokens*.
+            Resuming would corrupt the conversation; the caller should ask
+            the user to start a fresh session or shrink the prompt.
+    """
+    if _estimate_tokens(messages) <= max_tokens and (
+        not messages or messages[0].get("role") == "user"
+    ):
+        return messages, 0
+
+    units = _split_into_units(messages)
+    if not units:
+        return messages, 0
+
+    def _flat(us: list[list[dict[str, Any]]]) -> list[dict[str, Any]]:
+        return [m for unit in us for m in unit]
+
+    # Drop oldest units until we fit AND the resulting first message is a
+    # user turn — LLM APIs reject conversations that start with assistant
+    # or tool. If the most-recent unit alone can't satisfy both invariants,
+    # raise rather than corrupt the conversation.
+    dropped = 0
+    while units and (
+        _estimate_tokens(_flat(units)) > max_tokens or units[0][0].get("role") != "user"
+    ):
+        if len(units) == 1:
+            raise ValueError(
+                f"Session too large to resume: the most-recent turn alone "
+                f"exceeds the {max_tokens}-token budget, or doesn't start "
+                f"with a user message. Start a fresh session or trim the inputs."
+            )
+        dropped += len(units[0])
+        units.pop(0)
+
+    return _flat(units), dropped
+
+
+def replay_into_messages(
+    reader: SessionReader,
+    *,
+    expected_provider: str,
+) -> list[dict[str, Any]]:
+    """Convert a JSONL session stream back into the LLM messages list.
+
+    Each ``turn`` event contains the exact message dicts that run_agent appended
+    during one prompt/response cycle.  Replay concatenates them in order.
+
+    Args:
+        reader: SessionReader over a session JSONL file.
+        expected_provider: The provider the current client uses. Raises
+            ValueError if it doesn't match the session's recorded provider.
+
+    Returns:
+        List of ``{"role": ..., "content": ...}`` dicts suitable for
+        passing to ``run_agent(..., messages=...)``.
+    """
+    messages: list[dict[str, Any]] = []
+    meta_seen = False
+
+    for event in reader:
+        event_type = event.get("type", "")
+
+        if event_type == "meta":
+            if not meta_seen:
+                meta_seen = True
+                recorded_provider = event.get("provider", "")
+                if recorded_provider and recorded_provider != expected_provider:
+                    raise ValueError(
+                        f"Cannot resume session {event.get('session_id', '?')}: "
+                        f"it was started with {recorded_provider}; "
+                        f"current default is {expected_provider}. "
+                        f"Run with --model {recorded_provider}:<model> to resume."
+                    )
+            continue
+
+        elif event_type == "turn":
+            turn_messages = event.get("messages", [])
+            if isinstance(turn_messages, list):
+                messages.extend(turn_messages)
+        else:
+            raise ValueError(
+                f"Unknown JSONL event type {event_type!r} in session file. "
+                "The session file may be from a newer version of dft."
+            )
+
+    return messages
+
+
+# ---------------------------------------------------------------------------
+# Load helpers used by the CLI
+# ---------------------------------------------------------------------------
+
+
+def new_session(
+    cwd: Path,
+    *,
+    provider: str,
+    model: str = "",
+) -> tuple[SessionWriter, SessionIndex]:
+    """Create a new session writer and register it in the index.
+
+    Returns:
+        (writer, index)
+    """
+    session_id = uuid.uuid4().hex
+    started_at = _now_iso()
+    writer = SessionWriter(
+        session_id=session_id,
+        cwd=cwd,
+        provider=provider,
+        model=model,
+        started_at=started_at,
+    )
+    index = SessionIndex()
+    index.record(
+        session_id=session_id,
+        cwd=cwd,
+        started_at=started_at,
+        path=writer.path.relative_to(dft_home()),
+    )
+    return writer, index
+
+
+def load_session_for_resume(
+    session_id: str,
+    *,
+    expected_provider: str,
+    max_tokens: int | None = None,
+) -> tuple[list[dict[str, Any]], SessionMeta, int]:
+    """Load and replay a session by id, optionally trimming to fit context.
+
+    Args:
+        session_id: The session id to load.
+        expected_provider: Provider to validate against.
+        max_tokens: When set, oldest tool_call/tool_result-preserving units
+            are dropped until the message count fits. ``None`` skips trim.
+
+    Returns:
+        ``(messages, meta, dropped_count)`` — replayed messages, the session
+        meta dict, and the number of messages dropped to fit ``max_tokens``
+        (0 when no trim was needed or requested).
+
+    Raises:
+        ValueError: If session not found, cross-provider mismatch, or the
+            most-recent unit alone exceeds ``max_tokens``.
+    """
+    index = SessionIndex()
+    meta = index.find(session_id)
+    if meta is None:
+        raise ValueError(f"No session found with id {session_id}")
+
+    session_path_str = meta.get("path", "")
+    if not session_path_str:
+        raise ValueError(f"No session found with id {session_id}")
+
+    session_path = dft_home() / session_path_str
+    if not session_path.exists():
+        raise ValueError(f"No session found with id {session_id}")
+
+    reader = SessionReader(session_path)
+    messages = replay_into_messages(reader, expected_provider=expected_provider)
+    dropped = 0
+    if max_tokens is not None:
+        messages, dropped = trim_messages_to_context(messages, max_tokens=max_tokens)
+    return messages, meta, dropped

dataface/agent_api/_state.py

@@ -0,0 +1,28 @@
+"""On-disk state directory for dft.
+
+Default: ~/.dft/
+Override via DFT_HOME (unconditional) or XDG_CONFIG_HOME (maps to <XDG>/dft).
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+
+def dft_home() -> Path:
+    if dft := os.environ.get("DFT_HOME"):
+        return Path(dft)
+    if xdg := os.environ.get("XDG_CONFIG_HOME"):
+        return Path(xdg) / "dft"
+    return Path.home() / ".dft"
+
+
+def dft_state_path(*parts: str) -> Path:
+    if not parts:
+        raise ValueError(
+            "dft_state_path requires at least one path component; call dft_home() for the bare home dir"
+        )
+    path = dft_home().joinpath(*parts)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    return path