tor-browser

ChatStore.sys.mjs (23676B)

---

/*
This Source Code Form is subject to the terms of the Mozilla Public
* License, v. 2.0. If a copy of the MPL was not distributed with this
* file, You can obtain one at https://mozilla.org/MPL/2.0/. */

const lazy = {};

ChromeUtils.defineESModuleGetters(lazy, {
 Sqlite: "resource://gre/modules/Sqlite.sys.mjs",
});

ChromeUtils.defineLazyGetter(lazy, "log", function () {
 return console.createInstance({
   prefix: "ChatStore",
   maxLogLevelPref: "browser.aiwindow.chatStore.loglevel",
 });
});

import {
 CONVERSATION_TABLE,
 CONVERSATION_UPDATED_DATE_INDEX,
 CONVERSATION_INSERT,
 MESSAGE_TABLE,
 MESSAGE_ORDINAL_INDEX,
 MESSAGE_URL_INDEX,
 MESSAGE_CREATED_DATE_INDEX,
 MESSAGE_CONV_ID_INDEX,
 MESSAGE_INSERT,
 CONVERSATIONS_MOST_RECENT,
 CONVERSATION_BY_ID,
 CONVERSATIONS_BY_DATE,
 CONVERSATIONS_BY_URL,
 CONVERSATIONS_CONTENT_SEARCH,
 CONVERSATIONS_CONTENT_SEARCH_BY_ROLE,
 CONVERSATIONS_HISTORY_SEARCH,
 MESSAGES_BY_DATE,
 MESSAGES_BY_DATE_AND_ROLE,
 DELETE_CONVERSATION_BY_ID,
 CONVERSATIONS_OLDEST,
 CONVERSATION_HISTORY,
 ESCAPE_CHAR,
 getConversationMessagesSql,
} from "./ChatSql.sys.mjs";

import { ChatMinimal } from "./ChatMessage.sys.mjs";

export { ChatConversation } from "./ChatConversation.sys.mjs";
export { ChatMessage, ChatMinimal } from "./ChatMessage.sys.mjs";
export {
 CONVERSATION_STATUS,
 MESSAGE_ROLE,
 INSIGHTS_FLAG_SOURCE,
} from "./ChatConstants.sys.mjs";

import {
 CURRENT_SCHEMA_VERSION,
 DB_FOLDER_PATH,
 DB_FILE_NAME,
 PREF_BRANCH,
 CONVERSATION_STATUS,
} from "./ChatConstants.sys.mjs";

import {
 parseConversationRow,
 parseMessageRows,
 parseChatHistoryViewRows,
 toJSONOrNull,
} from "./ChatUtils.sys.mjs";

// NOTE: Reference to migrations file, migrations.mjs has an example
// migration function set up for a migration, and the eslint-disable-next-line
// should be removed once we create the first migration.
//
// eslint-disable-next-line no-unused-vars
import { migrations } from "./ChatMigrations.sys.mjs";

const MAX_DB_SIZE_BYTES = 75 * 1024 * 1024;
const SORTS = ["ASC", "DESC"];

/**
* Simple interface to store and retrieve chat conversations and messages.
*
* @todo Bug 2005409
* Move this documentation to Firefox source docs
*
* See: https://docs.google.com/document/d/1VlwmGbMhPIe-tmeKWinHuPh50VC9QrWEeQQ5V-UvEso/edit?tab=t.klqqibndv3zk
*
* @example
* let { ChatStore, ChatConversation, ChatMessage, MESSAGE_ROLE } =
*   ChromeUtils.importESModule("resource:///modules/aiwindow/ui/modules/ChatStore.sys.mjs");
* const chatStore = new ChatStore();
* const conversation = new ChatConversation({
*   title: "title",
*   description: "description",
*   pageUrl: new URL("https://mozilla.com/"),
*   pageMeta: { one: 1, two: 2 },
* });
* const msg1 = new ChatMessage({
*   ordinal: 0,
*   role: MESSAGE_ROLE.USER,
*   modelId: "test",
*   params: { one: "one" },
*   usage: { two: "two", content: "some content" },
* });
* const msg2 = new ChatMessage({
*   ordinal: 1,
*   role: MESSAGE_ROLE.ASSISTANT,
*   modelId: "test",
*   params: { one: "one" },
*   usage: { two: "two", content: "some content 2" },
* });
* conversation.messages = [msg1, msg2];
* await chatStore.updateConversation(conversation);
* // Or findConversationsByDate, findConversationsByURL.
* const foundConversation =
*   await chatStore.findConversationById(conversation.id);
*
* @typedef {object} ChatStore
*
* @property {*} x ?
*/
export class ChatStore {
 #asyncShutdownBlocker;
 #conn;
 #promiseConn;

 constructor() {
   this.#asyncShutdownBlocker = async () => {
     await this.#closeConnection();
   };
 }

 /**
  * Updates a conversation's saved state in the SQLite db
  *
  * @param {ChatConversation} conversation
  */
 async updateConversation(conversation) {
   await this.#ensureDatabase().catch(e => {
     lazy.log.error("Could not ensure a database connection.");
     throw e;
   });

   const pageUrl = URL.parse(conversation.pageUrl);

   await this.#conn
     .executeTransaction(async () => {
       await this.#conn.executeCached(CONVERSATION_INSERT, {
         conv_id: conversation.id,
         title: conversation.title,
         description: conversation.description,
         page_url: pageUrl?.href ?? null,
         page_meta: toJSONOrNull(conversation.pageMeta),
         created_date: conversation.createdDate,
         updated_date: conversation.updatedDate,
         status: conversation.status,
         active_branch_tip_message_id: conversation.activeBranchTipMessageId,
       });

       const messages = conversation.messages.map(m => ({
         message_id: m.id,
         conv_id: conversation.id,
         created_date: m.createdDate,
         parent_message_id: m.parentMessageId,
         revision_root_message_id: m.revisionRootMessageId,
         ordinal: m.ordinal,
         is_active_branch: m.isActiveBranch ? 1 : 0,
         role: m.role,
         model_id: m.modelId,
         params: toJSONOrNull(m.params),
         content: toJSONOrNull(m.content),
         usage: toJSONOrNull(m.usage),
         page_url: m.pageUrl?.href || "",
         turn_index: m.turnIndex,
         insights_enabled: m.insightsEnabled,
         insights_flag_source: m.insightsFlagSource,
         insights_applied_jsonb: toJSONOrNull(m.insightsApplied),
         web_search_queries_jsonb: toJSONOrNull(m.webSearchQueries),
       }));
       await this.#conn.executeCached(MESSAGE_INSERT, messages);
     })
     .catch(e => {
       lazy.log.error("Transaction failed to execute");
       throw e;
     });
 }

 /**
  * Gets a list of oldest conversations
  *
  * @param {number} numberOfConversations - How many conversations to retrieve
  * @returns {Array<ChatMinimal>} - List of ChatMinimal items
  */
 async findOldestConversations(numberOfConversations) {
   await this.#ensureDatabase().catch(e => {
     lazy.log.error("Could not ensure a database connection.");
     throw e;
   });

   const rows = await this.#conn
     .executeCached(CONVERSATIONS_OLDEST, {
       limit: numberOfConversations,
     })
     .catch(e => {
       lazy.log.error("Could not retrieve oldest conversations.");
       throw e;
     });

   return rows.map(row => {
     return new ChatMinimal({
       convId: row.getResultByName("conv_id"),
       title: row.getResultByName("title"),
     });
   });
 }

 /**
  * Gets a list of most recent conversations
  *
  * @param {number} numberOfConversations - How many conversations to retrieve
  * @returns {Array<ChatMinimal>} - List of ChatMinimal items
  */
 async findRecentConversations(numberOfConversations) {
   await this.#ensureDatabase().catch(e => {
     lazy.log.error("Could not ensure a database connection.");
     throw e;
   });

   const rows = await this.#conn
     .executeCached(CONVERSATIONS_MOST_RECENT, {
       limit: numberOfConversations,
     })
     .catch(e => {
       lazy.log.error("Could not retrieve most recent conversations.");
       throw e;
     });

   return rows.map(row => {
     return new ChatMinimal({
       convId: row.getResultByName("conv_id"),
       title: row.getResultByName("title"),
     });
   });
 }

 /**
  * Gets a Conversation using it's id
  *
  * @param {string} conversationId - The ID of the conversation to retrieve
  *
  * @returns {ChatConversation} - The conversation and its messages
  */
 async findConversationById(conversationId) {
   const conversations = await this.#findConversationsWithMessages(
     CONVERSATION_BY_ID,
     {
       conv_id: conversationId,
     }
   );

   return conversations[0] ?? null;
 }

 /**
  * Finds conversations between a specified start and end date
  *
  * @param {number} startDate - Start time epoch format
  * @param {number} endDate - End time epoch format
  *
  * @returns {Array<ChatConversation>} - The conversations and their messages
  */
 async findConversationsByDate(startDate, endDate) {
   return this.#findConversationsWithMessages(CONVERSATIONS_BY_DATE, {
     start_date: startDate,
     end_date: endDate,
   });
 }

 /**
  * Finds conversations between a specified start and end date
  *
  * @param {URL} pageUrl - The URL to find conversations for
  *
  * @returns {Array<ChatConversation>} - The conversations and their messages
  */
 async findConversationsByURL(pageUrl) {
   return this.#findConversationsWithMessages(CONVERSATIONS_BY_URL, {
     page_url: pageUrl.href,
   });
 }

 /**
  * Search for messages that happened between the specified start
  * and end dates, optionally, filter the messages by a specific
  * message role type.
  *
  * @param {Date} startDate - The start date, inclusive
  * @param {Date} [endDate=new Date()] - The end date, inclusive
  * @param {MessageRole} [role=-1] - The message role type to filter by, one of 0|1|2|3
  * as defined by the constant MESSAGE_ROLE
  * @param {number} [limit=-1] - The max number of messages to retrieve
  * @param {number} [offset=-1] - The number or messages to skip from the result set
  *
  * @returns {Array<ChatMessage>} - An array of ChatMessage entries
  */
 async findMessagesByDate(
   startDate,
   endDate = new Date(),
   role = -1,
   limit = -1,
   offset = -1
 ) {
   const params = {
     start_date: startDate.getTime(),
     end_date: endDate.getTime(),
     limit,
     offset,
   };

   let sql = MESSAGES_BY_DATE;
   if (role > -1) {
     sql = MESSAGES_BY_DATE_AND_ROLE;
     params.role = role;
   }

   let rows = await this.#conn.executeCached(sql, params);

   return parseMessageRows(rows);
 }

 #escapeForLike(searchString) {
   return searchString
     .replaceAll(ESCAPE_CHAR, `${ESCAPE_CHAR}${ESCAPE_CHAR}`)
     .replaceAll("%", `${ESCAPE_CHAR}%`)
     .replaceAll("_", `${ESCAPE_CHAR}_`);
 }

 /**
  * Searches through the message.content JSON object to find a particular
  * object path that contains a partial string match of a value.
  *
  * @param {string} keyChain - The object key chain to look through,
  * like obj.field1.field2
  * @param {MessageRole} [role=-1] - A message role to search for
  *
  * @returns {Array<ChatConversation>} - An array of conversations with messages
  * that contain a message that matches the search string at the given content
  * object path
  */
 async searchContent(keyChain, role = -1) {
   const path = `$.${keyChain}`;

   const query =
     role > -1
       ? CONVERSATIONS_CONTENT_SEARCH_BY_ROLE
       : CONVERSATIONS_CONTENT_SEARCH;

   const params = { path };

   if (role > -1) {
     params.role = role;
   }

   const rows = await this.#conn.executeCached(query, params);

   if (!rows.length) {
     return [];
   }

   const conversations = rows.map(parseConversationRow);

   return await this.#getMessagesForConversations(conversations);
 }

 /**
  * Searches for conversations where the conversation title, or the conversation
  * contains a user message where the search string contains a partial match
  * in the message.content.body field
  *
  * @param {string} searchString - The string to search with for conversations
  *
  * @returns {Array<ChatConversation>} - An array of conversations with messages
  * that contain a message that matches the search string in the conversation
  * titles
  */
 async search(searchString) {
   const path = `$.body`;
   const pattern = `%${this.#escapeForLike(searchString)}%`;

   const rows = await this.#conn.executeCached(CONVERSATIONS_HISTORY_SEARCH, {
     path,
     pattern,
   });

   if (!rows.length) {
     return [];
   }

   const conversations = rows.map(parseConversationRow);

   return await this.#getMessagesForConversations(conversations);
 }

 /**
  * Gets a list of chat history items to display in Chat History view.
  *
  * @param {number} [pageNumber=1] - The page number to get, 1 based indexing
  * @param {number} [pageSize=20] - Number of items to get per page
  * @param {string} [sort="desc"] - desc|asc The sorting order based on updated_date for conversations
  */
 async chatHistoryView(pageNumber = 1, pageSize = 20, sort = "desc") {
   const sorting = SORTS.find(item => item === sort.toUpperCase()) ?? "DESC";
   const offset = pageSize * (pageNumber - 1);
   const limit = pageSize;
   const params = {
     limit,
     offset,
   };

   const rows = await this.#conn.executeCached(
     CONVERSATION_HISTORY.replace("{sort}", sorting),
     params
   );

   return parseChatHistoryViewRows(rows);
 }

 /**
  * Prunes the database of old conversations in order to get the
  * database file size to the specified maximum size.
  *
  * @todo Bug 2005411
  * Review the requirements for db pruning and set up invocation schedule, and refactor
  * to use dbstat
  *
  * @param {number} [reduceByPercentage=0.05] - Percentage to reduce db file size by
  * @param {number} [maxDbSizeBytes=MAX_DB_SIZE_BYTES] - Db max file size
  */
 async pruneDatabase(
   reduceByPercentage = 0.05,
   maxDbSizeBytes = MAX_DB_SIZE_BYTES
 ) {
   if (!IOUtils.exists(this.databaseFilePath)) {
     return;
   }

   const DELETE_BATCH_SIZE = 50;

   const getPragmaInt = async name => {
     const result = await this.#conn.execute(`PRAGMA ${name}`);
     return result[0].getInt32(0);
   };

   // compute the logical DB size in bytes using SQLite's page_size,
   // page_count, and freelist_count
   const getLogicalDbSizeBytes = async () => {
     const pageSize = await getPragmaInt("page_size");
     const pageCount = await getPragmaInt("page_count");
     const freelistCount = await getPragmaInt("freelist_count");

     // Logical used pages = total pages - free pages
     const usedPages = pageCount - freelistCount;
     const lSize = usedPages * pageSize;

     return lSize;
   };

   let logicalSize = await getLogicalDbSizeBytes();
   if (logicalSize < maxDbSizeBytes) {
     return;
   }

   const targetLogicalSize = Math.max(
     0,
     logicalSize * (1 - reduceByPercentage)
   );

   const MAX_ITERATIONS = 100;
   // how many "no file size change" batches we tolerate
   const MAX_STAGNANT = 5;
   let iterations = 0;
   let stagnantIterations = 0;

   while (
     logicalSize > targetLogicalSize &&
     iterations < MAX_ITERATIONS &&
     stagnantIterations < MAX_STAGNANT
   ) {
     iterations++;

     const recentChats = await this.findOldestConversations(DELETE_BATCH_SIZE);

     if (!recentChats.length) {
       break;
     }

     for (const chat of recentChats) {
       await this.deleteConversationById(chat.id);
     }

     const newLogicalSize = await getLogicalDbSizeBytes();
     if (newLogicalSize >= logicalSize) {
       stagnantIterations++;
     } else {
       stagnantIterations = 0;
     }

     logicalSize = newLogicalSize;
   }

   // Actually reclaim disk space.
   await this.#conn.execute("PRAGMA incremental_vacuum;");
 }

 /**
  * Returns the file size of the database.
  * Establishes a connection first to make sure the
  * database exists.
  *
  * @returns {number} - The file size in bytes
  */
 async getDatabaseSize() {
   await this.#ensureDatabase();

   const stats = await IOUtils.stat(this.databaseFilePath);
   return stats.size;
 }

 /**
  * Deletes a particular conversation using it's id
  *
  * @param {string} id - The conv_id of a conversation row to delete
  */
 async deleteConversationById(id) {
   await this.#ensureDatabase();

   await this.#conn.execute(DELETE_CONVERSATION_BY_ID, {
     conv_id: id,
   });
 }

 /**
  * This method is meant to only be used for testing cleanup
  */
 async destroyDatabase() {
   await this.#removeDatabaseFiles();
 }

 /**
  * Gets the version of the schema currently set in the database.
  *
  * @returns {number}
  */
 async getDatabaseSchemaVersion() {
   if (!this.#conn) {
     await this.#ensureDatabase();
   }

   return this.#conn.getSchemaVersion();
 }

 async #getMessagesForConversations(conversations) {
   const convs = conversations.reduce((convMap, conv) => {
     convMap[conv.id] = conv;

     return convMap;
   }, {});

   // Find all the messages for all the conversations.
   const rows = await this.#conn
     .executeCached(
       getConversationMessagesSql(conversations.length),
       conversations.map(c => c.id)
     )
     .catch(e => {
       lazy.log.error("Could not retrieve messages for conversatons");
       lazy.log.error(`${e.message}\n${e.stack}`);

       return [];
     });

   // TODO: retrieve TTL content.

   parseMessageRows(rows).forEach(message => {
     const conversation = convs[message.convId];
     if (conversation) {
       conversation.messages.push(message);
     }
   });

   return conversations;
 }

 async #openConnection() {
   lazy.log.debug("Opening new connection");

   try {
     const confConfig = { path: this.databaseFilePath };
     this.#conn = await lazy.Sqlite.openConnection(confConfig);
   } catch (e) {
     lazy.log.error("openConnection() could not open db:", e.message);
     throw e;
   }

   lazy.Sqlite.shutdown.addBlocker(
     "ChatStore: Shutdown",
     this.#asyncShutdownBlocker
   );

   try {
     // TODO: remove this after switching pruneDatabase() to use dbstat
     await this.#conn.execute("PRAGMA page_size = 4096;");
     // Setup WAL journaling, as it is generally faster.
     await this.#conn.execute("PRAGMA journal_mode = WAL;");
     await this.#conn.execute("PRAGMA wal_autocheckpoint = 16;");

     // Store VACUUM information to be used by the VacuumManager.
     await this.#conn.execute("PRAGMA auto_vacuum = INCREMENTAL;");
     await this.#conn.execute("PRAGMA foreign_keys = ON;");
   } catch (e) {
     lazy.log.warn("Configuring SQLite PRAGMA settings: ", e.message);
   }
 }

 async #closeConnection() {
   if (!this.#conn) {
     return;
   }

   lazy.log.debug("Closing connection");
   lazy.Sqlite.shutdown.removeBlocker(this.#asyncShutdownBlocker);
   try {
     await this.#conn.close();
   } catch (e) {
     lazy.log.warn(`Error closing connection: ${e.message}`);
   }
   this.#conn = null;
 }

 /**
  * @todo Bug 2005412
  * Discuss implications of multiple instances of ChatStore
  * and the potential issues with migrations/schemas.
  */
 async #ensureDatabase() {
   if (this.#promiseConn) {
     return this.#promiseConn;
   }

   let deferred = Promise.withResolvers();
   this.#promiseConn = deferred.promise;
   if (this.#removeDatabaseOnStartup) {
     lazy.log.debug("Removing database on startup");
     try {
       await this.#removeDatabaseFiles();
     } catch (e) {
       deferred.reject(new Error("Could not remove the database files"));
       return deferred.promise;
     }
   }

   try {
     await this.#openConnection();
   } catch (e) {
     if (
       e.result == Cr.NS_ERROR_FILE_CORRUPTED ||
       e.errors?.some(error => error.result == Ci.mozIStorageError.NOTADB)
     ) {
       lazy.log.warn("Invalid database detected, removing it.", e);
       await this.#removeDatabaseFiles();
     }
   }

   if (!this.#conn) {
     try {
       await this.#openConnection();
     } catch (e) {
       lazy.log.error("Could not open the database connection.", e);
       deferred.reject(new Error("Could not open the database connection"));
       return deferred.promise;
     }
   }

   try {
     await this.#initializeSchema();
   } catch (e) {
     lazy.log.warn(
       "Failed to initialize the database schema, recreating the database.",
       e
     );
     // If the schema cannot be initialized try to create a new database file.
     await this.#removeDatabaseFiles();
   }

   deferred.resolve(this.#conn);
   return this.#promiseConn;
 }

 async setSchemaVersion(version) {
   await this.#conn.setSchemaVersion(version);
 }

 async #initializeSchema() {
   const version = await this.getDatabaseSchemaVersion();

   if (version == this.CURRENT_SCHEMA_VERSION) {
     return;
   }

   if (version > this.CURRENT_SCHEMA_VERSION) {
     await this.setSchemaVersion(this.CURRENT_SCHEMA_VERSION);
     return;
   }

   // Must migrate the schema.
   await this.#conn.executeTransaction(async () => {
     if (version == 0) {
       // This is a newly created database, just create the entities.
       await this.#createDatabaseEntities();
       await this.#conn.setSchemaVersion(this.CURRENT_SCHEMA_VERSION);
       // eslint-disable-next-line no-useless-return
       return;
     }

     await this.applyMigrations();
     await this.setSchemaVersion(this.CURRENT_SCHEMA_VERSION);
   });
 }

 async applyMigrations() {
   for (const migration of migrations) {
     if (typeof migration !== "function") {
       continue;
     }

     await migration(this.#conn, this.CURRENT_SCHEMA_VERSION);
   }
 }

 async #removeDatabaseFiles() {
   lazy.log.debug("Removing database files");
   await this.#closeConnection();
   try {
     for (let file of [
       this.databaseFilePath,
       PathUtils.join(DB_FOLDER_PATH, this.databaseFileName + "-wal"),
       PathUtils.join(DB_FOLDER_PATH, this.databaseFileName + "-shm"),
     ]) {
       lazy.log.debug(`Removing ${file}`);
       await IOUtils.remove(file, {
         retryReadonly: true,
         recursive: true,
         ignoreAbsent: true,
       });
     }
     this.#removeDatabaseOnStartup = false;
   } catch (e) {
     lazy.log.warn("Failed to remove database files", e);
     // Try to clear on next startup.
     this.#removeDatabaseOnStartup = true;
     // Re-throw the exception for the caller.
     throw e;
   }
 }

 async #findConversationsWithMessages(sql, queryParams) {
   await this.#ensureDatabase().catch(e => {
     lazy.log.error("Could not ensure a database connection.");
     lazy.log.error(`${e.message}\n${e.stack}`);

     return [];
   });

   // @todo Bug 2005414
   // Check summary first, find the one with the largest end_ordinal.
   // If not found retrieve all messages.
   // If found compare end_ordinal of the summary with active branch ordinal
   // to determine if extra messages must be retrieved.
   let rows = await this.#conn.executeCached(sql, queryParams);

   const conversations = rows.map(parseConversationRow);

   return await this.#getMessagesForConversations(conversations);
 }

 async #createDatabaseEntities() {
   await this.#conn.execute(CONVERSATION_TABLE);
   await this.#conn.execute(CONVERSATION_UPDATED_DATE_INDEX);
   await this.#conn.execute(MESSAGE_TABLE);
   await this.#conn.execute(MESSAGE_ORDINAL_INDEX);
   await this.#conn.execute(MESSAGE_URL_INDEX);
   await this.#conn.execute(MESSAGE_CREATED_DATE_INDEX);
   await this.#conn.execute(MESSAGE_CONV_ID_INDEX);
 }

 get #removeDatabaseOnStartup() {
   return Services.prefs.getBoolPref(
     `${PREF_BRANCH}.removeDatabaseOnStartup`,
     false
   );
 }

 set #removeDatabaseOnStartup(value) {
   lazy.log.debug(`Setting removeDatabaseOnStartup to ${value}`);
   Services.prefs.setBoolPref(`${PREF_BRANCH}.removeDatabaseOnStartup`, value);
 }

 static get CONVERSATION_STATUS() {
   return CONVERSATION_STATUS;
 }

 get CURRENT_SCHEMA_VERSION() {
   return CURRENT_SCHEMA_VERSION;
 }

 get connection() {
   return this.#conn;
 }

 get databaseFileName() {
   return DB_FILE_NAME;
 }

 get databaseFilePath() {
   return PathUtils.join(PathUtils.profileDir, this.databaseFileName);
 }
}