aline-ai 0.5.9__py3-none-any.whl → 0.5.10__py3-none-any.whl

realign/db/sqlite_db.py

@@ -20,6 +20,8 @@ from .base import (
     TurnRecord,
     EventRecord,
     LockRecord,
+    AgentRecord,
+    AgentContextRecord,
 )
 from .schema import (
     INIT_SCRIPTS,
@@ -2024,6 +2026,572 @@ class SQLiteDatabase(DatabaseInterface):
             self._connection.close()
             self._connection = None
 
+    # -------------------------------------------------------------------------
+    # Agent methods (Schema V15 - replaces terminal.json)
+    # -------------------------------------------------------------------------
+
+    def get_or_create_agent(
+        self,
+        agent_id: str,
+        provider: str,
+        session_type: str,
+        *,
+        session_id: Optional[str] = None,
+        context_id: Optional[str] = None,
+        transcript_path: Optional[str] = None,
+        cwd: Optional[str] = None,
+        project_dir: Optional[str] = None,
+        source: Optional[str] = None,
+        status: str = "active",
+        attention: Optional[str] = None,
+    ) -> AgentRecord:
+        """Get existing agent or create new one."""
+        conn = self._get_connection()
+        cursor = conn.cursor()
+
+        try:
+            cursor.execute("SELECT * FROM agents WHERE id = ?", (agent_id,))
+            row = cursor.fetchone()
+            if row:
+                return self._row_to_agent(row)
+        except sqlite3.OperationalError:
+            # Table may not exist in older schema
+            raise
+
+        # Get user identity from config
+        try:
+            from ..config import ReAlignConfig
+
+            config = ReAlignConfig.load()
+            creator_name = config.user_name
+            creator_id = config.user_id
+        except Exception:
+            creator_name = None
+            creator_id = None
+
+        now = datetime.now()
+        cursor.execute(
+            """
+            INSERT INTO agents (
+                id, provider, session_type, session_id, context_id,
+                transcript_path, cwd, project_dir, status, attention, source,
+                created_at, updated_at, creator_name, creator_id
+            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+            """,
+            (
+                agent_id,
+                provider,
+                session_type,
+                session_id,
+                context_id,
+                transcript_path,
+                cwd,
+                project_dir,
+                status,
+                attention,
+                source,
+                now,
+                now,
+                creator_name,
+                creator_id,
+            ),
+        )
+        conn.commit()
+
+        return AgentRecord(
+            id=agent_id,
+            provider=provider,
+            session_type=session_type,
+            session_id=session_id,
+            context_id=context_id,
+            transcript_path=transcript_path,
+            cwd=cwd,
+            project_dir=project_dir,
+            status=status,
+            attention=attention,
+            source=source,
+            created_at=now,
+            updated_at=now,
+            creator_name=creator_name,
+            creator_id=creator_id,
+        )
+
+    def update_agent(
+        self,
+        agent_id: str,
+        *,
+        provider: Optional[str] = None,
+        session_type: Optional[str] = None,
+        session_id: Optional[str] = None,
+        context_id: Optional[str] = None,
+        transcript_path: Optional[str] = None,
+        cwd: Optional[str] = None,
+        project_dir: Optional[str] = None,
+        status: Optional[str] = None,
+        attention: Optional[str] = None,
+        source: Optional[str] = None,
+    ) -> Optional[AgentRecord]:
+        """Update an existing agent. Returns the updated record or None if not found."""
+        conn = self._get_connection()
+        cursor = conn.cursor()
+
+        # Build dynamic UPDATE statement
+        updates: List[str] = []
+        params: List[Any] = []
+
+        if provider is not None:
+            updates.append("provider = ?")
+            params.append(provider)
+        if session_type is not None:
+            updates.append("session_type = ?")
+            params.append(session_type)
+        if session_id is not None:
+            updates.append("session_id = ?")
+            params.append(session_id)
+        if context_id is not None:
+            updates.append("context_id = ?")
+            params.append(context_id)
+        if transcript_path is not None:
+            updates.append("transcript_path = ?")
+            params.append(transcript_path)
+        if cwd is not None:
+            updates.append("cwd = ?")
+            params.append(cwd)
+        if project_dir is not None:
+            updates.append("project_dir = ?")
+            params.append(project_dir)
+        if status is not None:
+            updates.append("status = ?")
+            params.append(status)
+        if attention is not None:
+            updates.append("attention = ?")
+            params.append(attention if attention != "" else None)
+        if source is not None:
+            updates.append("source = ?")
+            params.append(source)
+
+        if not updates:
+            # Nothing to update, just return existing
+            return self.get_agent_by_id(agent_id)
+
+        updates.append("updated_at = datetime('now')")
+        params.append(agent_id)
+
+        try:
+            cursor.execute(
+                f"UPDATE agents SET {', '.join(updates)} WHERE id = ?",
+                params,
+            )
+            conn.commit()
+            return self.get_agent_by_id(agent_id)
+        except sqlite3.OperationalError:
+            conn.rollback()
+            return None
+
+    def get_agent_by_id(self, agent_id: str) -> Optional[AgentRecord]:
+        """Get an agent by its ID."""
+        conn = self._get_connection()
+        cursor = conn.cursor()
+        try:
+            cursor.execute("SELECT * FROM agents WHERE id = ?", (agent_id,))
+            row = cursor.fetchone()
+            if row:
+                return self._row_to_agent(row)
+        except sqlite3.OperationalError:
+            pass
+        return None
+
+    def list_agents(
+        self,
+        *,
+        status: Optional[str] = None,
+        context_id: Optional[str] = None,
+        limit: int = 100,
+    ) -> List[AgentRecord]:
+        """List agents, optionally filtered by status or context."""
+        conn = self._get_connection()
+        cursor = conn.cursor()
+
+        where_clauses: List[str] = []
+        params: List[Any] = []
+
+        if status:
+            where_clauses.append("status = ?")
+            params.append(status)
+        if context_id:
+            where_clauses.append("context_id = ?")
+            params.append(context_id)
+
+        where_sql = ""
+        if where_clauses:
+            where_sql = "WHERE " + " AND ".join(where_clauses)
+
+        params.append(limit)
+
+        try:
+            cursor.execute(
+                f"""
+                SELECT * FROM agents
+                {where_sql}
+                ORDER BY updated_at DESC
+                LIMIT ?
+                """,
+                params,
+            )
+            return [self._row_to_agent(row) for row in cursor.fetchall()]
+        except sqlite3.OperationalError:
+            return []
+
+    def delete_agent(self, agent_id: str) -> bool:
+        """Delete an agent by ID."""
+        conn = self._get_connection()
+        try:
+            cursor = conn.execute("DELETE FROM agents WHERE id = ?", (agent_id,))
+            conn.commit()
+            return cursor.rowcount > 0
+        except sqlite3.OperationalError:
+            conn.rollback()
+            return False
+
+    # -------------------------------------------------------------------------
+    # Agent Context methods (Schema V15 - replaces load.json)
+    # -------------------------------------------------------------------------
+
+    def get_or_create_agent_context(
+        self,
+        context_id: str,
+        *,
+        workspace: Optional[str] = None,
+        loaded_at: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> AgentContextRecord:
+        """Get existing agent context or create new one."""
+        conn = self._get_connection()
+        cursor = conn.cursor()
+
+        try:
+            cursor.execute("SELECT * FROM agent_contexts WHERE id = ?", (context_id,))
+            row = cursor.fetchone()
+            if row:
+                return self._row_to_agent_context(row)
+        except sqlite3.OperationalError:
+            raise
+
+        now = datetime.now()
+        metadata_json = json.dumps(metadata or {}, ensure_ascii=False)
+
+        cursor.execute(
+            """
+            INSERT INTO agent_contexts (id, workspace, loaded_at, created_at, updated_at, metadata)
+            VALUES (?, ?, ?, ?, ?, ?)
+            """,
+            (context_id, workspace, loaded_at, now, now, metadata_json),
+        )
+        conn.commit()
+
+        return AgentContextRecord(
+            id=context_id,
+            workspace=workspace,
+            loaded_at=loaded_at,
+            created_at=now,
+            updated_at=now,
+            metadata=metadata,
+            session_ids=[],
+            event_ids=[],
+        )
+
+    def update_agent_context(
+        self,
+        context_id: str,
+        *,
+        workspace: Optional[str] = None,
+        loaded_at: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> Optional[AgentContextRecord]:
+        """Update an existing agent context."""
+        conn = self._get_connection()
+        cursor = conn.cursor()
+
+        updates: List[str] = []
+        params: List[Any] = []
+
+        if workspace is not None:
+            updates.append("workspace = ?")
+            params.append(workspace)
+        if loaded_at is not None:
+            updates.append("loaded_at = ?")
+            params.append(loaded_at)
+        if metadata is not None:
+            updates.append("metadata = ?")
+            params.append(json.dumps(metadata, ensure_ascii=False))
+
+        if not updates:
+            return self.get_agent_context_by_id(context_id)
+
+        updates.append("updated_at = datetime('now')")
+        params.append(context_id)
+
+        try:
+            cursor.execute(
+                f"UPDATE agent_contexts SET {', '.join(updates)} WHERE id = ?",
+                params,
+            )
+            conn.commit()
+            return self.get_agent_context_by_id(context_id)
+        except sqlite3.OperationalError:
+            conn.rollback()
+            return None
+
+    def get_agent_context_by_id(self, context_id: str) -> Optional[AgentContextRecord]:
+        """Get an agent context by its ID, including linked sessions and events."""
+        conn = self._get_connection()
+        cursor = conn.cursor()
+
+        try:
+            cursor.execute("SELECT * FROM agent_contexts WHERE id = ?", (context_id,))
+            row = cursor.fetchone()
+            if not row:
+                return None
+
+            record = self._row_to_agent_context(row)
+
+            # Fetch linked session IDs
+            cursor.execute(
+                "SELECT session_id FROM agent_context_sessions WHERE context_id = ?",
+                (context_id,),
+            )
+            record.session_ids = [r[0] for r in cursor.fetchall()]
+
+            # Fetch linked event IDs
+            cursor.execute(
+                "SELECT event_id FROM agent_context_events WHERE context_id = ?",
+                (context_id,),
+            )
+            record.event_ids = [r[0] for r in cursor.fetchall()]
+
+            return record
+        except sqlite3.OperationalError:
+            return None
+
+    def get_agent_context_by_workspace(
+        self, workspace: str
+    ) -> Optional[AgentContextRecord]:
+        """Get an agent context by workspace path."""
+        conn = self._get_connection()
+        cursor = conn.cursor()
+
+        try:
+            cursor.execute(
+                "SELECT * FROM agent_contexts WHERE workspace = ?", (workspace,)
+            )
+            row = cursor.fetchone()
+            if not row:
+                return None
+
+            record = self._row_to_agent_context(row)
+            context_id = record.id
+
+            # Fetch linked session IDs
+            cursor.execute(
+                "SELECT session_id FROM agent_context_sessions WHERE context_id = ?",
+                (context_id,),
+            )
+            record.session_ids = [r[0] for r in cursor.fetchall()]
+
+            # Fetch linked event IDs
+            cursor.execute(
+                "SELECT event_id FROM agent_context_events WHERE context_id = ?",
+                (context_id,),
+            )
+            record.event_ids = [r[0] for r in cursor.fetchall()]
+
+            return record
+        except sqlite3.OperationalError:
+            return None
+
+    def list_agent_contexts(self, limit: int = 100) -> List[AgentContextRecord]:
+        """List all agent contexts."""
+        conn = self._get_connection()
+        cursor = conn.cursor()
+
+        try:
+            cursor.execute(
+                """
+                SELECT * FROM agent_contexts
+                ORDER BY updated_at DESC
+                LIMIT ?
+                """,
+                (limit,),
+            )
+            contexts = []
+            for row in cursor.fetchall():
+                record = self._row_to_agent_context(row)
+                context_id = record.id
+
+                # Fetch linked session IDs
+                cursor.execute(
+                    "SELECT session_id FROM agent_context_sessions WHERE context_id = ?",
+                    (context_id,),
+                )
+                record.session_ids = [r[0] for r in cursor.fetchall()]
+
+                # Fetch linked event IDs
+                cursor.execute(
+                    "SELECT event_id FROM agent_context_events WHERE context_id = ?",
+                    (context_id,),
+                )
+                record.event_ids = [r[0] for r in cursor.fetchall()]
+
+                contexts.append(record)
+            return contexts
+        except sqlite3.OperationalError:
+            return []
+
+    def delete_agent_context(self, context_id: str) -> bool:
+        """Delete an agent context and all its links."""
+        conn = self._get_connection()
+        try:
+            conn.execute(
+                "DELETE FROM agent_context_sessions WHERE context_id = ?", (context_id,)
+            )
+            conn.execute(
+                "DELETE FROM agent_context_events WHERE context_id = ?", (context_id,)
+            )
+            cursor = conn.execute(
+                "DELETE FROM agent_contexts WHERE id = ?", (context_id,)
+            )
+            conn.commit()
+            return cursor.rowcount > 0
+        except sqlite3.OperationalError:
+            conn.rollback()
+            return False
+
+    def link_session_to_agent_context(self, context_id: str, session_id: str) -> bool:
+        """Link a session to an agent context.
+
+        Silently skips if session doesn't exist (FK constraint).
+        """
+        conn = self._get_connection()
+        try:
+            conn.execute(
+                """
+                INSERT OR IGNORE INTO agent_context_sessions (context_id, session_id)
+                VALUES (?, ?)
+                """,
+                (context_id, session_id),
+            )
+            conn.commit()
+            return True
+        except (sqlite3.OperationalError, sqlite3.IntegrityError):
+            conn.rollback()
+            return False
+
+    def unlink_session_from_agent_context(
+        self, context_id: str, session_id: str
+    ) -> bool:
+        """Unlink a session from an agent context."""
+        conn = self._get_connection()
+        try:
+            conn.execute(
+                "DELETE FROM agent_context_sessions WHERE context_id = ? AND session_id = ?",
+                (context_id, session_id),
+            )
+            conn.commit()
+            return True
+        except (sqlite3.OperationalError, sqlite3.IntegrityError):
+            conn.rollback()
+            return False
+
+    def link_event_to_agent_context(self, context_id: str, event_id: str) -> bool:
+        """Link an event to an agent context.
+
+        Silently skips if event doesn't exist (FK constraint).
+        """
+        conn = self._get_connection()
+        try:
+            conn.execute(
+                """
+                INSERT OR IGNORE INTO agent_context_events (context_id, event_id)
+                VALUES (?, ?)
+                """,
+                (context_id, event_id),
+            )
+            conn.commit()
+            return True
+        except (sqlite3.OperationalError, sqlite3.IntegrityError):
+            conn.rollback()
+            return False
+
+    def unlink_event_from_agent_context(self, context_id: str, event_id: str) -> bool:
+        """Unlink an event from an agent context."""
+        conn = self._get_connection()
+        try:
+            conn.execute(
+                "DELETE FROM agent_context_events WHERE context_id = ? AND event_id = ?",
+                (context_id, event_id),
+            )
+            conn.commit()
+            return True
+        except (sqlite3.OperationalError, sqlite3.IntegrityError):
+            conn.rollback()
+            return False
+
+    def set_agent_context_sessions(
+        self, context_id: str, session_ids: List[str]
+    ) -> bool:
+        """Replace all sessions linked to a context.
+
+        Silently skips sessions that don't exist in DB (FK constraint).
+        """
+        conn = self._get_connection()
+        try:
+            conn.execute(
+                "DELETE FROM agent_context_sessions WHERE context_id = ?", (context_id,)
+            )
+            if session_ids:
+                # Insert one at a time to skip FK failures
+                for sid in session_ids:
+                    try:
+                        conn.execute(
+                            "INSERT OR IGNORE INTO agent_context_sessions (context_id, session_id) VALUES (?, ?)",
+                            (context_id, sid),
+                        )
+                    except sqlite3.IntegrityError:
+                        # Session doesn't exist in DB, skip
+                        pass
+            conn.commit()
+            return True
+        except (sqlite3.OperationalError, sqlite3.IntegrityError):
+            conn.rollback()
+            return False
+
+    def set_agent_context_events(self, context_id: str, event_ids: List[str]) -> bool:
+        """Replace all events linked to a context.
+
+        Silently skips events that don't exist in DB (FK constraint).
+        """
+        conn = self._get_connection()
+        try:
+            conn.execute(
+                "DELETE FROM agent_context_events WHERE context_id = ?", (context_id,)
+            )
+            if event_ids:
+                # Insert one at a time to skip FK failures
+                for eid in event_ids:
+                    try:
+                        conn.execute(
+                            "INSERT OR IGNORE INTO agent_context_events (context_id, event_id) VALUES (?, ?)",
+                            (context_id, eid),
+                        )
+                    except sqlite3.IntegrityError:
+                        # Event doesn't exist in DB, skip
+                        pass
+            conn.commit()
+            return True
+        except (sqlite3.OperationalError, sqlite3.IntegrityError):
+            conn.rollback()
+            return False
+
     # Helper methods for row mapping
     def _row_to_project(self, row: sqlite3.Row) -> ProjectRecord:
         return ProjectRecord(
@@ -2255,3 +2823,52 @@ class SQLiteDatabase(DatabaseInterface):
             creator_name=creator_name,
             creator_id=creator_id,
         )
+
+    def _row_to_agent(self, row: sqlite3.Row) -> AgentRecord:
+        """Convert a database row to an AgentRecord."""
+        creator_name = None
+        creator_id = None
+        try:
+            creator_name = row["creator_name"]
+            creator_id = row["creator_id"]
+        except (IndexError, KeyError):
+            pass
+
+        return AgentRecord(
+            id=row["id"],
+            provider=row["provider"],
+            session_type=row["session_type"],
+            session_id=row["session_id"],
+            context_id=row["context_id"],
+            transcript_path=row["transcript_path"],
+            cwd=row["cwd"],
+            project_dir=row["project_dir"],
+            status=row["status"] or "active",
+            attention=row["attention"],
+            source=row["source"],
+            created_at=self._parse_datetime(row["created_at"]),
+            updated_at=self._parse_datetime(row["updated_at"]),
+            creator_name=creator_name,
+            creator_id=creator_id,
+        )
+
+    def _row_to_agent_context(self, row: sqlite3.Row) -> AgentContextRecord:
+        """Convert a database row to an AgentContextRecord."""
+        metadata = None
+        try:
+            metadata_raw = row["metadata"]
+            if metadata_raw:
+                metadata = json.loads(metadata_raw)
+        except (IndexError, KeyError, json.JSONDecodeError):
+            pass
+
+        return AgentContextRecord(
+            id=row["id"],
+            workspace=row["workspace"],
+            loaded_at=row["loaded_at"],
+            created_at=self._parse_datetime(row["created_at"]),
+            updated_at=self._parse_datetime(row["updated_at"]),
+            metadata=metadata,
+            session_ids=None,  # Populated separately
+            event_ids=None,  # Populated separately
+        )