gdmcode 0.1.0__py3-none-any.whl

src/memory/session_store.py

@@ -0,0 +1,201 @@
+"""SessionStore — persists and resumes conversation sessions via GdmDatabase.
+
+Handles all session lifecycle: create, resume, save checkpoint, close.
+Works in tandem with TranscriptStore (in-memory) — SessionStore is the
+durable layer.
+"""
+from __future__ import annotations
+
+import json
+import logging
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from src.agent.transcript import Turn, TranscriptStore
+from src.exceptions import DatabaseError
+
+if TYPE_CHECKING:
+    from src.memory.db import GdmDatabase
+
+__all__ = ["Session", "SessionStore"]
+
+log = logging.getLogger(__name__)
+
+
+@dataclass
+class Session:
+    """Lightweight descriptor for a persisted session."""
+
+    session_id: str
+    project_id: str
+    created_at: datetime
+    updated_at: datetime
+    turn_count: int
+    cost_usd: float
+
+
+class SessionStore:
+    """Creates, persists, and resumes gdm sessions.
+
+    Usage::
+
+        store = SessionStore(db, project_root=Path.cwd())
+        session_id = store.create()
+        transcript = store.load_transcript(session_id, max_tokens=120_000)
+        # ... agent loop ...
+        store.save_transcript(session_id, transcript)
+        store.close(session_id, cost_usd=0.04)
+    """
+
+    def __init__(self, db: GdmDatabase, project_root: Path) -> None:
+        self._db = db
+        self._project_root = project_root
+        self._project_id = self._ensure_project()
+
+    # ------------------------------------------------------------------
+    # Project helpers
+    # ------------------------------------------------------------------
+
+    def _ensure_project(self) -> str:
+        """Return existing project_id or create a new row."""
+        root_str = str(self._project_root)
+        row = self._db.execute_one(
+            "SELECT project_id FROM projects WHERE root_path = ?", (root_str,)
+        )
+        if row:
+            return row["project_id"]
+
+        import uuid
+
+        pid = str(uuid.uuid4())
+        self._db.execute(
+            "INSERT INTO projects (project_id, root_path, name) VALUES (?, ?, ?)",
+            (pid, root_str, self._project_root.name),
+        )
+        return pid
+
+    # ------------------------------------------------------------------
+    # Session lifecycle
+    # ------------------------------------------------------------------
+
+    def create(self) -> str:
+        """Create a new session row. Returns session_id (UUID)."""
+        import uuid
+
+        sid = str(uuid.uuid4())
+        now = datetime.now(timezone.utc).isoformat()
+        try:
+            self._db.execute(
+                """INSERT INTO sessions
+                       (session_id, project_id, created_at, updated_at, turn_count, cost_usd)
+                   VALUES (?, ?, ?, ?, 0, 0.0)""",
+                (sid, self._project_id, now, now),
+            )
+        except Exception as exc:
+            raise DatabaseError(f"Failed to create session: {exc}") from exc
+        log.debug("Created session %s", sid)
+        return sid
+
+    def load_transcript(self, session_id: str, max_tokens: int = 120_000) -> TranscriptStore:
+        """Load a persisted transcript from the database.
+
+        Returns an empty TranscriptStore if no turns exist.
+        """
+        rows = self._db.execute_all(
+            "SELECT role, content, tokens, tool_name, tool_call_id, tool_calls_json "
+            "FROM memory WHERE session_id = ? ORDER BY id ASC",
+            (session_id,),
+        )
+        turns: list[Turn] = []
+        for row in rows:
+            tool_calls = None
+            if row["tool_calls_json"]:
+                try:
+                    tool_calls = json.loads(row["tool_calls_json"])
+                except (json.JSONDecodeError, TypeError):
+                    log.warning("Corrupt tool_calls_json for session %s — skipping", session_id)
+            turns.append(
+                Turn(
+                    role=row["role"],
+                    content=row["content"],
+                    tokens=row["tokens"] or 0,
+                    tool_name=row["tool_name"],
+                    tool_call_id=row["tool_call_id"],
+                    tool_calls=tool_calls,
+                )
+            )
+        return TranscriptStore.from_turns(turns, max_tokens=max_tokens)
+
+    def save_transcript(self, session_id: str, store: TranscriptStore) -> None:
+        """Overwrite the persisted transcript with the current in-memory state.
+
+        Routes through memory_save_turns() so turn_index is assigned correctly
+        and the UNIQUE(session_id, turn_index) constraint is satisfied on every call.
+        """
+        now = datetime.now(timezone.utc).isoformat()
+        try:
+            turns = [
+                {
+                    "role": t.role,
+                    "content": t.content,
+                    "tokens": t.tokens,
+                    "tool_name": t.tool_name,
+                    "tool_call_id": t.tool_call_id,
+                    "tool_calls": t.tool_calls,
+                }
+                for t in store
+            ]
+            self._db.memory_save_turns(session_id, turns)
+            self._db.execute(
+                "UPDATE sessions SET updated_at = ?, turn_count = ? WHERE session_id = ?",
+                (now, store.turn_count, session_id),
+            )
+        except Exception as exc:
+            raise DatabaseError(f"Failed to save transcript for {session_id}: {exc}") from exc
+
+    def close(self, session_id: str, cost_tracker: object = None) -> None:
+        """Mark session closed. Delegates final cost write to cost_tracker if provided.
+
+        Args:
+            session_id: The session to close.
+            cost_tracker: CostTracker instance — calls flush_to_db() to persist
+                          the authoritative total. If None, cost_usd is left as-is.
+        """
+        now = datetime.now(timezone.utc).isoformat()
+        self._db.execute(
+            "UPDATE sessions SET updated_at = ? WHERE session_id = ?",
+            (now, session_id),
+        )
+        if cost_tracker is not None:
+            cost_tracker.flush_to_db(self._db)  # type: ignore[union-attr]
+            log.info(
+                "Session %s closed. Total cost: $%.4f",
+                session_id,
+                getattr(cost_tracker, "session_total_usd", 0.0),
+            )
+        else:
+            log.info("Session %s closed.", session_id)
+
+    def list_recent(self, limit: int = 10) -> list[Session]:
+        """Return the N most recent sessions for this project."""
+        rows = self._db.execute_all(
+            """SELECT session_id, project_id, created_at, updated_at, turn_count, cost_usd
+               FROM sessions WHERE project_id = ?
+               ORDER BY updated_at DESC LIMIT ?""",
+            (self._project_id, limit),
+        )
+        sessions: list[Session] = []
+        for row in rows:
+            sessions.append(
+                Session(
+                    session_id=row["session_id"],
+                    project_id=row["project_id"],
+                    created_at=datetime.fromisoformat(row["created_at"]),
+                    updated_at=datetime.fromisoformat(row["updated_at"]),
+                    turn_count=row["turn_count"],
+                    cost_usd=row["cost_usd"],
+                )
+            )
+        return sessions