openai-agents 0.2.8__py3-none-any.whl → 0.6.8__py3-none-any.whl

agents/extensions/memory/advanced_sqlite_session.py

@@ -0,0 +1,1285 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import threading
+from contextlib import closing
+from pathlib import Path
+from typing import Any, Union, cast
+
+from agents.result import RunResult
+from agents.usage import Usage
+
+from ...items import TResponseInputItem
+from ...memory import SQLiteSession
+
+
+class AdvancedSQLiteSession(SQLiteSession):
+    """Enhanced SQLite session with conversation branching and usage analytics."""
+
+    def __init__(
+        self,
+        *,
+        session_id: str,
+        db_path: str | Path = ":memory:",
+        create_tables: bool = False,
+        logger: logging.Logger | None = None,
+        **kwargs,
+    ):
+        """Initialize the AdvancedSQLiteSession.
+
+        Args:
+            session_id: The ID of the session
+            db_path: The path to the SQLite database file. Defaults to `:memory:` for in-memory storage
+            create_tables: Whether to create the structure tables
+            logger: The logger to use. Defaults to the module logger
+            **kwargs: Additional keyword arguments to pass to the superclass
+        """  # noqa: E501
+        super().__init__(session_id, db_path, **kwargs)
+        if create_tables:
+            self._init_structure_tables()
+        self._current_branch_id = "main"
+        self._logger = logger or logging.getLogger(__name__)
+
+    def _init_structure_tables(self):
+        """Add structure and usage tracking tables.
+
+        Creates the message_structure and turn_usage tables with appropriate
+        indexes for conversation branching and usage analytics.
+        """
+        conn = self._get_connection()
+
+        # Message structure with branch support
+        conn.execute("""
+            CREATE TABLE IF NOT EXISTS message_structure (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                session_id TEXT NOT NULL,
+                message_id INTEGER NOT NULL,
+                branch_id TEXT NOT NULL DEFAULT 'main',
+                message_type TEXT NOT NULL,
+                sequence_number INTEGER NOT NULL,
+                user_turn_number INTEGER,
+                branch_turn_number INTEGER,
+                tool_name TEXT,
+                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                FOREIGN KEY (session_id) REFERENCES agent_sessions(session_id) ON DELETE CASCADE,
+                FOREIGN KEY (message_id) REFERENCES agent_messages(id) ON DELETE CASCADE
+            )
+        """)
+
+        # Turn-level usage tracking with branch support and full JSON details
+        conn.execute("""
+            CREATE TABLE IF NOT EXISTS turn_usage (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                session_id TEXT NOT NULL,
+                branch_id TEXT NOT NULL DEFAULT 'main',
+                user_turn_number INTEGER NOT NULL,
+                requests INTEGER DEFAULT 0,
+                input_tokens INTEGER DEFAULT 0,
+                output_tokens INTEGER DEFAULT 0,
+                total_tokens INTEGER DEFAULT 0,
+                input_tokens_details JSON,
+                output_tokens_details JSON,
+                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                FOREIGN KEY (session_id) REFERENCES agent_sessions(session_id) ON DELETE CASCADE,
+                UNIQUE(session_id, branch_id, user_turn_number)
+            )
+        """)
+
+        # Indexes
+        conn.execute("""
+            CREATE INDEX IF NOT EXISTS idx_structure_session_seq
+            ON message_structure(session_id, sequence_number)
+        """)
+        conn.execute("""
+            CREATE INDEX IF NOT EXISTS idx_structure_branch
+            ON message_structure(session_id, branch_id)
+        """)
+        conn.execute("""
+            CREATE INDEX IF NOT EXISTS idx_structure_turn
+            ON message_structure(session_id, branch_id, user_turn_number)
+        """)
+        conn.execute("""
+            CREATE INDEX IF NOT EXISTS idx_structure_branch_seq
+            ON message_structure(session_id, branch_id, sequence_number)
+        """)
+        conn.execute("""
+            CREATE INDEX IF NOT EXISTS idx_turn_usage_session_turn
+            ON turn_usage(session_id, branch_id, user_turn_number)
+        """)
+
+        conn.commit()
+
+    async def add_items(self, items: list[TResponseInputItem]) -> None:
+        """Add items to the session.
+
+        Args:
+            items: The items to add to the session
+        """
+        # Add to base table first
+        await super().add_items(items)
+
+        # Extract structure metadata with precise sequencing
+        if items:
+            await self._add_structure_metadata(items)
+
+    async def get_items(
+        self,
+        limit: int | None = None,
+        branch_id: str | None = None,
+    ) -> list[TResponseInputItem]:
+        """Get items from current or specified branch.
+
+        Args:
+            limit: Maximum number of items to return. If None, returns all items.
+            branch_id: Branch to get items from. If None, uses current branch.
+
+        Returns:
+            List of conversation items from the specified branch.
+        """
+        if branch_id is None:
+            branch_id = self._current_branch_id
+
+            # Get all items for this branch
+            def _get_all_items_sync():
+                """Synchronous helper to get all items for a branch."""
+                conn = self._get_connection()
+                # TODO: Refactor SQLiteSession to use asyncio.Lock instead of threading.Lock and update this code  # noqa: E501
+                with self._lock if self._is_memory_db else threading.Lock():
+                    with closing(conn.cursor()) as cursor:
+                        if limit is None:
+                            cursor.execute(
+                                """
+                                SELECT m.message_data
+                                FROM agent_messages m
+                                JOIN message_structure s ON m.id = s.message_id
+                                WHERE m.session_id = ? AND s.branch_id = ?
+                                ORDER BY s.sequence_number ASC
+                            """,
+                                (self.session_id, branch_id),
+                            )
+                        else:
+                            cursor.execute(
+                                """
+                                SELECT m.message_data
+                                FROM agent_messages m
+                                JOIN message_structure s ON m.id = s.message_id
+                                WHERE m.session_id = ? AND s.branch_id = ?
+                                ORDER BY s.sequence_number DESC
+                                LIMIT ?
+                            """,
+                                (self.session_id, branch_id, limit),
+                            )
+
+                        rows = cursor.fetchall()
+                        if limit is not None:
+                            rows = list(reversed(rows))
+
+                    items = []
+                    for (message_data,) in rows:
+                        try:
+                            item = json.loads(message_data)
+                            items.append(item)
+                        except json.JSONDecodeError:
+                            continue
+                    return items
+
+            return await asyncio.to_thread(_get_all_items_sync)
+
+        def _get_items_sync():
+            """Synchronous helper to get items for a specific branch."""
+            conn = self._get_connection()
+            # TODO: Refactor SQLiteSession to use asyncio.Lock instead of threading.Lock and update this code  # noqa: E501
+            with self._lock if self._is_memory_db else threading.Lock():
+                with closing(conn.cursor()) as cursor:
+                    # Get message IDs in correct order for this branch
+                    if limit is None:
+                        cursor.execute(
+                            """
+                            SELECT m.message_data
+                            FROM agent_messages m
+                            JOIN message_structure s ON m.id = s.message_id
+                            WHERE m.session_id = ? AND s.branch_id = ?
+                            ORDER BY s.sequence_number ASC
+                        """,
+                            (self.session_id, branch_id),
+                        )
+                    else:
+                        cursor.execute(
+                            """
+                            SELECT m.message_data
+                            FROM agent_messages m
+                            JOIN message_structure s ON m.id = s.message_id
+                            WHERE m.session_id = ? AND s.branch_id = ?
+                            ORDER BY s.sequence_number DESC
+                            LIMIT ?
+                        """,
+                            (self.session_id, branch_id, limit),
+                        )
+
+                    rows = cursor.fetchall()
+                    if limit is not None:
+                        rows = list(reversed(rows))
+
+                items = []
+                for (message_data,) in rows:
+                    try:
+                        item = json.loads(message_data)
+                        items.append(item)
+                    except json.JSONDecodeError:
+                        continue
+                return items
+
+        return await asyncio.to_thread(_get_items_sync)
+
+    async def store_run_usage(self, result: RunResult) -> None:
+        """Store usage data for the current conversation turn.
+
+        This is designed to be called after `Runner.run()` completes.
+        Session-level usage can be aggregated from turn data when needed.
+
+        Args:
+            result: The result from the run
+        """
+        try:
+            if result.context_wrapper.usage is not None:
+                # Get the current turn number for this branch
+                current_turn = self._get_current_turn_number()
+                # Only update turn-level usage - session usage is aggregated on demand
+                await self._update_turn_usage_internal(current_turn, result.context_wrapper.usage)
+        except Exception as e:
+            self._logger.error(f"Failed to store usage for session {self.session_id}: {e}")
+
+    def _get_next_turn_number(self, branch_id: str) -> int:
+        """Get the next turn number for a specific branch.
+
+        Args:
+            branch_id: The branch ID to get the next turn number for.
+
+        Returns:
+            The next available turn number for the specified branch.
+        """
+        conn = self._get_connection()
+        with closing(conn.cursor()) as cursor:
+            cursor.execute(
+                """
+                SELECT COALESCE(MAX(user_turn_number), 0)
+                FROM message_structure
+                WHERE session_id = ? AND branch_id = ?
+            """,
+                (self.session_id, branch_id),
+            )
+            result = cursor.fetchone()
+            max_turn = result[0] if result else 0
+            return max_turn + 1
+
+    def _get_next_branch_turn_number(self, branch_id: str) -> int:
+        """Get the next branch turn number for a specific branch.
+
+        Args:
+            branch_id: The branch ID to get the next branch turn number for.
+
+        Returns:
+            The next available branch turn number for the specified branch.
+        """
+        conn = self._get_connection()
+        with closing(conn.cursor()) as cursor:
+            cursor.execute(
+                """
+                SELECT COALESCE(MAX(branch_turn_number), 0)
+                FROM message_structure
+                WHERE session_id = ? AND branch_id = ?
+            """,
+                (self.session_id, branch_id),
+            )
+            result = cursor.fetchone()
+            max_turn = result[0] if result else 0
+            return max_turn + 1
+
+    def _get_current_turn_number(self) -> int:
+        """Get the current turn number for the current branch.
+
+        Returns:
+            The current turn number for the active branch.
+        """
+        conn = self._get_connection()
+        with closing(conn.cursor()) as cursor:
+            cursor.execute(
+                """
+                SELECT COALESCE(MAX(user_turn_number), 0)
+                FROM message_structure
+                WHERE session_id = ? AND branch_id = ?
+                """,
+                (self.session_id, self._current_branch_id),
+            )
+            result = cursor.fetchone()
+            return result[0] if result else 0
+
+    async def _add_structure_metadata(self, items: list[TResponseInputItem]) -> None:
+        """Extract structure metadata with branch-aware turn tracking.
+
+        This method:
+        - Assigns turn numbers per branch (not globally)
+        - Assigns explicit sequence numbers for precise ordering
+        - Links messages to their database IDs for structure tracking
+        - Handles multiple user messages in a single batch correctly
+
+        Args:
+            items: The items to add to the session
+        """
+
+        def _add_structure_sync():
+            """Synchronous helper to add structure metadata to database."""
+            conn = self._get_connection()
+            # TODO: Refactor SQLiteSession to use asyncio.Lock instead of threading.Lock and update this code  # noqa: E501
+            with self._lock if self._is_memory_db else threading.Lock():
+                # Get the IDs of messages we just inserted, in order
+                with closing(conn.cursor()) as cursor:
+                    cursor.execute(
+                        f"SELECT id FROM {self.messages_table} "
+                        f"WHERE session_id = ? ORDER BY id DESC LIMIT ?",
+                        (self.session_id, len(items)),
+                    )
+                    message_ids = [row[0] for row in cursor.fetchall()]
+                    message_ids.reverse()  # Match order of items
+
+                # Get current max sequence number (global)
+                with closing(conn.cursor()) as cursor:
+                    cursor.execute(
+                        """
+                        SELECT COALESCE(MAX(sequence_number), 0)
+                        FROM message_structure
+                        WHERE session_id = ?
+                    """,
+                        (self.session_id,),
+                    )
+                    seq_start = cursor.fetchone()[0]
+
+                # Get current turn numbers atomically with a single query
+                with closing(conn.cursor()) as cursor:
+                    cursor.execute(
+                        """
+                        SELECT
+                            COALESCE(MAX(user_turn_number), 0) as max_global_turn,
+                            COALESCE(MAX(branch_turn_number), 0) as max_branch_turn
+                        FROM message_structure
+                        WHERE session_id = ? AND branch_id = ?
+                    """,
+                        (self.session_id, self._current_branch_id),
+                    )
+                    result = cursor.fetchone()
+                    current_turn = result[0] if result else 0
+                    current_branch_turn = result[1] if result else 0
+
+                # Process items and assign turn numbers correctly
+                structure_data = []
+                user_message_count = 0
+
+                for i, (item, msg_id) in enumerate(zip(items, message_ids)):
+                    msg_type = self._classify_message_type(item)
+                    tool_name = self._extract_tool_name(item)
+
+                    # If this is a user message, increment turn counters
+                    if self._is_user_message(item):
+                        user_message_count += 1
+                        item_turn = current_turn + user_message_count
+                        item_branch_turn = current_branch_turn + user_message_count
+                    else:
+                        # Non-user messages inherit the turn number of the most recent user message
+                        item_turn = current_turn + user_message_count
+                        item_branch_turn = current_branch_turn + user_message_count
+
+                    structure_data.append(
+                        (
+                            self.session_id,
+                            msg_id,
+                            self._current_branch_id,
+                            msg_type,
+                            seq_start + i + 1,  # Global sequence
+                            item_turn,  # Global turn number
+                            item_branch_turn,  # Branch-specific turn number
+                            tool_name,
+                        )
+                    )
+
+                with closing(conn.cursor()) as cursor:
+                    cursor.executemany(
+                        """
+                        INSERT INTO message_structure
+                        (session_id, message_id, branch_id, message_type, sequence_number,
+                         user_turn_number, branch_turn_number, tool_name)
+                        VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+                    """,
+                        structure_data,
+                    )
+                    conn.commit()
+
+        try:
+            await asyncio.to_thread(_add_structure_sync)
+        except Exception as e:
+            self._logger.error(
+                f"Failed to add structure metadata for session {self.session_id}: {e}"
+            )
+            # Try to clean up any orphaned messages to maintain consistency
+            try:
+                await self._cleanup_orphaned_messages()
+            except Exception as cleanup_error:
+                self._logger.error(f"Failed to cleanup orphaned messages: {cleanup_error}")
+            # Don't re-raise - structure metadata is supplementary
+
+    async def _cleanup_orphaned_messages(self) -> None:
+        """Remove messages that exist in agent_messages but not in message_structure.
+
+        This can happen if _add_structure_metadata fails after super().add_items() succeeds.
+        Used for maintaining data consistency.
+        """
+
+        def _cleanup_sync():
+            """Synchronous helper to cleanup orphaned messages."""
+            conn = self._get_connection()
+            # TODO: Refactor SQLiteSession to use asyncio.Lock instead of threading.Lock and update this code  # noqa: E501
+            with self._lock if self._is_memory_db else threading.Lock():
+                with closing(conn.cursor()) as cursor:
+                    # Find messages without structure metadata
+                    cursor.execute(
+                        """
+                        SELECT am.id
+                        FROM agent_messages am
+                        LEFT JOIN message_structure ms ON am.id = ms.message_id
+                        WHERE am.session_id = ? AND ms.message_id IS NULL
+                    """,
+                        (self.session_id,),
+                    )
+
+                    orphaned_ids = [row[0] for row in cursor.fetchall()]
+
+                    if orphaned_ids:
+                        # Delete orphaned messages
+                        placeholders = ",".join("?" * len(orphaned_ids))
+                        cursor.execute(
+                            f"DELETE FROM agent_messages WHERE id IN ({placeholders})", orphaned_ids
+                        )
+
+                        deleted_count = cursor.rowcount
+                        conn.commit()
+
+                        self._logger.info(f"Cleaned up {deleted_count} orphaned messages")
+                        return deleted_count
+
+                    return 0
+
+        return await asyncio.to_thread(_cleanup_sync)
+
+    def _classify_message_type(self, item: TResponseInputItem) -> str:
+        """Classify the type of a message item.
+
+        Args:
+            item: The message item to classify.
+
+        Returns:
+            String representing the message type (user, assistant, etc.).
+        """
+        if isinstance(item, dict):
+            if item.get("role") == "user":
+                return "user"
+            elif item.get("role") == "assistant":
+                return "assistant"
+            elif item.get("type"):
+                return str(item.get("type"))
+        return "other"
+
+    def _extract_tool_name(self, item: TResponseInputItem) -> str | None:
+        """Extract tool name if this is a tool call/output.
+
+        Args:
+            item: The message item to extract tool name from.
+
+        Returns:
+            Tool name if item is a tool call, None otherwise.
+        """
+        if isinstance(item, dict):
+            item_type = item.get("type")
+
+            # For MCP tools, try to extract from server_label if available
+            if item_type in {"mcp_call", "mcp_approval_request"} and "server_label" in item:
+                server_label = item.get("server_label")
+                tool_name = item.get("name")
+                if tool_name and server_label:
+                    return f"{server_label}.{tool_name}"
+                elif server_label:
+                    return str(server_label)
+                elif tool_name:
+                    return str(tool_name)
+
+            # For tool types without a 'name' field, derive from the type
+            elif item_type in {
+                "computer_call",
+                "file_search_call",
+                "web_search_call",
+                "code_interpreter_call",
+            }:
+                return item_type
+
+            # Most other tool calls have a 'name' field
+            elif "name" in item:
+                name = item.get("name")
+                return str(name) if name is not None else None
+
+        return None
+
+    def _is_user_message(self, item: TResponseInputItem) -> bool:
+        """Check if this is a user message.
+
+        Args:
+            item: The message item to check.
+
+        Returns:
+            True if the item is a user message, False otherwise.
+        """
+        return isinstance(item, dict) and item.get("role") == "user"
+
+    async def create_branch_from_turn(
+        self, turn_number: int, branch_name: str | None = None
+    ) -> str:
+        """Create a new branch starting from a specific user message turn.
+
+        Args:
+            turn_number: The branch turn number of the user message to branch from
+            branch_name: Optional name for the branch (auto-generated if None)
+
+        Returns:
+            The branch_id of the newly created branch
+
+        Raises:
+            ValueError: If turn doesn't exist or doesn't contain a user message
+        """
+        import time
+
+        # Validate the turn exists and contains a user message
+        def _validate_turn():
+            """Synchronous helper to validate turn exists and contains user message."""
+            conn = self._get_connection()
+            with closing(conn.cursor()) as cursor:
+                cursor.execute(
+                    """
+                    SELECT am.message_data
+                    FROM message_structure ms
+                    JOIN agent_messages am ON ms.message_id = am.id
+                    WHERE ms.session_id = ? AND ms.branch_id = ?
+                    AND ms.branch_turn_number = ? AND ms.message_type = 'user'
+                    """,
+                    (self.session_id, self._current_branch_id, turn_number),
+                )
+
+                result = cursor.fetchone()
+                if not result:
+                    raise ValueError(
+                        f"Turn {turn_number} does not contain a user message "
+                        f"in branch '{self._current_branch_id}'"
+                    )
+
+                message_data = result[0]
+                try:
+                    content = json.loads(message_data).get("content", "")
+                    return content[:50] + "..." if len(content) > 50 else content
+                except Exception:
+                    return "Unable to parse content"
+
+        turn_content = await asyncio.to_thread(_validate_turn)
+
+        # Generate branch name if not provided
+        if branch_name is None:
+            timestamp = int(time.time())
+            branch_name = f"branch_from_turn_{turn_number}_{timestamp}"
+
+        # Copy messages before the branch point to the new branch
+        await self._copy_messages_to_new_branch(branch_name, turn_number)
+
+        # Switch to new branch
+        old_branch = self._current_branch_id
+        self._current_branch_id = branch_name
+
+        self._logger.debug(
+            f"Created branch '{branch_name}' from turn {turn_number} ('{turn_content}') in '{old_branch}'"  # noqa: E501
+        )
+        return branch_name
+
+    async def create_branch_from_content(
+        self, search_term: str, branch_name: str | None = None
+    ) -> str:
+        """Create branch from the first user turn matching the search term.
+
+        Args:
+            search_term: Text to search for in user messages.
+            branch_name: Optional name for the branch (auto-generated if None).
+
+        Returns:
+            The branch_id of the newly created branch.
+
+        Raises:
+            ValueError: If no matching turns are found.
+        """
+        matching_turns = await self.find_turns_by_content(search_term)
+        if not matching_turns:
+            raise ValueError(f"No user turns found containing '{search_term}'")
+
+        # Use the first (earliest) match
+        turn_number = matching_turns[0]["turn"]
+        return await self.create_branch_from_turn(turn_number, branch_name)
+
+    async def switch_to_branch(self, branch_id: str) -> None:
+        """Switch to a different branch.
+
+        Args:
+            branch_id: The branch to switch to.
+
+        Raises:
+            ValueError: If the branch doesn't exist.
+        """
+
+        # Validate branch exists
+        def _validate_branch():
+            """Synchronous helper to validate branch exists."""
+            conn = self._get_connection()
+            with closing(conn.cursor()) as cursor:
+                cursor.execute(
+                    """
+                    SELECT COUNT(*) FROM message_structure
+                    WHERE session_id = ? AND branch_id = ?
+                """,
+                    (self.session_id, branch_id),
+                )
+
+                count = cursor.fetchone()[0]
+                if count == 0:
+                    raise ValueError(f"Branch '{branch_id}' does not exist")
+
+        await asyncio.to_thread(_validate_branch)
+
+        old_branch = self._current_branch_id
+        self._current_branch_id = branch_id
+        self._logger.info(f"Switched from branch '{old_branch}' to '{branch_id}'")
+
+    async def delete_branch(self, branch_id: str, force: bool = False) -> None:
+        """Delete a branch and all its associated data.
+
+        Args:
+            branch_id: The branch to delete.
+            force: If True, allows deleting the current branch (will switch to 'main').
+
+        Raises:
+            ValueError: If branch doesn't exist, is 'main', or is current branch without force.
+        """
+        if not branch_id or not branch_id.strip():
+            raise ValueError("Branch ID cannot be empty")
+
+        branch_id = branch_id.strip()
+
+        # Protect main branch
+        if branch_id == "main":
+            raise ValueError("Cannot delete the 'main' branch")
+
+        # Check if trying to delete current branch
+        if branch_id == self._current_branch_id:
+            if not force:
+                raise ValueError(
+                    f"Cannot delete current branch '{branch_id}'. Use force=True or switch branches first"  # noqa: E501
+                )
+            else:
+                # Switch to main before deleting
+                await self.switch_to_branch("main")
+
+        def _delete_sync():
+            """Synchronous helper to delete branch and associated data."""
+            conn = self._get_connection()
+            # TODO: Refactor SQLiteSession to use asyncio.Lock instead of threading.Lock and update this code  # noqa: E501
+            with self._lock if self._is_memory_db else threading.Lock():
+                with closing(conn.cursor()) as cursor:
+                    # First verify the branch exists
+                    cursor.execute(
+                        """
+                        SELECT COUNT(*) FROM message_structure
+                        WHERE session_id = ? AND branch_id = ?
+                    """,
+                        (self.session_id, branch_id),
+                    )
+
+                    count = cursor.fetchone()[0]
+                    if count == 0:
+                        raise ValueError(f"Branch '{branch_id}' does not exist")
+
+                    # Delete from turn_usage first (foreign key constraint)
+                    cursor.execute(
+                        """
+                        DELETE FROM turn_usage
+                        WHERE session_id = ? AND branch_id = ?
+                    """,
+                        (self.session_id, branch_id),
+                    )
+
+                    usage_deleted = cursor.rowcount
+
+                    # Delete from message_structure
+                    cursor.execute(
+                        """
+                        DELETE FROM message_structure
+                        WHERE session_id = ? AND branch_id = ?
+                    """,
+                        (self.session_id, branch_id),
+                    )
+
+                    structure_deleted = cursor.rowcount
+
+                    conn.commit()
+
+                    return usage_deleted, structure_deleted
+
+        usage_deleted, structure_deleted = await asyncio.to_thread(_delete_sync)
+
+        self._logger.info(
+            f"Deleted branch '{branch_id}': {structure_deleted} message entries, {usage_deleted} usage entries"  # noqa: E501
+        )
+
+    async def list_branches(self) -> list[dict[str, Any]]:
+        """List all branches in this session.
+
+        Returns:
+            List of dicts with branch info containing:
+                - 'branch_id': Branch identifier
+                - 'message_count': Number of messages in branch
+                - 'user_turns': Number of user turns in branch
+                - 'is_current': Whether this is the current branch
+                - 'created_at': When the branch was first created
+        """
+
+        def _list_branches_sync():
+            """Synchronous helper to list all branches."""
+            conn = self._get_connection()
+            with closing(conn.cursor()) as cursor:
+                cursor.execute(
+                    """
+                    SELECT
+                        ms.branch_id,
+                        COUNT(*) as message_count,
+                        COUNT(CASE WHEN ms.message_type = 'user' THEN 1 END) as user_turns,
+                        MIN(ms.created_at) as created_at
+                    FROM message_structure ms
+                    WHERE ms.session_id = ?
+                    GROUP BY ms.branch_id
+                    ORDER BY created_at
+                """,
+                    (self.session_id,),
+                )
+
+                branches = []
+                for row in cursor.fetchall():
+                    branch_id, msg_count, user_turns, created_at = row
+                    branches.append(
+                        {
+                            "branch_id": branch_id,
+                            "message_count": msg_count,
+                            "user_turns": user_turns,
+                            "is_current": branch_id == self._current_branch_id,
+                            "created_at": created_at,
+                        }
+                    )
+
+                return branches
+
+        return await asyncio.to_thread(_list_branches_sync)
+
+    async def _copy_messages_to_new_branch(self, new_branch_id: str, from_turn_number: int) -> None:
+        """Copy messages before the branch point to the new branch.
+
+        Args:
+            new_branch_id: The ID of the new branch to copy messages to.
+            from_turn_number: The turn number to copy messages up to (exclusive).
+        """
+
+        def _copy_sync():
+            """Synchronous helper to copy messages to new branch."""
+            conn = self._get_connection()
+            # TODO: Refactor SQLiteSession to use asyncio.Lock instead of threading.Lock and update this code  # noqa: E501
+            with self._lock if self._is_memory_db else threading.Lock():
+                with closing(conn.cursor()) as cursor:
+                    # Get all messages before the branch point
+                    cursor.execute(
+                        """
+                        SELECT
+                            ms.message_id,
+                            ms.message_type,
+                            ms.sequence_number,
+                            ms.user_turn_number,
+                            ms.branch_turn_number,
+                            ms.tool_name
+                        FROM message_structure ms
+                        WHERE ms.session_id = ? AND ms.branch_id = ?
+                        AND ms.branch_turn_number < ?
+                        ORDER BY ms.sequence_number
+                    """,
+                        (self.session_id, self._current_branch_id, from_turn_number),
+                    )
+
+                    messages_to_copy = cursor.fetchall()
+
+                    if messages_to_copy:
+                        # Get the max sequence number for the new inserts
+                        cursor.execute(
+                            """
+                            SELECT COALESCE(MAX(sequence_number), 0)
+                            FROM message_structure
+                            WHERE session_id = ?
+                        """,
+                            (self.session_id,),
+                        )
+
+                        seq_start = cursor.fetchone()[0]
+
+                        # Insert copied messages with new branch_id
+                        new_structure_data = []
+                        for i, (
+                            msg_id,
+                            msg_type,
+                            _,
+                            user_turn,
+                            branch_turn,
+                            tool_name,
+                        ) in enumerate(messages_to_copy):
+                            new_structure_data.append(
+                                (
+                                    self.session_id,
+                                    msg_id,  # Same message_id (sharing the actual message data)
+                                    new_branch_id,
+                                    msg_type,
+                                    seq_start + i + 1,  # New sequence number
+                                    user_turn,  # Keep same global turn number
+                                    branch_turn,  # Keep same branch turn number
+                                    tool_name,
+                                )
+                            )
+
+                        cursor.executemany(
+                            """
+                            INSERT INTO message_structure
+                            (session_id, message_id, branch_id, message_type, sequence_number,
+                             user_turn_number, branch_turn_number, tool_name)
+                            VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+                        """,
+                            new_structure_data,
+                        )
+
+                    conn.commit()
+
+        await asyncio.to_thread(_copy_sync)
+
+    async def get_conversation_turns(self, branch_id: str | None = None) -> list[dict[str, Any]]:
+        """Get user turns with content for easy browsing and branching decisions.
+
+        Args:
+            branch_id: Branch to get turns from (current branch if None).
+
+        Returns:
+            List of dicts with turn info containing:
+                - 'turn': Branch turn number
+                - 'content': User message content (truncated)
+                - 'full_content': Full user message content
+                - 'timestamp': When the turn was created
+                - 'can_branch': Always True (all user messages can branch)
+        """
+        if branch_id is None:
+            branch_id = self._current_branch_id
+
+        def _get_turns_sync():
+            """Synchronous helper to get conversation turns."""
+            conn = self._get_connection()
+            with closing(conn.cursor()) as cursor:
+                cursor.execute(
+                    """
+                    SELECT
+                        ms.branch_turn_number,
+                        am.message_data,
+                        ms.created_at
+                    FROM message_structure ms
+                    JOIN agent_messages am ON ms.message_id = am.id
+                    WHERE ms.session_id = ? AND ms.branch_id = ?
+                    AND ms.message_type = 'user'
+                    ORDER BY ms.branch_turn_number
+                """,
+                    (self.session_id, branch_id),
+                )
+
+                turns = []
+                for row in cursor.fetchall():
+                    turn_num, message_data, created_at = row
+                    try:
+                        content = json.loads(message_data).get("content", "")
+                        turns.append(
+                            {
+                                "turn": turn_num,
+                                "content": content[:100] + "..." if len(content) > 100 else content,
+                                "full_content": content,
+                                "timestamp": created_at,
+                                "can_branch": True,
+                            }
+                        )
+                    except (json.JSONDecodeError, AttributeError):
+                        continue
+
+                return turns
+
+        return await asyncio.to_thread(_get_turns_sync)
+
+    async def find_turns_by_content(
+        self, search_term: str, branch_id: str | None = None
+    ) -> list[dict[str, Any]]:
+        """Find user turns containing specific content.
+
+        Args:
+            search_term: Text to search for in user messages.
+            branch_id: Branch to search in (current branch if None).
+
+        Returns:
+            List of matching turns with same format as get_conversation_turns().
+        """
+        if branch_id is None:
+            branch_id = self._current_branch_id
+
+        def _search_sync():
+            """Synchronous helper to search turns by content."""
+            conn = self._get_connection()
+            with closing(conn.cursor()) as cursor:
+                cursor.execute(
+                    """
+                    SELECT
+                        ms.branch_turn_number,
+                        am.message_data,
+                        ms.created_at
+                    FROM message_structure ms
+                    JOIN agent_messages am ON ms.message_id = am.id
+                    WHERE ms.session_id = ? AND ms.branch_id = ?
+                    AND ms.message_type = 'user'
+                    AND am.message_data LIKE ?
+                    ORDER BY ms.branch_turn_number
+                """,
+                    (self.session_id, branch_id, f"%{search_term}%"),
+                )
+
+                matches = []
+                for row in cursor.fetchall():
+                    turn_num, message_data, created_at = row
+                    try:
+                        content = json.loads(message_data).get("content", "")
+                        matches.append(
+                            {
+                                "turn": turn_num,
+                                "content": content,
+                                "full_content": content,
+                                "timestamp": created_at,
+                                "can_branch": True,
+                            }
+                        )
+                    except (json.JSONDecodeError, AttributeError):
+                        continue
+
+                return matches
+
+        return await asyncio.to_thread(_search_sync)
+
+    async def get_conversation_by_turns(
+        self, branch_id: str | None = None
+    ) -> dict[int, list[dict[str, str | None]]]:
+        """Get conversation grouped by user turns for specified branch.
+
+        Args:
+            branch_id: Branch to get conversation from (current branch if None).
+
+        Returns:
+            Dictionary mapping turn numbers to lists of message metadata.
+        """
+        if branch_id is None:
+            branch_id = self._current_branch_id
+
+        def _get_conversation_sync():
+            """Synchronous helper to get conversation by turns."""
+            conn = self._get_connection()
+            with closing(conn.cursor()) as cursor:
+                cursor.execute(
+                    """
+                    SELECT user_turn_number, message_type, tool_name
+                    FROM message_structure
+                    WHERE session_id = ? AND branch_id = ?
+                    ORDER BY sequence_number
+                """,
+                    (self.session_id, branch_id),
+                )
+
+                turns: dict[int, list[dict[str, str | None]]] = {}
+                for row in cursor.fetchall():
+                    turn_num, msg_type, tool_name = row
+                    if turn_num not in turns:
+                        turns[turn_num] = []
+                    turns[turn_num].append({"type": msg_type, "tool_name": tool_name})
+                return turns
+
+        return await asyncio.to_thread(_get_conversation_sync)
+
+    async def get_tool_usage(self, branch_id: str | None = None) -> list[tuple[str, int, int]]:
+        """Get all tool usage by turn for specified branch.
+
+        Args:
+            branch_id: Branch to get tool usage from (current branch if None).
+
+        Returns:
+            List of tuples containing (tool_name, usage_count, turn_number).
+        """
+        if branch_id is None:
+            branch_id = self._current_branch_id
+
+        def _get_tool_usage_sync():
+            """Synchronous helper to get tool usage statistics."""
+            conn = self._get_connection()
+            with closing(conn.cursor()) as cursor:
+                cursor.execute(
+                    """
+                    SELECT tool_name, COUNT(*), user_turn_number
+                    FROM message_structure
+                    WHERE session_id = ? AND branch_id = ? AND message_type IN (
+                        'tool_call', 'function_call', 'computer_call', 'file_search_call',
+                        'web_search_call', 'code_interpreter_call', 'custom_tool_call',
+                        'mcp_call', 'mcp_approval_request'
+                    )
+                    GROUP BY tool_name, user_turn_number
+                    ORDER BY user_turn_number
+                """,
+                    (self.session_id, branch_id),
+                )
+                return cursor.fetchall()
+
+        return await asyncio.to_thread(_get_tool_usage_sync)
+
+    async def get_session_usage(self, branch_id: str | None = None) -> dict[str, int] | None:
+        """Get cumulative usage for session or specific branch.
+
+        Args:
+            branch_id: If provided, only get usage for that branch. If None, get all branches.
+
+        Returns:
+            Dictionary with usage statistics or None if no usage data found.
+        """
+
+        def _get_usage_sync():
+            """Synchronous helper to get session usage data."""
+            conn = self._get_connection()
+            # TODO: Refactor SQLiteSession to use asyncio.Lock instead of threading.Lock and update this code  # noqa: E501
+            with self._lock if self._is_memory_db else threading.Lock():
+                if branch_id:
+                    # Branch-specific usage
+                    query = """
+                        SELECT
+                            SUM(requests) as total_requests,
+                            SUM(input_tokens) as total_input_tokens,
+                            SUM(output_tokens) as total_output_tokens,
+                            SUM(total_tokens) as total_total_tokens,
+                            COUNT(*) as total_turns
+                        FROM turn_usage
+                        WHERE session_id = ? AND branch_id = ?
+                    """
+                    params: tuple[str, ...] = (self.session_id, branch_id)
+                else:
+                    # All branches
+                    query = """
+                        SELECT
+                            SUM(requests) as total_requests,
+                            SUM(input_tokens) as total_input_tokens,
+                            SUM(output_tokens) as total_output_tokens,
+                            SUM(total_tokens) as total_total_tokens,
+                            COUNT(*) as total_turns
+                        FROM turn_usage
+                        WHERE session_id = ?
+                    """
+                    params = (self.session_id,)
+
+                with closing(conn.cursor()) as cursor:
+                    cursor.execute(query, params)
+                    row = cursor.fetchone()
+
+                    if row and row[0] is not None:
+                        return {
+                            "requests": row[0] or 0,
+                            "input_tokens": row[1] or 0,
+                            "output_tokens": row[2] or 0,
+                            "total_tokens": row[3] or 0,
+                            "total_turns": row[4] or 0,
+                        }
+                    return None
+
+        result = await asyncio.to_thread(_get_usage_sync)
+
+        return cast(Union[dict[str, int], None], result)
+
+    async def get_turn_usage(
+        self,
+        user_turn_number: int | None = None,
+        branch_id: str | None = None,
+    ) -> list[dict[str, Any]] | dict[str, Any]:
+        """Get usage statistics by turn with full JSON token details.
+
+        Args:
+            user_turn_number: Specific turn to get usage for. If None, returns all turns.
+            branch_id: Branch to get usage from (current branch if None).
+
+        Returns:
+            Dictionary with usage data for specific turn, or list of dictionaries for all turns.
+        """
+
+        if branch_id is None:
+            branch_id = self._current_branch_id
+
+        def _get_turn_usage_sync():
+            """Synchronous helper to get turn usage statistics."""
+            conn = self._get_connection()
+
+            if user_turn_number is not None:
+                query = """
+                    SELECT requests, input_tokens, output_tokens, total_tokens,
+                           input_tokens_details, output_tokens_details
+                    FROM turn_usage
+                    WHERE session_id = ? AND branch_id = ? AND user_turn_number = ?
+                """
+
+                with closing(conn.cursor()) as cursor:
+                    cursor.execute(query, (self.session_id, branch_id, user_turn_number))
+                    row = cursor.fetchone()
+
+                    if row:
+                        # Parse JSON details if present
+                        input_details = None
+                        output_details = None
+
+                        if row[4]:  # input_tokens_details
+                            try:
+                                input_details = json.loads(row[4])
+                            except json.JSONDecodeError:
+                                pass
+
+                        if row[5]:  # output_tokens_details
+                            try:
+                                output_details = json.loads(row[5])
+                            except json.JSONDecodeError:
+                                pass
+
+                        return {
+                            "requests": row[0],
+                            "input_tokens": row[1],
+                            "output_tokens": row[2],
+                            "total_tokens": row[3],
+                            "input_tokens_details": input_details,
+                            "output_tokens_details": output_details,
+                        }
+                    return {}
+            else:
+                query = """
+                    SELECT user_turn_number, requests, input_tokens, output_tokens,
+                           total_tokens, input_tokens_details, output_tokens_details
+                    FROM turn_usage
+                    WHERE session_id = ? AND branch_id = ?
+                    ORDER BY user_turn_number
+                """
+
+                with closing(conn.cursor()) as cursor:
+                    cursor.execute(query, (self.session_id, branch_id))
+                    results = []
+                    for row in cursor.fetchall():
+                        # Parse JSON details if present
+                        input_details = None
+                        output_details = None
+
+                        if row[5]:  # input_tokens_details
+                            try:
+                                input_details = json.loads(row[5])
+                            except json.JSONDecodeError:
+                                pass
+
+                        if row[6]:  # output_tokens_details
+                            try:
+                                output_details = json.loads(row[6])
+                            except json.JSONDecodeError:
+                                pass
+
+                        results.append(
+                            {
+                                "user_turn_number": row[0],
+                                "requests": row[1],
+                                "input_tokens": row[2],
+                                "output_tokens": row[3],
+                                "total_tokens": row[4],
+                                "input_tokens_details": input_details,
+                                "output_tokens_details": output_details,
+                            }
+                        )
+                    return results
+
+        result = await asyncio.to_thread(_get_turn_usage_sync)
+
+        return cast(Union[list[dict[str, Any]], dict[str, Any]], result)
+
+    async def _update_turn_usage_internal(self, user_turn_number: int, usage_data: Usage) -> None:
+        """Internal method to update usage for a specific turn with full JSON details.
+
+        Args:
+            user_turn_number: The turn number to update usage for.
+            usage_data: The usage data to store.
+        """
+
+        def _update_sync():
+            """Synchronous helper to update turn usage data."""
+            conn = self._get_connection()
+            # TODO: Refactor SQLiteSession to use asyncio.Lock instead of threading.Lock and update this code  # noqa: E501
+            with self._lock if self._is_memory_db else threading.Lock():
+                # Serialize token details as JSON
+                input_details_json = None
+                output_details_json = None
+
+                if hasattr(usage_data, "input_tokens_details") and usage_data.input_tokens_details:
+                    try:
+                        input_details_json = json.dumps(usage_data.input_tokens_details.__dict__)
+                    except (TypeError, ValueError) as e:
+                        self._logger.warning(f"Failed to serialize input tokens details: {e}")
+                        input_details_json = None
+
+                    if (
+                        hasattr(usage_data, "output_tokens_details")
+                        and usage_data.output_tokens_details
+                    ):
+                        try:
+                            output_details_json = json.dumps(
+                                usage_data.output_tokens_details.__dict__
+                            )
+                        except (TypeError, ValueError) as e:
+                            self._logger.warning(f"Failed to serialize output tokens details: {e}")
+                            output_details_json = None
+
+                with closing(conn.cursor()) as cursor:
+                    cursor.execute(
+                        """
+                        INSERT OR REPLACE INTO turn_usage
+                        (session_id, branch_id, user_turn_number, requests, input_tokens, output_tokens,
+                         total_tokens, input_tokens_details, output_tokens_details)
+                        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
+                    """,  # noqa: E501
+                        (
+                            self.session_id,
+                            self._current_branch_id,
+                            user_turn_number,
+                            usage_data.requests or 0,
+                            usage_data.input_tokens or 0,
+                            usage_data.output_tokens or 0,
+                            usage_data.total_tokens or 0,
+                            input_details_json,
+                            output_details_json,
+                        ),
+                    )
+                    conn.commit()
+
+        await asyncio.to_thread(_update_sync)