prompt_objects 0.1.0

data/lib/prompt_objects/session/store.rb

@@ -0,0 +1,826 @@
+# frozen_string_literal: true
+
+require "sqlite3"
+require "json"
+require "securerandom"
+
+module PromptObjects
+  module Session
+    # SQLite-based session storage for conversation history.
+    # Each environment has its own sessions.db file (gitignored for privacy).
+    class Store
+      SCHEMA_VERSION = 4
+
+      # Thread types for conversation branching
+      THREAD_TYPES = %w[root continuation delegation fork].freeze
+
+      # Valid source values for session tracking
+      SOURCES = %w[tui mcp api web cli].freeze
+
+      # @param db_path [String] Path to the SQLite database file
+      def initialize(db_path)
+        @db_path = db_path
+        @db = SQLite3::Database.new(db_path)
+        @db.results_as_hash = true
+
+        # Enable WAL mode for better concurrent access (TUI + MCP can access simultaneously)
+        @db.execute("PRAGMA journal_mode=WAL")
+
+        # Set busy timeout to 5 seconds - wait for locks instead of failing immediately
+        @db.busy_timeout = 5000
+
+        setup_schema
+      end
+
+      # Close the database connection.
+      def close
+        @db.close if @db
+      end
+
+      # --- Session/Thread CRUD ---
+
+      # Create a new session (thread) for a PO.
+      # @param po_name [String] Name of the prompt object
+      # @param name [String, nil] Optional session name
+      # @param source [String] Source interface (tui, mcp, api, web, cli)
+      # @param source_client [String, nil] Client identifier (e.g., "claude-desktop", "cursor")
+      # @param metadata [Hash] Optional metadata
+      # @param parent_session_id [String, nil] Parent thread ID (for branching)
+      # @param parent_message_id [Integer, nil] Message ID that spawned this thread
+      # @param parent_po [String, nil] PO that created this thread (for cross-PO delegation)
+      # @param thread_type [String] Type of thread: root, continuation, delegation, fork
+      # @return [String] Session ID
+      def create_session(po_name:, name: nil, source: "tui", source_client: nil, metadata: {},
+                         parent_session_id: nil, parent_message_id: nil, parent_po: nil, thread_type: "root")
+        id = SecureRandom.uuid
+        now = Time.now.utc.iso8601
+
+        @db.execute(<<~SQL, [id, po_name, name, source, source_client, source, now, now, metadata.to_json, parent_session_id, parent_message_id, parent_po, thread_type])
+          INSERT INTO sessions (id, po_name, name, source, source_client, last_message_source, created_at, updated_at, metadata, parent_session_id, parent_message_id, parent_po, thread_type)
+          VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+        SQL
+
+        id
+      end
+
+      # Create a new thread (alias for create_session with thread semantics).
+      # @param po_name [String] Name of the prompt object
+      # @param parent_session_id [String, nil] Parent thread ID
+      # @param parent_message_id [Integer, nil] Message ID that spawned this thread
+      # @param parent_po [String, nil] PO that created this thread
+      # @param thread_type [String] Type: root, continuation, delegation, fork
+      # @param opts [Hash] Additional options passed to create_session
+      # @return [String] Thread/Session ID
+      def create_thread(po_name:, parent_session_id: nil, parent_message_id: nil, parent_po: nil, thread_type: "root", **opts)
+        create_session(
+          po_name: po_name,
+          parent_session_id: parent_session_id,
+          parent_message_id: parent_message_id,
+          parent_po: parent_po,
+          thread_type: thread_type,
+          **opts
+        )
+      end
+
+      # Get a session by ID.
+      # @param id [String] Session ID
+      # @return [Hash, nil] Session data or nil if not found
+      def get_session(id)
+        row = @db.get_first_row("SELECT * FROM sessions WHERE id = ?", [id])
+        return nil unless row
+
+        parse_session_row(row)
+      end
+
+      # Get the most recent session for a PO, or create one if none exists.
+      # @param po_name [String] Name of the prompt object
+      # @param source [String] Source interface for new session
+      # @param source_client [String, nil] Client identifier for new session
+      # @return [Hash] Session data
+      def get_or_create_session(po_name:, source: "tui", source_client: nil)
+        session = get_latest_session(po_name: po_name)
+        return session if session
+
+        id = create_session(po_name: po_name, source: source, source_client: source_client)
+        get_session(id)
+      end
+
+      # Get the most recent session for a PO.
+      # @param po_name [String] Name of the prompt object
+      # @return [Hash, nil] Session data or nil
+      def get_latest_session(po_name:)
+        row = @db.get_first_row(<<~SQL, [po_name])
+          SELECT * FROM sessions
+          WHERE po_name = ?
+          ORDER BY updated_at DESC
+          LIMIT 1
+        SQL
+
+        return nil unless row
+
+        parse_session_row(row)
+      end
+
+      # List all sessions for a PO.
+      # @param po_name [String] Name of the prompt object
+      # @return [Array<Hash>] Session data with message counts
+      def list_sessions(po_name:)
+        rows = @db.execute(<<~SQL, [po_name])
+          SELECT s.*, COUNT(m.id) as message_count
+          FROM sessions s
+          LEFT JOIN messages m ON m.session_id = s.id
+          WHERE s.po_name = ?
+          GROUP BY s.id
+          ORDER BY s.updated_at DESC
+        SQL
+
+        rows.map { |row| parse_session_row(row, include_count: true) }
+      end
+
+      # List all sessions across all POs.
+      # @param source [String, nil] Filter by source interface
+      # @param limit [Integer, nil] Maximum number of sessions
+      # @return [Array<Hash>] Session data with message counts
+      def list_all_sessions(source: nil, limit: nil)
+        sql = <<~SQL
+          SELECT s.*, COUNT(m.id) as message_count
+          FROM sessions s
+          LEFT JOIN messages m ON m.session_id = s.id
+        SQL
+
+        params = []
+        if source
+          sql += " WHERE s.source = ?"
+          params << source
+        end
+
+        sql += " GROUP BY s.id ORDER BY s.updated_at DESC"
+
+        if limit
+          sql += " LIMIT ?"
+          params << limit
+        end
+
+        rows = @db.execute(sql, params)
+        rows.map { |row| parse_session_row(row, include_count: true) }
+      end
+
+      # --- Thread Navigation ---
+
+      # Get all child threads of a session.
+      # @param session_id [String] Parent session ID
+      # @return [Array<Hash>] Child session data with message counts
+      def get_child_threads(session_id)
+        rows = @db.execute(<<~SQL, [session_id])
+          SELECT s.*, COUNT(m.id) as message_count
+          FROM sessions s
+          LEFT JOIN messages m ON m.session_id = s.id
+          WHERE s.parent_session_id = ?
+          GROUP BY s.id
+          ORDER BY s.created_at ASC
+        SQL
+
+        rows.map { |row| parse_session_row(row, include_count: true) }
+      end
+
+      # Get the lineage (path from root) for a thread.
+      # @param session_id [String] Session ID
+      # @return [Array<Hash>] Ancestors from root to current (inclusive)
+      def get_thread_lineage(session_id)
+        lineage = []
+        current_id = session_id
+
+        while current_id
+          session = get_session(current_id)
+          break unless session
+
+          lineage.unshift(session)
+          current_id = session[:parent_session_id]
+        end
+
+        lineage
+      end
+
+      # Get the full thread tree starting from a session.
+      # @param session_id [String] Root session ID
+      # @param max_depth [Integer] Maximum recursion depth
+      # @return [Hash] Tree structure with session and children
+      def get_thread_tree(session_id, max_depth: 10)
+        return nil if max_depth <= 0
+
+        session = get_session(session_id)
+        return nil unless session
+
+        # Add message count
+        session[:message_count] = message_count(session_id)
+
+        children = get_child_threads(session_id)
+        child_trees = children.map { |child| get_thread_tree(child[:id], max_depth: max_depth - 1) }.compact
+
+        {
+          session: session,
+          children: child_trees
+        }
+      end
+
+      # Get root threads for a PO (threads with no parent).
+      # @param po_name [String] Name of the prompt object
+      # @return [Array<Hash>] Root session data with message counts
+      def get_root_threads(po_name:)
+        rows = @db.execute(<<~SQL, [po_name])
+          SELECT s.*, COUNT(m.id) as message_count
+          FROM sessions s
+          LEFT JOIN messages m ON m.session_id = s.id
+          WHERE s.po_name = ? AND s.parent_session_id IS NULL
+          GROUP BY s.id
+          ORDER BY s.updated_at DESC
+        SQL
+
+        rows.map { |row| parse_session_row(row, include_count: true) }
+      end
+
+      # Auto-generate a name for a thread from its first message.
+      # @param session_id [String] Session ID
+      # @param first_message [String] First message content
+      # @param max_length [Integer] Maximum name length
+      def auto_name_thread(session_id, first_message, max_length: 40)
+        return unless first_message
+
+        auto_name = first_message.to_s.gsub(/\s+/, " ").strip[0, max_length]
+        auto_name += "..." if first_message.to_s.length > max_length
+        update_session(session_id, name: auto_name)
+      end
+
+      # Search sessions by message content using full-text search.
+      # @param query [String] Search query
+      # @param po_name [String, nil] Filter by PO
+      # @param source [String, nil] Filter by source
+      # @param limit [Integer] Maximum results
+      # @return [Array<Hash>] Sessions with match info
+      def search_sessions(query, po_name: nil, source: nil, limit: 50)
+        return [] if query.nil? || query.strip.empty?
+
+        # Use FTS5 MATCH syntax
+        # Escape special characters and add prefix matching
+        safe_query = query.gsub(/['"()]/, " ").strip
+        fts_query = safe_query.split.map { |term| "#{term}*" }.join(" ")
+
+        # First get matching message IDs from FTS
+        sql = <<~SQL
+          SELECT DISTINCT s.*, COUNT(m.id) as message_count
+          FROM sessions s
+          INNER JOIN messages m ON m.session_id = s.id
+          INNER JOIN messages_fts ON messages_fts.rowid = m.id
+          WHERE messages_fts MATCH ?
+        SQL
+
+        params = [fts_query]
+
+        if po_name
+          sql += " AND s.po_name = ?"
+          params << po_name
+        end
+
+        if source
+          sql += " AND s.source = ?"
+          params << source
+        end
+
+        sql += " GROUP BY s.id ORDER BY s.updated_at DESC LIMIT ?"
+        params << limit
+
+        rows = @db.execute(sql, params)
+        results = rows.map { |row| parse_session_row(row, include_count: true) }
+
+        # Get a snippet from matching messages for each session
+        results.each do |session|
+          snippet = get_match_snippet(session[:id], query)
+          session[:match_snippet] = snippet if snippet
+        end
+
+        results
+      rescue SQLite3::SQLException => e
+        # FTS table might not exist in older databases
+        if e.message.include?("no such table")
+          search_sessions_fallback(query, po_name: po_name, source: source, limit: limit)
+        else
+          raise
+        end
+      end
+
+      # Get a snippet of matching content from a session
+      # @param session_id [String] Session ID
+      # @param query [String] Search query
+      # @return [String, nil] Snippet with match highlighted
+      def get_match_snippet(session_id, query)
+        # Get first message that matches
+        row = @db.get_first_row(<<~SQL, [session_id, "%#{query}%"])
+          SELECT content FROM messages
+          WHERE session_id = ? AND content LIKE ?
+          LIMIT 1
+        SQL
+
+        return nil unless row && row["content"]
+
+        content = row["content"]
+        # Find match position and extract snippet
+        query_lower = query.downcase
+        pos = content.downcase.index(query_lower)
+        return content[0, 60] + "..." unless pos
+
+        # Extract snippet around match
+        start_pos = [pos - 20, 0].max
+        end_pos = [pos + query.length + 40, content.length].min
+        snippet = content[start_pos, end_pos - start_pos]
+
+        # Add ellipsis if truncated
+        snippet = "..." + snippet if start_pos > 0
+        snippet = snippet + "..." if end_pos < content.length
+
+        # Highlight match with markers
+        snippet.gsub(/#{Regexp.escape(query)}/i) { |m| ">>>#{m}<<<" }
+      end
+
+      # Fallback search without FTS (slower but works on older databases)
+      # @param query [String] Search query
+      # @param po_name [String, nil] Filter by PO
+      # @param source [String, nil] Filter by source
+      # @param limit [Integer] Maximum results
+      # @return [Array<Hash>] Sessions with match info
+      def search_sessions_fallback(query, po_name: nil, source: nil, limit: 50)
+        return [] if query.nil? || query.strip.empty?
+
+        sql = <<~SQL
+          SELECT DISTINCT s.*, COUNT(m.id) as message_count
+          FROM sessions s
+          INNER JOIN messages m ON m.session_id = s.id
+          WHERE m.content LIKE ?
+        SQL
+
+        params = ["%#{query}%"]
+
+        if po_name
+          sql += " AND s.po_name = ?"
+          params << po_name
+        end
+
+        if source
+          sql += " AND s.source = ?"
+          params << source
+        end
+
+        sql += " GROUP BY s.id ORDER BY s.updated_at DESC LIMIT ?"
+        params << limit
+
+        rows = @db.execute(sql, params)
+        rows.map { |row| parse_session_row(row, include_count: true) }
+      end
+
+      # Update a session's metadata.
+      # @param id [String] Session ID
+      # @param name [String, nil] New session name
+      # @param metadata [Hash, nil] New metadata (merged with existing)
+      # @param last_message_source [String, nil] Source of last message (tui, mcp, api)
+      def update_session(id, name: nil, metadata: nil, last_message_source: nil)
+        updates = ["updated_at = ?"]
+        params = [Time.now.utc.iso8601]
+
+        if name
+          updates << "name = ?"
+          params << name
+        end
+
+        if last_message_source
+          updates << "last_message_source = ?"
+          params << last_message_source
+        end
+
+        if metadata
+          # Merge with existing metadata
+          existing = get_session(id)
+          if existing
+            merged = (existing[:metadata] || {}).merge(metadata)
+            updates << "metadata = ?"
+            params << merged.to_json
+          end
+        end
+
+        params << id
+        @db.execute("UPDATE sessions SET #{updates.join(', ')} WHERE id = ?", params)
+      end
+
+      # Delete a session and all its messages.
+      # @param id [String] Session ID
+      def delete_session(id)
+        @db.execute("DELETE FROM messages WHERE session_id = ?", [id])
+        @db.execute("DELETE FROM sessions WHERE id = ?", [id])
+      end
+
+      # --- Message CRUD ---
+
+      # Add a message to a session.
+      # @param session_id [String] Session ID
+      # @param role [Symbol, String] Message role (:user, :assistant, :tool)
+      # @param content [String, nil] Message content
+      # @param from_po [String, nil] Source PO for delegation tracking
+      # @param tool_calls [Array, nil] Tool calls data
+      # @param tool_results [Array, nil] Tool results data
+      # @param source [String, nil] Source interface that added this message
+      # @return [Integer] Message ID
+      def add_message(session_id:, role:, content: nil, from_po: nil, tool_calls: nil, tool_results: nil, source: nil)
+        now = Time.now.utc.iso8601
+
+        params = [
+          session_id,
+          role.to_s,
+          content,
+          from_po,
+          tool_calls&.to_json,
+          tool_results&.to_json,
+          now
+        ]
+
+        @db.execute(<<~SQL, params)
+          INSERT INTO messages (session_id, role, content, from_po, tool_calls, tool_results, created_at)
+          VALUES (?, ?, ?, ?, ?, ?, ?)
+        SQL
+
+        # Update session's updated_at and optionally last_message_source
+        if source
+          @db.execute("UPDATE sessions SET updated_at = ?, last_message_source = ? WHERE id = ?", [now, source, session_id])
+        else
+          @db.execute("UPDATE sessions SET updated_at = ? WHERE id = ?", [now, session_id])
+        end
+
+        @db.last_insert_row_id
+      end
+
+      # Get all messages for a session.
+      # @param session_id [String] Session ID
+      # @return [Array<Hash>] Messages in chronological order
+      def get_messages(session_id)
+        rows = @db.execute(<<~SQL, [session_id])
+          SELECT * FROM messages
+          WHERE session_id = ?
+          ORDER BY id ASC
+        SQL
+
+        rows.map { |row| parse_message_row(row) }
+      end
+
+      # Get message count for a session.
+      # @param session_id [String] Session ID
+      # @return [Integer]
+      def message_count(session_id)
+        row = @db.get_first_row("SELECT COUNT(*) as count FROM messages WHERE session_id = ?", [session_id])
+        row["count"]
+      end
+
+      # Clear all messages from a session (but keep the session).
+      # @param session_id [String] Session ID
+      def clear_messages(session_id)
+        @db.execute("DELETE FROM messages WHERE session_id = ?", [session_id])
+        @db.execute("UPDATE sessions SET updated_at = ? WHERE id = ?", [Time.now.utc.iso8601, session_id])
+      end
+
+      # --- Stats ---
+
+      # Get total message count across all sessions.
+      # @return [Integer]
+      def total_messages
+        row = @db.get_first_row("SELECT COUNT(*) as count FROM messages")
+        row["count"]
+      end
+
+      # Get total session count.
+      # @return [Integer]
+      def total_sessions
+        row = @db.get_first_row("SELECT COUNT(*) as count FROM sessions")
+        row["count"]
+      end
+
+      # --- Export ---
+
+      # Export a session to JSON format.
+      # @param session_id [String] Session ID
+      # @return [Hash] Session data with messages
+      def export_session_json(session_id)
+        session = get_session(session_id)
+        return nil unless session
+
+        messages = get_messages(session_id)
+
+        {
+          id: session[:id],
+          po_name: session[:po_name],
+          name: session[:name],
+          source: session[:source],
+          source_client: session[:source_client],
+          created_at: session[:created_at]&.iso8601,
+          updated_at: session[:updated_at]&.iso8601,
+          metadata: session[:metadata],
+          messages: messages.map do |m|
+            {
+              role: m[:role].to_s,
+              content: m[:content],
+              from_po: m[:from_po],
+              tool_calls: m[:tool_calls],
+              tool_results: m[:tool_results],
+              created_at: m[:created_at]&.iso8601
+            }
+          end
+        }
+      end
+
+      # Export a session to Markdown format.
+      # @param session_id [String] Session ID
+      # @return [String] Markdown content
+      def export_session_markdown(session_id)
+        session = get_session(session_id)
+        return nil unless session
+
+        messages = get_messages(session_id)
+
+        lines = []
+        lines << "# Session: #{session[:name] || 'Unnamed'}"
+        lines << ""
+        lines << "- **PO**: #{session[:po_name]}"
+        lines << "- **Source**: #{session[:source]}"
+        lines << "- **Created**: #{session[:created_at]&.strftime('%Y-%m-%d %H:%M')}"
+        lines << "- **Updated**: #{session[:updated_at]&.strftime('%Y-%m-%d %H:%M')}"
+        lines << "- **Messages**: #{messages.length}"
+        lines << ""
+        lines << "---"
+        lines << ""
+
+        messages.each do |m|
+          timestamp = m[:created_at]&.strftime('%H:%M')
+          role_label = case m[:role].to_s
+                       when "user" then "**User**"
+                       when "assistant" then "**#{m[:from_po] || session[:po_name]}**"
+                       when "tool" then "*Tool*"
+                       else "**#{m[:role]}**"
+                       end
+
+          lines << "#{role_label} (#{timestamp}):"
+          lines << ""
+
+          if m[:content]
+            lines << m[:content]
+            lines << ""
+          end
+
+          if m[:tool_calls]
+            lines << "<details><summary>Tool calls</summary>"
+            lines << ""
+            lines << "```json"
+            lines << JSON.pretty_generate(m[:tool_calls])
+            lines << "```"
+            lines << "</details>"
+            lines << ""
+          end
+
+          if m[:tool_results]
+            lines << "<details><summary>Tool results</summary>"
+            lines << ""
+            lines << "```json"
+            lines << JSON.pretty_generate(m[:tool_results])
+            lines << "```"
+            lines << "</details>"
+            lines << ""
+          end
+
+          lines << "---"
+          lines << ""
+        end
+
+        lines.join("\n")
+      end
+
+      # Export all sessions for a PO.
+      # @param po_name [String] PO name
+      # @param format [Symbol] :json or :markdown
+      # @return [String] Exported content
+      def export_all_sessions(po_name:, format: :json)
+        sessions = list_sessions(po_name: po_name)
+
+        case format
+        when :json
+          exported = sessions.map { |s| export_session_json(s[:id]) }
+          JSON.pretty_generate(exported)
+        when :markdown
+          sessions.map { |s| export_session_markdown(s[:id]) }.join("\n\n")
+        else
+          raise ArgumentError, "Unknown format: #{format}"
+        end
+      end
+
+      # --- Import ---
+
+      # Import a session from JSON data.
+      # @param data [Hash] Session data (as returned by export_session_json)
+      # @param po_name [String, nil] Override PO name
+      # @return [String] New session ID
+      def import_session(data, po_name: nil)
+        data = data.transform_keys(&:to_sym) if data.is_a?(Hash)
+
+        # Create new session with new ID
+        new_id = create_session(
+          po_name: po_name || data[:po_name],
+          name: "#{data[:name]} (imported)",
+          source: "tui",
+          metadata: (data[:metadata] || {}).merge(
+            imported_from: data[:id],
+            imported_at: Time.now.utc.iso8601,
+            original_source: data[:source]
+          )
+        )
+
+        # Import messages
+        messages = data[:messages] || []
+        messages.each do |m|
+          m = m.transform_keys(&:to_sym) if m.is_a?(Hash)
+          add_message(
+            session_id: new_id,
+            role: m[:role],
+            content: m[:content],
+            from_po: m[:from_po],
+            tool_calls: m[:tool_calls],
+            tool_results: m[:tool_results]
+          )
+        end
+
+        new_id
+      end
+
+      private
+
+      def setup_schema
+        # Check if we need to create/migrate
+        version = get_schema_version
+
+        if version == 0
+          create_schema
+          set_schema_version(SCHEMA_VERSION)
+        elsif version < SCHEMA_VERSION
+          migrate_schema(version)
+          set_schema_version(SCHEMA_VERSION)
+        end
+      end
+
+      def get_schema_version
+        @db.get_first_value("PRAGMA user_version") || 0
+      end
+
+      def set_schema_version(version)
+        @db.execute("PRAGMA user_version = #{version}")
+      end
+
+      def create_schema
+        @db.execute_batch(<<~SQL)
+          CREATE TABLE IF NOT EXISTS sessions (
+            id TEXT PRIMARY KEY,
+            po_name TEXT NOT NULL,
+            name TEXT,
+            source TEXT DEFAULT 'tui',
+            source_client TEXT,
+            last_message_source TEXT,
+            created_at TEXT NOT NULL,
+            updated_at TEXT NOT NULL,
+            metadata TEXT DEFAULT '{}',
+            -- Thread/branching support (v4)
+            parent_session_id TEXT,
+            parent_message_id INTEGER,
+            parent_po TEXT,
+            thread_type TEXT DEFAULT 'root'
+          );
+
+          CREATE INDEX IF NOT EXISTS idx_sessions_po_name ON sessions(po_name);
+          CREATE INDEX IF NOT EXISTS idx_sessions_updated_at ON sessions(updated_at);
+          CREATE INDEX IF NOT EXISTS idx_sessions_source ON sessions(source);
+          CREATE INDEX IF NOT EXISTS idx_sessions_parent ON sessions(parent_session_id);
+
+          CREATE TABLE IF NOT EXISTS messages (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            session_id TEXT NOT NULL REFERENCES sessions(id),
+            role TEXT NOT NULL,
+            content TEXT,
+            from_po TEXT,
+            tool_calls TEXT,
+            tool_results TEXT,
+            created_at TEXT NOT NULL
+          );
+
+          CREATE INDEX IF NOT EXISTS idx_messages_session_id ON messages(session_id);
+
+          -- Full-text search index for message content
+          CREATE VIRTUAL TABLE IF NOT EXISTS messages_fts USING fts5(
+            content,
+            content='messages',
+            content_rowid='id'
+          );
+
+          -- Triggers to keep FTS in sync
+          CREATE TRIGGER IF NOT EXISTS messages_ai AFTER INSERT ON messages BEGIN
+            INSERT INTO messages_fts(rowid, content) VALUES (new.id, new.content);
+          END;
+
+          CREATE TRIGGER IF NOT EXISTS messages_ad AFTER DELETE ON messages BEGIN
+            INSERT INTO messages_fts(messages_fts, rowid, content) VALUES('delete', old.id, old.content);
+          END;
+
+          CREATE TRIGGER IF NOT EXISTS messages_au AFTER UPDATE ON messages BEGIN
+            INSERT INTO messages_fts(messages_fts, rowid, content) VALUES('delete', old.id, old.content);
+            INSERT INTO messages_fts(rowid, content) VALUES (new.id, new.content);
+          END;
+        SQL
+      end
+
+      def migrate_schema(from_version)
+        if from_version < 2
+          # Add source tracking columns
+          @db.execute_batch(<<~SQL)
+            ALTER TABLE sessions ADD COLUMN source TEXT DEFAULT 'tui';
+            ALTER TABLE sessions ADD COLUMN source_client TEXT;
+            ALTER TABLE sessions ADD COLUMN last_message_source TEXT;
+            CREATE INDEX IF NOT EXISTS idx_sessions_source ON sessions(source);
+          SQL
+        end
+
+        if from_version < 3
+          # Add full-text search for messages
+          @db.execute_batch(<<~SQL)
+            -- Full-text search index for message content
+            CREATE VIRTUAL TABLE IF NOT EXISTS messages_fts USING fts5(
+              content,
+              content='messages',
+              content_rowid='id'
+            );
+
+            -- Triggers to keep FTS in sync
+            CREATE TRIGGER IF NOT EXISTS messages_ai AFTER INSERT ON messages BEGIN
+              INSERT INTO messages_fts(rowid, content) VALUES (new.id, new.content);
+            END;
+
+            CREATE TRIGGER IF NOT EXISTS messages_ad AFTER DELETE ON messages BEGIN
+              INSERT INTO messages_fts(messages_fts, rowid, content) VALUES('delete', old.id, old.content);
+            END;
+
+            CREATE TRIGGER IF NOT EXISTS messages_au AFTER UPDATE ON messages BEGIN
+              INSERT INTO messages_fts(messages_fts, rowid, content) VALUES('delete', old.id, old.content);
+              INSERT INTO messages_fts(rowid, content) VALUES (new.id, new.content);
+            END;
+          SQL
+
+          # Populate FTS with existing messages
+          @db.execute("INSERT INTO messages_fts(rowid, content) SELECT id, content FROM messages WHERE content IS NOT NULL")
+        end
+
+        if from_version < 4
+          # Add thread/branching support columns
+          @db.execute_batch(<<~SQL)
+            ALTER TABLE sessions ADD COLUMN parent_session_id TEXT;
+            ALTER TABLE sessions ADD COLUMN parent_message_id INTEGER;
+            ALTER TABLE sessions ADD COLUMN parent_po TEXT;
+            ALTER TABLE sessions ADD COLUMN thread_type TEXT DEFAULT 'root';
+            CREATE INDEX IF NOT EXISTS idx_sessions_parent ON sessions(parent_session_id);
+          SQL
+        end
+      end
+
+      def parse_session_row(row, include_count: false)
+        result = {
+          id: row["id"],
+          po_name: row["po_name"],
+          name: row["name"],
+          source: row["source"] || "tui",
+          source_client: row["source_client"],
+          last_message_source: row["last_message_source"],
+          created_at: row["created_at"] ? Time.parse(row["created_at"]) : nil,
+          updated_at: row["updated_at"] ? Time.parse(row["updated_at"]) : nil,
+          metadata: row["metadata"] ? JSON.parse(row["metadata"], symbolize_names: true) : {},
+          # Thread fields (v4)
+          parent_session_id: row["parent_session_id"],
+          parent_message_id: row["parent_message_id"],
+          parent_po: row["parent_po"],
+          thread_type: row["thread_type"] || "root"
+        }
+        result[:message_count] = row["message_count"] if include_count && row["message_count"]
+        result
+      end
+
+      def parse_message_row(row)
+        {
+          id: row["id"],
+          session_id: row["session_id"],
+          role: row["role"].to_sym,
+          content: row["content"],
+          from_po: row["from_po"],
+          tool_calls: row["tool_calls"] ? JSON.parse(row["tool_calls"], symbolize_names: true) : nil,
+          tool_results: row["tool_results"] ? JSON.parse(row["tool_results"], symbolize_names: true) : nil,
+          created_at: row["created_at"] ? Time.parse(row["created_at"]) : nil
+        }
+      end
+    end
+  end
+end