deep-agent-cli 0.1.0__py3-none-any.whl

agent/session.py

@@ -0,0 +1,431 @@
+"""SQLite session catalog + LangGraph SqliteSaver for local persistence."""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from enum import StrEnum
+import hashlib
+import os
+from pathlib import Path
+import sqlite3
+from typing import Any, Callable, Literal
+from uuid import uuid4
+
+from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, ToolMessage
+from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
+from langgraph.checkpoint.sqlite import SqliteSaver
+
+from agent.attachments import (
+    AttachmentCleanupResult,
+    AttachmentStore,
+    ImageAttachmentRef,
+    find_attachment_storage_keys,
+    refs_from_message,
+)
+from agent.session_lock import SessionLease, SessionLockManager
+from agent.stream import message_text, reasoning_text, visible_text
+
+SessionStatus = Literal[
+    "running", "completed", "waiting", "cancelled", "failed",
+]
+
+
+class StopReason(StrEnum):
+    PENDING = "pending"
+    STOP = "stop"
+    ERROR = "error"
+    ABORTED = "aborted"
+    DEFERRED = "deferred"
+
+_CATALOG_DDL = """
+CREATE TABLE IF NOT EXISTS session_catalog (
+    id TEXT PRIMARY KEY,
+    title TEXT NOT NULL,
+    created_at TEXT NOT NULL,
+    updated_at TEXT NOT NULL,
+    model_id TEXT,
+    permission_mode TEXT,
+    last_run_status TEXT NOT NULL
+);
+CREATE INDEX IF NOT EXISTS session_catalog_updated_idx
+    ON session_catalog (updated_at DESC);
+"""
+
+
+@dataclass(frozen=True)
+class SessionInfo:
+    id: str
+    title: str
+    created_at: datetime
+    updated_at: datetime
+    status: SessionStatus
+    model_id: str | None = None
+    permission_mode: str | None = None
+    last_run_status: StopReason = StopReason.PENDING
+
+
+@dataclass
+class TranscriptBlock:
+    kind: str
+    content: str = ""
+    thinking: str = ""
+    tool_call_id: str = ""
+    name: str = ""
+    arguments: dict[str, Any] | None = None
+    is_error: bool = False
+    status: str = ""
+    attachments: tuple[ImageAttachmentRef, ...] = ()
+    exit_code: int | None = None
+    # Raw ToolMessage.artifact restored from the checkpoint; UI previews are
+    # re-derived from it instead of being persisted.
+    artifact: Any = None
+
+
+def workspace_state_path(workspace: Path, *, override: str | Path | None = None) -> Path:
+    if override is not None:
+        return Path(override).expanduser().resolve()
+    env = os.environ.get("DEEP_AGENT_STATE_PATH")
+    if env:
+        return Path(env).expanduser().resolve()
+    digest = hashlib.sha256(str(workspace.expanduser().resolve()).encode("utf-8")).hexdigest()[:16]
+    return Path.home() / ".deep-agent" / "sessions" / f"{digest}.sqlite3"
+
+
+def _utc_now() -> datetime:
+    return datetime.now(timezone.utc)
+
+
+def _parse_dt(value: str) -> datetime:
+    return datetime.fromisoformat(value)
+
+
+def _strict_serde() -> JsonPlusSerializer:
+    # Strict allowlist: only built-in safe msgpack types are reconstructed.
+    return JsonPlusSerializer(allowed_msgpack_modules=None, pickle_fallback=False)
+
+
+class SessionStore:
+    """Presentation-neutral session API backed by one SQLite file per workspace."""
+
+    def __init__(self, path: Path) -> None:
+        self.path = path
+        self.path.parent.mkdir(parents=True, mode=0o700, exist_ok=True)
+        try:
+            os.chmod(self.path.parent, 0o700)
+        except OSError:
+            pass
+        self._conn = sqlite3.connect(str(self.path), check_same_thread=False)
+        self._conn.execute("PRAGMA journal_mode=WAL;")
+        self._conn.executescript(_CATALOG_DDL)
+        self._migrate_catalog()
+        if self.path.exists():
+            try:
+                os.chmod(self.path, 0o600)
+            except OSError:
+                pass
+        self.checkpointer = SqliteSaver(self._conn, serde=_strict_serde())
+        self.session_locks = SessionLockManager(self.path)
+        self.attachment_store = AttachmentStore.beside_database(self.path)
+        self.startup_cleanup: AttachmentCleanupResult | None = None
+        try:
+            self.startup_cleanup = self.cleanup_attachments()
+        except Exception:  # noqa: BLE001 - persistence must remain usable if maintenance fails
+            self.startup_cleanup = None
+        self.attachment_store.acquire_runtime_lease()
+        self._closed = False
+
+    @classmethod
+    def for_workspace(
+        cls,
+        workspace: Path,
+        *,
+        override: str | Path | None = None,
+    ) -> "SessionStore":
+        return cls(workspace_state_path(workspace, override=override))
+
+    def try_acquire_session(self, thread_id: str) -> SessionLease | None:
+        """Take exclusive ownership of a persisted thread; None when busy."""
+        return self.session_locks.try_acquire(thread_id)
+
+    def wait_acquire_session(
+        self,
+        thread_id: str,
+        *,
+        cancelled: Callable[[], bool] | None = None,
+    ) -> SessionLease | None:
+        """Block until the thread is free, the wait is cancelled, or ownership lands."""
+        return self.session_locks.wait_acquire(thread_id, cancelled=cancelled)
+
+    def close(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+        self.session_locks.close()
+        self.attachment_store.close()
+        self._conn.close()
+
+    def _migrate_catalog(self) -> None:
+        # Hold a write reservation while checking the schema so concurrent
+        # processes cannot both decide to migrate the same workspace database.
+        self._conn.execute("BEGIN IMMEDIATE")
+        try:
+            columns = {row[1] for row in self._conn.execute("PRAGMA table_info(session_catalog)")}
+            if "status" not in columns:
+                self._conn.commit()
+                return
+            model = "model_id" if "model_id" in columns else "NULL"
+            permission = "permission_mode" if "permission_mode" in columns else "NULL"
+            last = "last_run_status" if "last_run_status" in columns else "NULL"
+            reason = (
+                f"CASE WHEN {last} IN ('pending', 'stop', 'error', 'aborted', 'deferred') THEN {last} "
+                "WHEN status = 'completed' OR status IN ('length', 'tool_use') THEN 'stop' "
+                "WHEN status = 'cancelled' THEN 'aborted' "
+                "WHEN status = 'failed' THEN 'error' "
+                "WHEN status IN ('waiting', 'interrupted') THEN 'deferred' "
+                "ELSE 'pending' END"
+            )
+            self._conn.execute(
+                "CREATE TABLE session_catalog_new ("
+                "id TEXT PRIMARY KEY, title TEXT NOT NULL, created_at TEXT NOT NULL, "
+                "updated_at TEXT NOT NULL, model_id TEXT, permission_mode TEXT, "
+                "last_run_status TEXT NOT NULL)"
+            )
+            self._conn.execute(
+                "INSERT INTO session_catalog_new "
+                "(id, title, created_at, updated_at, model_id, permission_mode, last_run_status) "
+                f"SELECT id, title, created_at, updated_at, {model}, {permission}, {reason} "
+                "FROM session_catalog"
+            )
+            self._conn.execute("DROP TABLE session_catalog")
+            self._conn.execute("ALTER TABLE session_catalog_new RENAME TO session_catalog")
+            self._conn.execute(
+                "CREATE INDEX session_catalog_updated_idx ON session_catalog (updated_at DESC)"
+            )
+            self._conn.commit()
+        except Exception:
+            self._conn.rollback()
+            raise
+
+    def attachment_storage_keys(self) -> set[str]:
+        """Return references from every retained checkpoint and pending write."""
+        live: set[str] = set()
+        for item in self.checkpointer.list(None):
+            live.update(find_attachment_storage_keys(item.checkpoint))
+            live.update(find_attachment_storage_keys(item.pending_writes))
+        return live
+
+    def cleanup_attachments(
+        self,
+        *,
+        protected: tuple[ImageAttachmentRef, ...] = (),
+        release_runtime_lease: bool = False,
+    ) -> AttachmentCleanupResult:
+        live = self.attachment_storage_keys()
+        return self.attachment_store.cleanup(
+            live,
+            protected_storage_keys=(ref.storage_key for ref in protected),
+            release_runtime_lease=release_runtime_lease,
+        )
+
+    def create_session(
+        self, *, title: str = "", session_id: str | None = None,
+        model_id: str | None = None, permission_mode: str = "ask",
+    ) -> SessionInfo:
+        now = _utc_now()
+        info = SessionInfo(
+            id=session_id or str(uuid4()),
+            title=title or "New session",
+            created_at=now,
+            updated_at=now,
+            status="running",
+            model_id=model_id,
+            permission_mode=permission_mode,
+        )
+        with self.checkpointer.lock:
+            self._conn.execute(
+                "INSERT INTO session_catalog (id, title, created_at, updated_at, model_id, permission_mode, last_run_status) VALUES (?, ?, ?, ?, ?, ?, ?)",
+                (info.id, info.title, info.created_at.isoformat(), info.updated_at.isoformat(),
+                 info.model_id, info.permission_mode, info.last_run_status.value),
+            )
+            self._conn.commit()
+        return info
+
+    def touch(
+        self,
+        session_id: str,
+        *,
+        title: str | None = None,
+        model_id: str | None = None,
+        permission_mode: str | None = None,
+        last_run_status: StopReason | None = None,
+    ) -> None:
+        now = _utc_now().isoformat()
+        with self.checkpointer.lock:
+            row = self._conn.execute(
+                "SELECT title, model_id, permission_mode, last_run_status FROM session_catalog WHERE id = ?", (session_id,),
+            ).fetchone()
+            if row is None:
+                return
+            new_title = title if title is not None else row[0]
+            self._conn.execute(
+                "UPDATE session_catalog SET title = ?, updated_at = ?, model_id = ?, permission_mode = ?, last_run_status = ? WHERE id = ?",
+                (new_title, now, model_id if model_id is not None else row[1],
+                 permission_mode if permission_mode is not None else row[2],
+                 last_run_status.value if last_run_status is not None else row[3], session_id),
+            )
+            self._conn.commit()
+
+    def list_sessions(self, *, limit: int = 50) -> list[SessionInfo]:
+        with self.checkpointer.lock:
+            rows = self._conn.execute(
+                "SELECT id, title, created_at, updated_at, model_id, permission_mode, last_run_status "
+                "FROM session_catalog ORDER BY updated_at DESC LIMIT ?",
+                (limit,),
+            ).fetchall()
+        return [_session_info(row) for row in rows]
+
+    def get(self, session_id: str) -> SessionInfo | None:
+        with self.checkpointer.lock:
+            row = self._conn.execute(
+                "SELECT id, title, created_at, updated_at, model_id, permission_mode, last_run_status FROM session_catalog WHERE id = ?",
+                (session_id,),
+            ).fetchone()
+        if row is None:
+            return None
+        return _session_info(row)
+
+    def resolve_prefix(self, prefix: str) -> SessionInfo | None:
+        prefix = prefix.strip()
+        if not prefix:
+            return None
+        exact = self.get(prefix)
+        if exact is not None:
+            return exact
+        with self.checkpointer.lock:
+            rows = self._conn.execute(
+                "SELECT id, title, created_at, updated_at, model_id, permission_mode, last_run_status FROM session_catalog WHERE id LIKE ?",
+                (f"{prefix}%",),
+            ).fetchall()
+        if len(rows) != 1:
+            return None
+        return _session_info(rows[0])
+
+
+def _session_info(row: Any) -> SessionInfo:
+    try:
+        reason = StopReason(row[6])
+    except ValueError:
+        reason = StopReason.PENDING
+    return SessionInfo(
+        id=row[0], title=row[1], created_at=_parse_dt(row[2]), updated_at=_parse_dt(row[3]),
+        status=_status_for_reason(reason), model_id=row[4], permission_mode=row[5], last_run_status=reason,
+    )
+
+
+def _status_for_reason(reason: StopReason) -> SessionStatus:
+    return {
+        StopReason.PENDING: "running",
+        StopReason.STOP: "completed",
+        StopReason.ERROR: "failed",
+        StopReason.ABORTED: "cancelled",
+        StopReason.DEFERRED: "waiting",
+    }[reason]
+
+
+def tool_message_is_error(message: ToolMessage) -> bool:
+    """Match live tool_completed error classification, including execute artifacts."""
+    status = str(getattr(message, "status", "") or "")
+    content = message_text(message)
+    name = str(getattr(message, "name", "") or "")
+    artifact = getattr(message, "artifact", None)
+    if name == "execute" and isinstance(artifact, dict):
+        exit_code = artifact.get("exit_code")
+        return (
+            status == "error"
+            or artifact.get("termination_reason") is not None
+            or (isinstance(exit_code, int) and exit_code != 0)
+        )
+    return (
+        status == "error"
+        or content.lower().startswith("error")
+        or "cancelled by user" in content.lower()
+    )
+
+
+def messages_to_transcript(messages: list[BaseMessage]) -> list[TranscriptBlock]:
+    """Rebuild a presentation-neutral transcript from LangGraph checkpoint messages."""
+    blocks: list[TranscriptBlock] = []
+    todo_call_ids: set[str] = set()
+    for message in messages:
+        if isinstance(message, HumanMessage):
+            blocks.append(TranscriptBlock(
+                kind="user",
+                content=message_text(message),
+                attachments=refs_from_message(message),
+            ))
+        elif isinstance(message, AIMessage):
+            thinking = reasoning_text(message)
+            text = visible_text(message)
+            if thinking or text:
+                blocks.append(TranscriptBlock(kind="assistant", content=text, thinking=thinking))
+            for call in message.tool_calls or []:
+                if call.get("name") == "write_todos":
+                    todo_call_ids.add(str(call.get("id") or ""))
+                    continue
+                blocks.append(TranscriptBlock(
+                    kind="tool",
+                    tool_call_id=str(call.get("id") or ""),
+                    name=str(call.get("name") or "tool"),
+                    arguments=call.get("args") if isinstance(call.get("args"), dict) else {},
+                    status="running",
+                ))
+        elif isinstance(message, ToolMessage):
+            content = message_text(message)
+            tool_call_id = str(getattr(message, "tool_call_id", "") or "")
+            if getattr(message, "name", None) == "write_todos" or tool_call_id in todo_call_ids:
+                continue
+            is_error = tool_message_is_error(message)
+            artifact = getattr(message, "artifact", None)
+            code = (
+                artifact.get("exit_code")
+                if isinstance(artifact, dict) and getattr(message, "name", None) == "execute"
+                else None
+            )
+            exit_code = code if isinstance(code, int) and not isinstance(code, bool) else None
+            updated = False
+            for block in reversed(blocks):
+                if block.kind == "tool" and block.tool_call_id == tool_call_id:
+                    block.content = content
+                    block.is_error = is_error
+                    block.status = "error" if is_error else "completed"
+                    block.exit_code = exit_code
+                    block.artifact = artifact
+                    updated = True
+                    break
+            if not updated:
+                blocks.append(TranscriptBlock(
+                    kind="tool",
+                    tool_call_id=tool_call_id,
+                    name=str(getattr(message, "name", "") or "tool"),
+                    content=content,
+                    is_error=is_error,
+                    status="error" if is_error else "completed",
+                    exit_code=exit_code,
+                    artifact=artifact,
+                ))
+    return blocks
+
+
+def settle_restored_tools(
+    blocks: list[TranscriptBlock], *, waiting_ids: set[str] | None = None,
+    waiting_human: bool = False,
+) -> None:
+    """Freeze tool calls without results when showing a saved checkpoint."""
+    pending = waiting_ids or set()
+    for block in blocks:
+        if block.kind != "tool" or block.status != "running":
+            continue
+        if block.tool_call_id in pending or (waiting_human and block.name == "request_human_input"):
+            block.status = "waiting"
+        else:
+            block.status = "interrupted"

agent/session_lock.py

@@ -0,0 +1,223 @@
+"""Exclusive thread ownership for one session database.
+
+A lease is an advisory ``flock`` on ``<database>.locks/<sha256(thread_id)>.lock``
+so two processes can never write the same persisted thread at once. Lock files
+are created once and never deleted; releasing only drops the kernel hold. The
+locks directory and files stay private to the database owner and every file
+descriptor is close-on-exec so agent child processes cannot inherit a hold.
+
+Runners without a persisted store share nothing across processes; for those,
+``InProcessSessionLockManager`` guards threads that flow through one
+checkpointer object within this process.
+"""
+from __future__ import annotations
+
+import fcntl
+import hashlib
+import os
+from pathlib import Path
+import threading
+import time
+from typing import Any, Callable
+from weakref import WeakKeyDictionary, ref
+
+WAIT_POLL_SECONDS = 0.05
+
+
+class SessionLockBusyError(RuntimeError):
+    """The thread is already owned by another runner or process."""
+
+
+class SessionLease:
+    """Ownership of one persisted thread; ``release`` is idempotent."""
+
+    def __init__(
+        self,
+        thread_id: str,
+        handle: Any,
+        on_release: Callable[["SessionLease"], None],
+    ) -> None:
+        self._thread_id = thread_id
+        self._handle = handle
+        self._on_release = on_release
+        self._released = False
+
+    @property
+    def thread_id(self) -> str:
+        return self._thread_id
+
+    @property
+    def released(self) -> bool:
+        return self._released
+
+    def release(self) -> None:
+        if self._released:
+            return
+        self._released = True
+        try:
+            fcntl.flock(self._handle.fileno(), fcntl.LOCK_UN)
+        finally:
+            self._handle.close()
+        self._on_release(self)
+
+    def fileno(self) -> int:
+        return self._handle.fileno()
+
+
+class SessionLockManager:
+    """flock-backed thread ownership beside one session database."""
+
+    def __init__(self, database: Path) -> None:
+        self._root = database.parent / f"{database.name}.locks"
+        self._root.mkdir(parents=True, mode=0o700, exist_ok=True)
+        try:
+            os.chmod(self._root, 0o700)
+        except OSError:
+            pass
+        self._mutex = threading.Lock()
+        self._leases: dict[str, "ref[SessionLease]"] = {}
+
+    def lock_path(self, thread_id: str) -> Path:
+        digest = hashlib.sha256(thread_id.encode("utf-8")).hexdigest()
+        return self._root / f"{digest}.lock"
+
+    def try_acquire(self, thread_id: str) -> SessionLease | None:
+        path = self.lock_path(thread_id)
+        fd = os.open(path, os.O_RDWR | os.O_CREAT | os.O_CLOEXEC, 0o600)
+        handle = os.fdopen(fd, "r+b")
+        try:
+            fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+        except BlockingIOError:
+            handle.close()
+            return None
+        try:
+            os.chmod(path, 0o600)
+        except OSError:
+            pass
+        lease = SessionLease(thread_id, handle, self._forget)
+        with self._mutex:
+            self._leases[thread_id] = ref(lease)
+        return lease
+
+    def wait_acquire(
+        self,
+        thread_id: str,
+        *,
+        cancelled: Callable[[], bool] | None = None,
+        poll_seconds: float = WAIT_POLL_SECONDS,
+    ) -> SessionLease | None:
+        while True:
+            if cancelled is not None and cancelled():
+                return None
+            lease = self.try_acquire(thread_id)
+            if lease is not None:
+                # Cancellation arriving between the check and the acquisition
+                # must still win: hand the lock straight back.
+                if cancelled is not None and cancelled():
+                    lease.release()
+                    return None
+                return lease
+            time.sleep(poll_seconds)
+
+    def close(self) -> None:
+        """Release every lease this manager still tracks; idempotent."""
+        with self._mutex:
+            leases = [item() for item in self._leases.values()]
+            self._leases.clear()
+        for lease in leases:
+            if lease is not None:
+                lease.release()
+
+    def _forget(self, lease: SessionLease) -> None:
+        with self._mutex:
+            held = self._leases.get(lease.thread_id)
+            if held is not None and held() is lease:
+                del self._leases[lease.thread_id]
+
+
+class InProcessSessionLease:
+    """Ownership of one in-memory thread; ``release`` is idempotent."""
+
+    def __init__(self, thread_id: str, owners: dict[str, "ref[InProcessSessionLease]"]) -> None:
+        self._thread_id = thread_id
+        self._owners = owners
+        self._released = False
+
+    @property
+    def thread_id(self) -> str:
+        return self._thread_id
+
+    @property
+    def released(self) -> bool:
+        return self._released
+
+    def release(self) -> None:
+        if self._released:
+            return
+        self._released = True
+        with _IN_PROCESS_GUARD:
+            held = self._owners.get(self._thread_id)
+            if held is not None and held() is self:
+                del self._owners[self._thread_id]
+
+
+_IN_PROCESS_GUARD = threading.Lock()
+_IN_PROCESS_OWNERS: "WeakKeyDictionary[Any, dict[str, ref[InProcessSessionLease]]]" = WeakKeyDictionary()
+
+
+class InProcessSessionLockManager:
+    """Guards threads shared through one checkpointer object within this process."""
+
+    def __init__(self, checkpointer: Any) -> None:
+        self._checkpointer = checkpointer
+
+    def try_acquire(self, thread_id: str) -> InProcessSessionLease | None:
+        if self._checkpointer is None:
+            # Nothing is shared when the graph has no checkpointer at all.
+            return InProcessSessionLease(thread_id, {})
+        with _IN_PROCESS_GUARD:
+            owners = _IN_PROCESS_OWNERS.get(self._checkpointer)
+            if owners is None:
+                owners = {}
+                _IN_PROCESS_OWNERS[self._checkpointer] = owners
+            held = owners.get(thread_id)
+            if held is not None:
+                lease = held()
+                if lease is not None and not lease.released:
+                    return None
+            lease = InProcessSessionLease(thread_id, owners)
+            owners[thread_id] = ref(lease)
+            return lease
+
+    def wait_acquire(
+        self,
+        thread_id: str,
+        *,
+        cancelled: Callable[[], bool] | None = None,
+        poll_seconds: float = WAIT_POLL_SECONDS,
+    ) -> InProcessSessionLease | None:
+        while True:
+            if cancelled is not None and cancelled():
+                return None
+            lease = self.try_acquire(thread_id)
+            if lease is not None:
+                # Cancellation arriving between the check and the acquisition
+                # must still win: hand the lock straight back.
+                if cancelled is not None and cancelled():
+                    lease.release()
+                    return None
+                return lease
+            time.sleep(poll_seconds)
+
+    def close(self) -> None:
+        if self._checkpointer is None:
+            return
+        with _IN_PROCESS_GUARD:
+            owners = _IN_PROCESS_OWNERS.get(self._checkpointer)
+            if owners is None:
+                return
+            leases = [item() for item in owners.values()]
+            owners.clear()
+        for lease in leases:
+            if lease is not None:
+                lease.release()