advanced_sqlite_session.py

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 ..._tool_identity import is_reserved_synthetic_tool_namespace, tool_qualified_name
from ...items import TResponseInputItem
from ...memory import SQLiteSession
from ...memory.session_settings import SessionSettings, resolve_session_limit


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,
        session_settings: SessionSettings | 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=session_id,
            db_path=db_path,
            session_settings=session_settings,
            **kwargs,
        )
        if create_tables:
            self._init_structure_tables()
        self._current_branch_id = "main"
        self._logger = logger or logging.getLogger(__name__)
        # Note: _file_db_lock is inherited from SQLiteSession base class for thread safety

    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(f"""
            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 {self.sessions_table}(session_id) ON DELETE CASCADE,
                FOREIGN KEY (message_id)
                    REFERENCES {self.messages_table}(id) ON DELETE CASCADE
            )
        """)

        # Turn-level usage tracking with branch support and full JSON details
        conn.execute(f"""
            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 {self.sessions_table}(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, uses session_settings.limit.
            branch_id: Branch to get items from. If None, uses current branch.

        Returns:
            List of conversation items from the specified branch.
        """
        session_limit = resolve_session_limit(limit, self.session_settings)

        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()
                # Use the shared instance lock inherited from base class for disk-based DBs
                with self._lock if self._is_memory_db else self._file_db_lock:
                    with closing(conn.cursor()) as cursor:
                        if session_limit is None:
                            cursor.execute(
                                f"""
                                SELECT m.message_data
                                FROM {self.messages_table} 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(
                                f"""
                                SELECT m.message_data
                                FROM {self.messages_table} 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, session_limit),
                            )

                        rows = cursor.fetchall()
                        if session_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()
            # Use the instance lock for disk-based DBs to ensure consistent locking
            with self._lock if self._is_memory_db else self._file_db_lock:
                with closing(conn.cursor()) as cursor:
                    # Get message IDs in correct order for this branch
                    if session_limit is None:
                        cursor.execute(
                            f"""
                            SELECT m.message_data
                            FROM {self.messages_table} 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(
                            f"""
                            SELECT m.message_data
                            FROM {self.messages_table} 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, session_limit),
                        )

                    rows = cursor.fetchall()
                    if session_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()
            # Use the instance lock for disk-based DBs to ensure consistent locking
            with self._lock if self._is_memory_db else self._file_db_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) -> int:
        """Remove messages that exist in the configured message table 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()
            # Use the instance lock for disk-based DBs to ensure consistent locking
            with self._lock if self._is_memory_db else self._file_db_lock:
                with closing(conn.cursor()) as cursor:
                    # Find messages without structure metadata
                    cursor.execute(
                        f"""
                        SELECT am.id
                        FROM {self.messages_table} 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 {self.messages_table} 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",
                "tool_search_call",
                "tool_search_output",
            }:
                if item_type in {"tool_search_call", "tool_search_output"}:
                    return "tool_search"
                return item_type

            # Most other tool calls have a 'name' field
            elif "name" in item:
                name = item.get("name")
                namespace = item.get("namespace")
                if name is not None:
                    name_str = str(name)
                    namespace_str = str(namespace) if namespace is not None else None
                    if is_reserved_synthetic_tool_namespace(name_str, namespace_str):
                        return name_str
                    qualified_name = tool_qualified_name(
                        name_str,
                        namespace_str,
                    )
                    return qualified_name or name_str
                return 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(
                    f"""
                    SELECT am.message_data
                    FROM message_structure ms
                    JOIN {self.messages_table} 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()
            # Use the instance lock for disk-based DBs to ensure consistent locking
            with self._lock if self._is_memory_db else self._file_db_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()
            # Use the instance lock for disk-based DBs to ensure consistent locking
            with self._lock if self._is_memory_db else self._file_db_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(
                    f"""
                    SELECT
                        ms.branch_turn_number,
                        am.message_data,
                        ms.created_at
                    FROM message_structure ms
                    JOIN {self.messages_table} 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(
                    f"""
                    SELECT
                        ms.branch_turn_number,
                        am.message_data,
                        ms.created_at
                    FROM message_structure ms
                    JOIN {self.messages_table} 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, SUM(usage_count), user_turn_number
                    FROM (
                        SELECT tool_name, 1 AS usage_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', 'tool_search_call',
                            'custom_tool_call', 'mcp_call', 'mcp_approval_request'
                        )

                        UNION ALL

                        SELECT ms.tool_name, 1 AS usage_count, ms.user_turn_number
                        FROM message_structure ms
                        WHERE ms.session_id = ? AND ms.branch_id = ?
                          AND ms.message_type = 'tool_search_output'
                          AND NOT EXISTS (
                              SELECT 1
                              FROM message_structure calls
                              WHERE calls.session_id = ms.session_id
                                AND calls.branch_id = ms.branch_id
                                AND calls.user_turn_number = ms.user_turn_number
                                AND calls.tool_name = ms.tool_name
                                AND calls.message_type = 'tool_search_call'
                          )
                    )
                    GROUP BY tool_name, user_turn_number
                    ORDER BY user_turn_number
                """,
                    (
                        self.session_id,
                        branch_id,
                        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()
            # Use the instance lock for disk-based DBs to ensure consistent locking
            with self._lock if self._is_memory_db else self._file_db_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()
            # Use the instance lock for disk-based DBs to ensure consistent locking
            with self._lock if self._is_memory_db else self._file_db_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)