commit 7793bcad62947819304c839402dd546491b30a79
parent e63b083d46f9cef1794222f77d13b1dc6124eaa9
Author: Omar Gonzalez <s9tpepper@apache.org>
Date: Fri, 12 Dec 2025 05:11:42 +0000
Bug 2000961 - ChatStore module for ai window persistence r=mak,Mardak,ai-frontend-reviewers
Differential Revision: https://phabricator.services.mozilla.com/D274644
Diffstat:
16 files changed, 3377 insertions(+), 0 deletions(-)
diff --git a/browser/app/profile/firefox.js b/browser/app/profile/firefox.js
@@ -2246,6 +2246,7 @@ pref("browser.ml.smartAssist.overrideNewTab", false);
// AI Window Feature
pref("browser.aiwindow.enabled", false);
+pref("browser.aiwindow.chatStore.loglevel", "Error");
// Block insecure active content on https pages
pref("security.mixed_content.block_active_content", true);
diff --git a/browser/base/content/test/static/browser_all_files_referenced.js b/browser/base/content/test/static/browser_all_files_referenced.js
@@ -335,6 +335,10 @@ var allowlist = [
{
file: "moz-src:///browser/components/aiwindow/models/IntentClassifier.sys.mjs",
},
+ // Bug 2000961 - Add ChatStore.sys.mjs module
+ {
+ file: "moz-src:///browser/components/aiwindow/ui/modules/ChatStore.sys.mjs",
+ },
// Bug 2003598 - Add Chat service with fetch with history (backed out due to unused file)
{
file: "moz-src:///browser/components/aiwindow/models/Chat.sys.mjs",
diff --git a/browser/components/aiwindow/ui/modules/ChatConstants.sys.mjs b/browser/components/aiwindow/ui/modules/ChatConstants.sys.mjs
@@ -0,0 +1,67 @@
+/*
+ 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/. */
+
+/**
+ * The current SQLite database schema version
+ */
+export const CURRENT_SCHEMA_VERSION = 1;
+
+/**
+ * The directory that the SQLite database lives in
+ */
+export const DB_FOLDER_PATH = PathUtils?.profileDir ?? "./";
+
+/**
+ * The name of the SQLite database file
+ */
+export const DB_FILE_NAME = "chat-store.sqlite";
+
+/**
+ * Preference branch for the Chat storage location
+ */
+export const PREF_BRANCH = "browser.aiWindow.chatHistory";
+
+/**
+ * @typedef ConversationStatus
+ * @property {number} ACTIVE - An active conversation
+ * @property {number} ARCHIVE - An archived conversation
+ * @property {number} DELETED - A deleted conversation
+ */
+
+/**
+ * @type {ConversationStatus}
+ */
+export const CONVERSATION_STATUS = Object.freeze({
+ ACTIVE: 0,
+ ARCHIVED: 1,
+ DELETED: 2,
+});
+
+/**
+ * @typedef {0 | 1 | 2 | 3} MessageRole
+ */
+
+/**
+ * @enum {MessageRole}
+ */
+export const MESSAGE_ROLE = Object.freeze({
+ USER: 0,
+ ASSISTANT: 1,
+ SYSTEM: 2,
+ TOOL: 3,
+});
+
+/**
+ * @typedef {0 | 1 | 2} InsightsFlagSource
+ */
+
+/**
+ * @type {InsightsFlagSource}
+ */
+export const INSIGHTS_FLAG_SOURCE = Object.freeze({
+ GLOBAL: 0,
+ CONVERSATION: 1,
+ MESSAGE_ONCE: 2,
+});
diff --git a/browser/components/aiwindow/ui/modules/ChatConversation.sys.mjs b/browser/components/aiwindow/ui/modules/ChatConversation.sys.mjs
@@ -0,0 +1,248 @@
+/*
+ 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/. */
+
+import { makeGuid } from "./ChatUtils.sys.mjs";
+import { CONVERSATION_STATUS, MESSAGE_ROLE } from "./ChatConstants.sys.mjs";
+import {
+ AssistantRoleOpts,
+ ChatMessage,
+ ToolRoleOpts,
+ UserRoleOpts,
+} from "./ChatMessage.sys.mjs";
+
+/**
+ * A conversation containing messages.
+ */
+export class ChatConversation {
+ id;
+ title;
+ description;
+ pageUrl;
+ pageMeta;
+ createdDate;
+ updatedDate;
+ status;
+ #messages;
+ activeBranchTipMessageId;
+
+ /**
+ * @param {object} params
+ * @param {string} [params.id]
+ * @param {string} params.title
+ * @param {string} params.description
+ * @param {URL} params.pageUrl
+ * @param {object} params.pageMeta
+ * @param {number} [params.createdDate]
+ * @param {number} [params.updatedDate]
+ * @param {CONVERSATION_STATUS} [params.status]
+ * @param {Array<ChatMessage>} [params.messages]
+ */
+ constructor(params) {
+ const {
+ id = makeGuid(),
+ title,
+ description,
+ pageUrl,
+ pageMeta,
+ createdDate = Date.now(),
+ updatedDate = Date.now(),
+ messages = [],
+ } = params;
+
+ this.id = id;
+ this.title = title;
+ this.description = description;
+ this.pageUrl = pageUrl;
+ this.pageMeta = pageMeta;
+ this.createdDate = createdDate;
+ this.updatedDate = updatedDate;
+ this.#messages = messages;
+
+ // NOTE: Destructuring params.status causes a linter error
+ this.status = params.status || CONVERSATION_STATUS.ACTIVE;
+ }
+
+ /**
+ * Adds a message to the conversation
+ *
+ * @param {ConversationRole} role - The type of conversation message
+ * @param {object} content - The conversation message contents
+ * @param {URL} pageUrl - The current page url when message was submitted
+ * @param {number} turnIndex - The current conversation turn/cycle
+ * @param {AssistantRoleOpts|ToolRoleOpts|UserRoleOpts} opts - Additional opts for the message
+ */
+ addMessage(role, content, pageUrl, turnIndex, opts = {}) {
+ if (role < 0 || role > MESSAGE_ROLE.TOOL) {
+ return;
+ }
+
+ if (turnIndex < 0) {
+ turnIndex = 0;
+ }
+
+ let parentMessageId = null;
+ if (this?.messages?.length) {
+ const lastMessageIndex = this.messages.length - 1;
+ parentMessageId = this.messages[lastMessageIndex].id;
+ }
+
+ const convId = this.id;
+ const currentMessages = this?.messages || [];
+ const ordinal = currentMessages.length ? currentMessages.length + 1 : 1;
+
+ const message_data = {
+ parentMessageId,
+ content,
+ ordinal,
+ pageUrl,
+ turnIndex,
+ role,
+ convId,
+ ...opts,
+ };
+
+ const newMessage = new ChatMessage(message_data);
+
+ this.messages.push(newMessage);
+ }
+
+ /**
+ * Add a user message to the conversation
+ *
+ * @todo Bug 2005424
+ * Limit/filter out data uris from message data
+ *
+ * @param {string} contentBody - The user message content
+ * @param {string?} [pageUrl=""] - The current page url when message was submitted
+ * @param {number?} [turnIndex=0] - The conversation turn/cycle
+ * @param {UserRoleOpts} [userOpts=new UserRoleOpts()] - User message options
+ */
+ addUserMessage(
+ contentBody,
+ pageUrl = "",
+ turnIndex = 0,
+ userOpts = new UserRoleOpts()
+ ) {
+ const content = {
+ type: "text",
+ body: contentBody,
+ };
+
+ let url = URL.parse(pageUrl);
+
+ this.addMessage(MESSAGE_ROLE.USER, content, url, turnIndex, userOpts);
+ }
+
+ /**
+ * Add an assistant message to the conversation
+ *
+ * @param {string} type - The assistant message type: text|function
+ * @param {string} contentBody - The assistant message content
+ * @param {number} turnIndex - The current conversation turn/cycle
+ * @param {AssistantRoleOpts} [assistantOpts=new AssistantRoleOpts()] - ChatMessage options specific to assistant messages
+ */
+ addAssistantMessage(
+ type,
+ contentBody,
+ turnIndex,
+ assistantOpts = new AssistantRoleOpts()
+ ) {
+ const content = {
+ type: "text",
+ body: contentBody,
+ };
+
+ this.addMessage(
+ MESSAGE_ROLE.ASSISTANT,
+ content,
+ "",
+ turnIndex,
+ assistantOpts
+ );
+ }
+
+ /**
+ * Add a tool call message to the conversation
+ *
+ * @param {object} content - The tool call object to be saved as JSON
+ * @param {number} turnIndex - The current conversation turn/cycle
+ * @param {ToolRoleOpts} [toolOpts=new ToolRoleOpts()] - Message opts for a tool role message
+ */
+ addToolCallMessage(content, turnIndex, toolOpts = new ToolRoleOpts()) {
+ this.addMessage(MESSAGE_ROLE.TOOL, content, "", turnIndex, toolOpts);
+ }
+
+ /**
+ * Add a system message to the conversation
+ *
+ * @param {string} type - The assistant message type: text|injected_insights|injected_real_time_info
+ * @param {string} contentBody - The system message object to be saved as JSON
+ * @param {number} turnIndex - The current conversation turn/cycle
+ */
+ addSystemMessage(type, contentBody, turnIndex) {
+ const content = { type, body: contentBody };
+
+ this.addMessage(MESSAGE_ROLE.SYSTEM, content, "", turnIndex);
+ }
+
+ /**
+ * Retrieves the list of visited sites during a conversation in visited order.
+ * Primarily used to retrieve external URLs that the user had a conversation
+ * around to display in Chat History view.
+ *
+ * @param {boolean} [includeInternal=false] - Whether to include internal Firefox URLs
+ *
+ * @returns {Array<URL>} - Ordered list of visited page URLs for this conversation
+ */
+ getSitesList(includeInternal = false) {
+ const seen = new Set();
+ const deduped = [];
+
+ this.messages.forEach(message => {
+ if (!message.pageUrl) {
+ return;
+ }
+
+ if (!includeInternal && !message.pageUrl.protocol.startsWith("http")) {
+ return;
+ }
+
+ if (!seen.has(message.pageUrl.href)) {
+ seen.add(message.pageUrl.href);
+ deduped.push(message.pageUrl);
+ }
+ });
+
+ return deduped;
+ }
+
+ /**
+ * Returns the most recently visited external sites during this conversation, or null
+ * if no external sites have been visited.
+ *
+ * @returns {URL|null}
+ */
+ getMostRecentPageVisited() {
+ const sites = this.getSitesList();
+
+ return sites.length ? sites.pop() : null;
+ }
+
+ #updateActiveBranchTipMessageId() {
+ this.activeBranchTipMessageId = this.messages
+ .filter(m => m.isActiveBranch)
+ .sort((a, b) => b.ordinal - a.ordinal)
+ .shift()?.id;
+ }
+
+ set messages(value) {
+ this.#messages = value;
+ this.#updateActiveBranchTipMessageId();
+ }
+
+ get messages() {
+ return this.#messages;
+ }
+}
diff --git a/browser/components/aiwindow/ui/modules/ChatMessage.sys.mjs b/browser/components/aiwindow/ui/modules/ChatMessage.sys.mjs
@@ -0,0 +1,213 @@
+/*
+ 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/. */
+
+import { makeGuid } from "./ChatUtils.sys.mjs";
+
+/**
+ * A message in a conversation.
+ */
+export class ChatMessage {
+ id;
+ createdDate;
+ parentMessageId;
+ revisionRootMessageId;
+ ordinal;
+ isActiveBranch;
+ role;
+ modelId;
+ params;
+ usage;
+ content;
+ convId;
+ pageUrl;
+ turnIndex;
+ insightsEnabled;
+ insightsFlagSource;
+ insightsApplied;
+ webSearchQueries;
+
+ /**
+ * @param {object} param
+ * @param {number} param.ordinal - The order of the message
+ * @param {MessageRole} param.role - The message role
+ * @param {object} param.content - The message content object
+ * @param {number} param.turnIndex - The message turn, different than ordinal,
+ * prompt/reply for example would be one turn
+ * @param {string} [param.pageUrl = null] - A URL object defining which page
+ * the user was on when submitting a message if role == user
+ * @param {string} [param.id = makeGuid()] - The row.message_id of the
+ * message in the database
+ * @param {number} [param.createdDate = Date.now()] - The date the message was
+ * sent/stored in the database
+ * @param {string} [param.parentMessageId = null] - The id of the message
+ * which came before this message when it was added to the conversation,
+ * null if its the first message of the converation
+ * @param {string} [param.convId = null] - The id of the conversation the
+ * message belongs to
+ * @param {?boolean} param.insightsEnabled - Whether insights were enabled
+ * when the message was submitted if role == assistant
+ * @param {InsightsFlagSource} param.insightsFlagSource - How the
+ * insightsEnabled flag was determined if role == assistant, one of
+ * INSIGHTS_FLAG_SOURCE.GLOBAL, INSIGHTS_FLAG_SOURCE.CONVERSATION,
+ * INSIGHTS_FLAG_SOURCE.MESSAGE_ONCE
+ * @param {?Array<string>} param.insightsApplied - List of strings of insights
+ * that were applied to a response if insightsEnabled == true
+ * @param {?Array<string>} param.webSearchQueries - List of strings of web
+ * search queries that were applied to a response if role == assistant
+ * @param {object} [param.params = null] - Model params used if role == assistant|tool
+ * @param {object} [param.usage = null] - Token usage data for the current
+ * response if role == assistant
+ * @param {string} [param.modelId = null] - The model used for content
+ * generation if role == assistant|tool
+ * @param {string} [param.revisionRootMessageId = id] - Reference to the root
+ * of this branch, which ID a message branched from. Should be set to the
+ * same value as id when a message is first created. If a message is
+ * edited/regenerated revisionRootMessageId should remain the same for
+ * subsequent edits/regenerations, the id would diverge for subsequent
+ * edits/regenerations.
+ * @param {boolean} [param.isActiveBranch = true] - Defaults to true when a
+ * message is originally generated. If a message is edited/regenerated, the
+ * edited message turns to false and the newly edited/regenerated message is
+ * the only message of the revision branch set to true.
+ */
+ constructor({
+ ordinal,
+ role,
+ content,
+ turnIndex,
+ pageUrl = null,
+ id = makeGuid(),
+ createdDate = Date.now(),
+ parentMessageId = null,
+ convId = null,
+ insightsEnabled = null,
+ insightsFlagSource = null,
+ insightsApplied = null,
+ webSearchQueries = null,
+ params = null,
+ usage = null,
+ modelId = null,
+ revisionRootMessageId = id,
+ isActiveBranch = true,
+ }) {
+ this.id = id;
+ this.createdDate = createdDate;
+ this.parentMessageId = parentMessageId;
+ this.revisionRootMessageId = revisionRootMessageId;
+ this.isActiveBranch = isActiveBranch;
+ this.ordinal = ordinal;
+ this.role = role;
+ this.modelId = modelId;
+ this.params = params;
+ this.usage = usage;
+ this.content = content;
+ this.convId = convId;
+ this.pageUrl = pageUrl ? new URL(pageUrl) : null;
+ this.turnIndex = turnIndex;
+ this.insightsEnabled = insightsEnabled;
+ this.insightsFlagSource = insightsFlagSource;
+ this.insightsApplied = insightsApplied;
+ this.webSearchQueries = webSearchQueries;
+ }
+}
+
+/**
+ * Options required for a conversation message with
+ * role of assistant
+ */
+export class AssistantRoleOpts {
+ insightsEnabled;
+ insightsFlagSource;
+ insightsApplied;
+ webSearchQueries;
+ params;
+ usage;
+ modelId;
+
+ /**
+ * @param {string} [modelId=null]
+ * @param {object} [params=null] - The model params used
+ * @param {object} [usage=null] - Token usage data for the current response
+ * @param {boolean} [insightsEnabled=false] - Whether insights were enabled when the message was submitted
+ * @param {import("moz-src:///browser/components/aiwindow/ui/modules/ChatStorage.sys.mjs").InsightsFlagSource} [insightsFlagSource=null] - How the insightsEnabled flag was determined
+ * @param {?Array<string>} [insightsApplied=[]] - List of strings of insights that were applied to a response
+ * @param {?Array<string>} [webSearchQueries=[]] - List of strings of web search queries that were applied to a response
+ */
+ constructor(
+ modelId = null,
+ params = null,
+ usage = null,
+ insightsEnabled = false,
+ insightsFlagSource = null,
+ insightsApplied = [],
+ webSearchQueries = []
+ ) {
+ this.insightsEnabled = insightsEnabled;
+ this.insightsFlagSource = insightsFlagSource;
+ this.insightsApplied = insightsApplied;
+ this.webSearchQueries = webSearchQueries;
+ this.params = params;
+ this.usage = usage;
+ this.modelId = modelId;
+ }
+}
+
+/**
+ * Options required for a conversation message with
+ * role of assistant
+ */
+export class ToolRoleOpts {
+ modelId;
+
+ /**
+ * @param {string} [modelId=null]
+ */
+ constructor(modelId = null) {
+ this.modelId = modelId;
+ }
+}
+
+/**
+ * Options required for a conversation message with
+ * role of user
+ */
+export class UserRoleOpts {
+ revisionRootMessageId;
+
+ /**
+ * @param {string} [revisionRootMessageId=undefined]
+ */
+ constructor(revisionRootMessageId) {
+ if (revisionRootMessageId) {
+ this.revisionRootMessageId = revisionRootMessageId;
+ }
+ }
+}
+
+/**
+ * Used to retrieve chat entries for the History app menu
+ */
+export class ChatMinimal {
+ #id;
+ #title;
+
+ /**
+ * @param {object} params
+ * @param {string} params.convId
+ * @param {string} params.title
+ */
+ constructor({ convId, title }) {
+ this.#id = convId;
+ this.#title = title;
+ }
+
+ get id() {
+ return this.#id;
+ }
+
+ get title() {
+ return this.#title;
+ }
+}
diff --git a/browser/components/aiwindow/ui/modules/ChatMigrations.sys.mjs b/browser/components/aiwindow/ui/modules/ChatMigrations.sys.mjs
@@ -0,0 +1,47 @@
+/*
+ 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/. */
+
+/**
+ * Please refer to sql.mjs for details on creating new migrations.
+ *
+ * - List each change here and what it's for.
+ *
+ * @param {OpenedConnection} _conn - The SQLite connection to use for the migration
+ * @param {number} _version - The version number of the current schema
+ */
+async function applyV2(_conn, _version) {
+ // if (version < 2) {
+ // // Drop the summary table, it is no longer needed
+ // await conn.execute(`DROP TABLE IF EXISTS summary;`);
+ //
+ // try {
+ // await conn.execute("SELECT page_url FROM message LIMIT 1;");
+ // } catch (e) {
+ // // Add a page_url column to the message table to keep track
+ // // of the pages visitedjduring a conversation and when they
+ // // were visited
+ // await conn.execute(`
+ // ALTER TABLE message ADD COLUMN page_url TEXT;
+ // `);
+ // }
+ //
+ // try {
+ // await conn.execute("SELECT turn_index FROM message LIMIT 1;");
+ // } catch (e) {
+ // // Add a turn_index column to group all types of messages into a
+ // // particular turn, ex: prompt -> response would be a single turn.
+ // await conn.execute(`
+ // ALTER TABLE message ADD COLUMN turn_index INTEGER;
+ // `);
+ // }
+ // }
+}
+
+/**
+ * Array of migration functions to run in the order they should be run in.
+ *
+ * @returns {Array<Function>}
+ */
+export const migrations = [applyV2];
diff --git a/browser/components/aiwindow/ui/modules/ChatSql.sys.mjs b/browser/components/aiwindow/ui/modules/ChatSql.sys.mjs
@@ -0,0 +1,244 @@
+/*
+ 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/. */
+
+// Every time the schema or the underlying data changes, you must bump up the
+// schema version.
+
+// Remember to:
+// 1. Bump up the version number
+// 2. Add a migration function to migrate the data to the new schema.
+// 3. Update #createDatabaseEntities and #checkDatabaseHealth
+// 4. Add a test to check that the migration works correctly.
+
+// Note: migrations should be reasonably re-entry-friendly. If the user
+// downgrades, the schema version is decreased, and upon a subsequent upgrade,
+// the migration step is reapplied.
+// This ensures that any necessary conversions are performed, even for entries
+// added after the downgrade.
+// In practice, schema changes should be additive, allowing newer versions to
+// operate on older schemas, albeit with potentially reduced functionality.
+
+export const ESCAPE_CHAR = "/";
+
+export const CONVERSATION_TABLE = `
+CREATE TABLE conversation (
+ conv_id TEXT PRIMARY KEY,
+ title TEXT,
+ description TEXT,
+ page_url TEXT,
+ page_meta_jsonb BLOB,
+ created_date INTEGER NOT NULL,
+ updated_date INTEGER NOT NULL,
+ status INTEGER NOT NULL DEFAULT 0,
+ active_branch_tip_message_id TEXT -- no foreign here, as we insert messages later.
+) WITHOUT ROWID;
+`;
+
+export const CONVERSATION_UPDATED_DATE_INDEX = `
+CREATE INDEX conversation_updated_date_idx ON conversation(updated_date);
+`;
+
+export const MESSAGE_TABLE = `
+CREATE TABLE message (
+ message_id TEXT PRIMARY KEY,
+ conv_id TEXT NOT NULL REFERENCES conversation(conv_id) ON DELETE CASCADE,
+ created_date INTEGER NOT NULL,
+ parent_message_id TEXT REFERENCES message(message_id) ON DELETE CASCADE,
+ revision_root_message_id TEXT REFERENCES message(message_id) ON DELETE CASCADE,
+ ordinal INTEGER NOT NULL CHECK(ordinal >= 0),
+ is_active_branch INTEGER NOT NULL,
+ role INTEGER NOT NULL,
+ model_id TEXT,
+ params_jsonb BLOB,
+ content_jsonb BLOB,
+ usage_jsonb BLOB,
+ page_url TEXT,
+ turn_index INTEGER,
+ insights_enabled BOOLEAN,
+ insights_flag_source INTEGER,
+ insights_applied_jsonb BLOB,
+ web_search_queries_jsonb BLOB
+) WITHOUT ROWID;
+`;
+
+export const MESSAGE_ORDINAL_INDEX = `
+CREATE INDEX message_ordinal_idx ON message(ordinal);
+`;
+
+// @todo Bug 2005423
+// Maybe add hashed url column to optimize message_url_idx
+export const MESSAGE_URL_INDEX = `
+CREATE INDEX message_url_idx ON message(page_url);
+`;
+
+export const MESSAGE_CREATED_DATE_INDEX = `
+CREATE INDEX message_created_date_idx ON message(created_date);
+`;
+
+export const CONVERSATION_INSERT = `
+INSERT INTO conversation (
+ conv_id, title, description, page_url, page_meta_jsonb,
+ created_date, updated_date, status, active_branch_tip_message_id
+) VALUES (
+ :conv_id, :title, :description, :page_url, jsonb(:page_meta),
+ :created_date, :updated_date, :status, :active_branch_tip_message_id
+)
+ON CONFLICT(conv_id) DO UPDATE
+ SET title = :title,
+ updated_date = :updated_date,
+ status = :status,
+ active_branch_tip_message_id = :active_branch_tip_message_id;
+`;
+
+export const MESSAGE_INSERT = `
+INSERT INTO message (
+ message_id, conv_id, created_date, parent_message_id,
+ revision_root_message_id, ordinal, is_active_branch, role,
+ model_id, params_jsonb, content_jsonb, usage_jsonb, page_url, turn_index,
+ insights_enabled, insights_flag_source, insights_applied_jsonb,
+ web_search_queries_jsonb
+) VALUES (
+ :message_id, :conv_id, :created_date, :parent_message_id,
+ :revision_root_message_id, :ordinal, :is_active_branch, :role,
+ :model_id, jsonb(:params), jsonb(:content), jsonb(:usage), :page_url, :turn_index,
+ :insights_enabled, :insights_flag_source, jsonb(:insights_applied_jsonb),
+ jsonb(:web_search_queries_jsonb)
+)
+ON CONFLICT(message_id) DO UPDATE SET
+ is_active_branch = :is_active_branch;
+`;
+
+export const CONVERSATIONS_MOST_RECENT = `
+SELECT conv_id, title
+FROM conversation
+ORDER BY updated_date DESC
+LIMIT :limit;
+`;
+
+export const CONVERSATIONS_OLDEST = `
+SELECT conv_id, title
+FROM conversation
+ORDER BY updated_date ASC
+LIMIT :limit;
+`;
+
+export const CONVERSATION_BY_ID = `
+SELECT conv_id, title, description, page_url,
+ json(page_meta_jsonb) AS page_meta, created_date, updated_date,
+ status, active_branch_tip_message_id
+FROM conversation WHERE conv_id = :conv_id;
+`;
+
+export const CONVERSATIONS_BY_DATE = `
+SELECT conv_id, title, description, page_url,
+ json(page_meta_jsonb) AS page_meta, created_date, updated_date,
+ status, active_branch_tip_message_id
+FROM conversation
+WHERE updated_date >= :start_date AND updated_date <= :end_date
+ORDER BY updated_date DESC;
+`;
+
+export const CONVERSATIONS_BY_URL = `
+SELECT c.conv_id, c.title, c.description, c.page_url,
+ json(c.page_meta_jsonb) AS page_meta, c.created_date, c.updated_date,
+ c.status, c.active_branch_tip_message_id
+FROM conversation c
+WHERE EXISTS (
+ SELECT 1
+ FROM message m
+ WHERE m.conv_id = c.conv_id
+ AND m.page_url = :page_url
+)
+ORDER BY c.updated_date DESC;
+`;
+
+/**
+ * Get all messages for multiple conversations
+ *
+ * @param {number} amount - The number of conversation IDs to get messages for
+ */
+export function getConversationMessagesSql(amount) {
+ return `
+ SELECT
+ message_id, created_date, parent_message_id, revision_root_message_id,
+ ordinal, is_active_branch, role, model_id, conv_id,
+ json(params_jsonb) AS params, json(usage_jsonb) AS usage,
+ page_url, turn_index, insights_enabled, insights_flag_source,
+ json(insights_applied_jsonb) AS insights_applied,
+ json(web_search_queries_jsonb) AS web_search_queries,
+ json(content_jsonb) AS content
+ FROM message
+ WHERE conv_id IN(${new Array(amount).fill("?").join(",")})
+ ORDER BY ordinal ASC;
+ `;
+}
+
+export const CONVERSATIONS_CONTENT_SEARCH = `
+SELECT c.conv_id, c.title, c.description, c.page_url,
+ json(c.page_meta_jsonb) AS page_meta, c.created_date, c.updated_date,
+ c.status, c.active_branch_tip_message_id
+FROM conversation c
+JOIN message m ON m.conv_id = c.conv_id
+WHERE json_type(m.content_jsonb, :path) IS NOT NULL;
+`;
+
+export const CONVERSATIONS_CONTENT_SEARCH_BY_ROLE = `
+SELECT c.conv_id, c.title, c.description, c.page_url,
+ json(c.page_meta_jsonb) AS page_meta, c.created_date, c.updated_date,
+ c.status, c.active_branch_tip_message_id
+FROM conversation c
+JOIN message m ON m.conv_id = c.conv_id
+WHERE m.role = :role
+ AND json_type(m.content_jsonb, :path) IS NOT NULL;
+`;
+
+export const CONVERSATIONS_HISTORY_SEARCH = `
+SELECT c.conv_id, c.title, c.description, c.page_url,
+ json(c.page_meta_jsonb) AS page_meta, c.created_date, c.updated_date,
+ c.status, c.active_branch_tip_message_id
+FROM conversation c
+JOIN message m ON m.conv_id = c.conv_id
+WHERE m.role = 0
+ AND (
+ CAST(json_extract(m.content_jsonb, :path) AS TEXT) LIKE :pattern ESCAPE '/'
+ OR
+ c.title LIKE :pattern ESCAPE '/'
+ );
+`;
+
+export const MESSAGES_BY_DATE = `
+SELECT
+ message_id, created_date, parent_message_id, revision_root_message_id,
+ ordinal, is_active_branch, role, model_id, conv_id,
+ json(params_jsonb) AS params, json(usage_jsonb) AS usage,
+ page_url, turn_index, insights_enabled, insights_flag_source,
+ json(insights_applied_jsonb) AS insights_applied,
+ json(web_search_queries_jsonb) AS web_search_queries,
+ json(content_jsonb) AS content
+FROM message
+WHERE created_date >= :start_date AND created_date <= :end_date
+ORDER BY created_date DESC
+LIMIT :limit OFFSET :offset;
+`;
+
+export const MESSAGES_BY_DATE_AND_ROLE = `
+SELECT
+ message_id, created_date, parent_message_id, revision_root_message_id,
+ ordinal, is_active_branch, role, model_id, conv_id,
+ json(params_jsonb) AS params, json(usage_jsonb) AS usage,
+ page_url, turn_index, insights_enabled, insights_flag_source,
+ json(insights_applied_jsonb) AS insights_applied,
+ json(web_search_queries_jsonb) AS web_search_queries,
+ json(content_jsonb) AS content
+FROM message
+WHERE role = :role
+ AND created_date >= :start_date AND created_date <= :end_date
+ORDER BY created_date DESC
+LIMIT :limit OFFSET :offset;
+`;
+
+export const DELETE_CONVERSATION_BY_ID = `
+DELETE FROM conversation WHERE conv_id = :conv_id;
+`;
diff --git a/browser/components/aiwindow/ui/modules/ChatStore.sys.mjs b/browser/components/aiwindow/ui/modules/ChatStore.sys.mjs
@@ -0,0 +1,798 @@
+/*
+ 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_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,
+ 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,
+ 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;
+
+/**
+ * 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);
+ }
+
+ /**
+ * 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);
+ }
+
+ 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);
+ }
+}
diff --git a/browser/components/aiwindow/ui/modules/ChatUtils.sys.mjs b/browser/components/aiwindow/ui/modules/ChatUtils.sys.mjs
@@ -0,0 +1,132 @@
+/*
+ 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, {
+ CryptoUtils: "moz-src:///services/crypto/modules/utils.sys.mjs",
+});
+
+import { MESSAGE_ROLE } from "./ChatConstants.sys.mjs";
+import { ChatConversation } from "./ChatConversation.sys.mjs";
+import { ChatMessage } from "./ChatMessage.sys.mjs";
+
+/**
+ * Creates a 12 characters GUID with 72 bits of entropy.
+ *
+ * @returns {string} A base64url encoded GUID.
+ */
+export function makeGuid() {
+ return ChromeUtils.base64URLEncode(lazy.CryptoUtils.generateRandomBytes(9), {
+ pad: false,
+ });
+}
+
+/**
+ * Parse a conversation row from the database into a ChatConversation
+ * object.
+ *
+ * @param {object} row - The database row to parse.
+ * @returns {ChatConversation} The parsed conversation object.
+ */
+export function parseConversationRow(row) {
+ return new ChatConversation({
+ id: row.getResultByName("conv_id"),
+ title: row.getResultByName("title"),
+ description: row.getResultByName("description"),
+ pageUrl: URL.parse(row.getResultByName("page_url")),
+ pageMeta: parseJSONOrNull(row.getResultByName("page_meta")),
+ createdDate: row.getResultByName("created_date"),
+ updatedDate: row.getResultByName("updated_date"),
+ status: row.getResultByName("status"),
+ });
+}
+
+/**
+ * Parse message rows from the database into an array of ChatMessage
+ * objects.
+ *
+ * @param {object} rows - The database rows to parse.
+ * @returns {Array<ChatMessage>} The parsed message objects.
+ */
+export function parseMessageRows(rows) {
+ return rows.map(row => {
+ return new ChatMessage({
+ id: row.getResultByName("message_id"),
+ createdDate: row.getResultByName("created_date"),
+ parentMessageId: row.getResultByName("parent_message_id"),
+ revisionRootMessageId: row.getResultByName("revision_root_message_id"),
+ ordinal: row.getResultByName("ordinal"),
+ isActiveBranch: !!row.getResultByName("is_active_branch"),
+ role: row.getResultByName("role"),
+ modelId: row.getResultByName("model_id"),
+ params: parseJSONOrNull(row.getResultByName("params")),
+ usage: parseJSONOrNull(row.getResultByName("usage")),
+ content: parseJSONOrNull(row.getResultByName("content")),
+ convId: row.getResultByName("conv_id"),
+ pageUrl: URL.parse(row.getResultByName("page_url")),
+ turnIndex: row.getResultByName("turn_index"),
+ insightsEnabled: row.getResultByName("insights_enabled"),
+ insightsFlagSource: row.getResultByName("insights_flag_source"),
+ insightsApplied: parseJSONOrNull(row.getResultByName("insights_applied")),
+ webSearchQueries: parseJSONOrNull(
+ row.getResultByName("web_search_queries")
+ ),
+ });
+ });
+}
+
+/**
+ * Try to parse a JSON string, returning null if it fails or the value is falsy.
+ *
+ * @param {string} value - The JSON string to parse.
+ * @returns {object|null} The parsed object or null.
+ */
+export function parseJSONOrNull(value) {
+ if (!value) {
+ return null;
+ }
+ try {
+ return JSON.parse(value);
+ } catch (e) {
+ return null;
+ }
+}
+
+/**
+ * Try to stringify a value if it is truthy, otherwise return null.
+ *
+ * @param {*} value - A value to JSON.stringify()
+ *
+ * @returns {string|null} - JSON string
+ */
+export function toJSONOrNull(value) {
+ return value ? JSON.stringify(value) : null;
+}
+
+/**
+ * Converts the different types of message roles from
+ * the database numeric type to a string label
+ *
+ * @param {number} role - The database numeric role type
+ * @returns {string} - A human readable role label
+ */
+export function getRoleLabel(role) {
+ switch (role) {
+ case MESSAGE_ROLE.USER:
+ return "User";
+
+ case MESSAGE_ROLE.ASSISTANT:
+ return "Assistant";
+
+ case MESSAGE_ROLE.SYSTEM:
+ return "System";
+
+ case MESSAGE_ROLE.TOOL:
+ return "Tool";
+ }
+
+ return "";
+}
diff --git a/browser/components/aiwindow/ui/moz.build b/browser/components/aiwindow/ui/moz.build
@@ -11,6 +11,15 @@ MOZ_SRC_FILES += [
"actors/AIChatContentChild.sys.mjs",
"actors/AIChatContentParent.sys.mjs",
"modules/AIWindow.sys.mjs",
+ "modules/ChatConstants.sys.mjs",
+ "modules/ChatConversation.sys.mjs",
+ "modules/ChatMessage.sys.mjs",
+ "modules/ChatMigrations.sys.mjs",
+ "modules/ChatSql.sys.mjs",
+ "modules/ChatStore.sys.mjs",
+ "modules/ChatUtils.sys.mjs",
]
JAR_MANIFESTS += ["jar.mn"]
+
+XPCSHELL_TESTS_MANIFESTS += ["test/xpcshell/xpcshell.toml"]
diff --git a/browser/components/aiwindow/ui/test/xpcshell/asserts.js b/browser/components/aiwindow/ui/test/xpcshell/asserts.js
@@ -0,0 +1,256 @@
+/* Any copyright is dedicated to the Public Domain.
+ * http://creativecommons.org/publicdomain/zero/1.0/ */
+
+"use strict";
+
+/**
+ * @todo Bug 2005403
+ * Move this to main assert module, could look like:
+ * await Assert.continueOnAssert(async function() {
+ * Assert.ok();
+ * await Assert.rejects();
+ * ...
+ * });
+ */
+function SoftAssert() {
+ this.failures = [];
+}
+
+SoftAssert.prototype._recordFailure = function (err, context) {
+ this.failures.push({
+ message: err.message,
+ name: err.name,
+ stack: err.stack,
+ context,
+ });
+};
+
+SoftAssert.prototype.throws = function (
+ block,
+ expectedError,
+ message = "throws failed"
+) {
+ try {
+ Assert.throws(block, expectedError, message);
+
+ this._recordFailure(new Error("Did not throw"), {
+ type: "throws",
+ expectedError,
+ message,
+ });
+ } catch {}
+};
+
+SoftAssert.prototype.ok = function (condition, message = "ok failed") {
+ try {
+ Assert.ok(condition, message);
+ } catch (err) {
+ this._recordFailure(err, { type: "ok", condition, message });
+ }
+};
+
+SoftAssert.prototype.equal = function (
+ actual,
+ expected,
+ message = "equal failed"
+) {
+ try {
+ Assert.equal(actual, expected, message);
+ } catch (err) {
+ this._recordFailure(err, { type: "equal", actual, expected, message });
+ }
+};
+
+SoftAssert.prototype.notEqual = function (
+ actual,
+ expected,
+ message = "notEqual failed"
+) {
+ try {
+ Assert.notEqual(actual, expected, message);
+ } catch (err) {
+ this._recordFailure(err, { type: "notEqual", actual, expected, message });
+ }
+};
+
+SoftAssert.prototype.deepEqual = function (
+ actual,
+ expected,
+ message = "deepEqual failed"
+) {
+ try {
+ Assert.deepEqual(actual, expected, message);
+ } catch (err) {
+ this._recordFailure(err, { type: "deepEqual", actual, expected, message });
+ }
+};
+
+SoftAssert.prototype.notDeepEqual = function (
+ actual,
+ expected,
+ message = "notDeepEqual failed"
+) {
+ try {
+ Assert.notDeepEqual(actual, expected, message);
+ } catch (err) {
+ this._recordFailure(err, {
+ type: "notDeepEqual",
+ actual,
+ expected,
+ message,
+ });
+ }
+};
+
+SoftAssert.prototype.strictEqual = function (
+ actual,
+ expected,
+ message = "strictEqual failed"
+) {
+ try {
+ Assert.strictEqual(actual, expected, message);
+ } catch (err) {
+ this._recordFailure(err, {
+ type: "strictEqual",
+ actual,
+ expected,
+ message,
+ });
+ }
+};
+
+SoftAssert.prototype.notStrictEqual = function (
+ actual,
+ expected,
+ message = "notStrictEqual failed"
+) {
+ try {
+ Assert.notStrictEqual(actual, expected, message);
+ } catch (err) {
+ this._recordFailure(err, {
+ type: "notStrictEqual",
+ actual,
+ expected,
+ message,
+ });
+ }
+};
+
+SoftAssert.prototype.rejects = async function (
+ actual,
+ expected,
+ message = "rejects failed"
+) {
+ try {
+ await Assert.rejects(actual, expected, message);
+ } catch (err) {
+ this._recordFailure(err, { type: "rejects", actual, expected, message });
+ }
+};
+
+SoftAssert.prototype.greater = function (
+ actual,
+ expected,
+ message = "greater failed"
+) {
+ try {
+ Assert.greater(actual, expected, message);
+ } catch (err) {
+ this._recordFailure(err, { type: "greater", actual, expected, message });
+ }
+};
+
+SoftAssert.prototype.greaterOrEqual = function (
+ actual,
+ expected,
+ message = "greaterOrEqual failed"
+) {
+ try {
+ Assert.greaterOrEqual(actual, expected, message);
+ } catch (err) {
+ this._recordFailure(err, {
+ type: "greaterOrEqual",
+ actual,
+ expected,
+ message,
+ });
+ }
+};
+
+SoftAssert.prototype.less = function (
+ actual,
+ expected,
+ message = "less failed"
+) {
+ try {
+ Assert.less(actual, expected, message);
+ } catch (err) {
+ this._recordFailure(err, { type: "less", actual, expected, message });
+ }
+};
+
+SoftAssert.prototype.lessOrEqual = function (
+ actual,
+ expected,
+ message = "lessOrEqual failed"
+) {
+ try {
+ Assert.lessOrEqual(actual, expected, message);
+ } catch (err) {
+ this._recordFailure(err, {
+ type: "lessOrEqual",
+ actual,
+ expected,
+ message,
+ });
+ }
+};
+
+// Call this at the end to fail the test if any soft asserts failed.
+SoftAssert.prototype.assertAll = function (label = "Soft assertion failures") {
+ if (!this.failures.length) {
+ return;
+ }
+
+ let details = `${label} (${this.failures.length}):\n`;
+ for (let i = 0; i < this.failures.length; i++) {
+ const f = this.failures[i];
+ details += `\n[${i + 1}] ${f.message}\n`;
+ if (f.context) {
+ details += ` Context: ${JSON.stringify(f.context)}\n`;
+ }
+ if (f.stack) {
+ details += ` Stack: ${f.stack}\n`;
+ }
+ }
+
+ // One combined failure at the end:
+ throw new Error(details);
+};
+
+function softAssertions(Assert) {
+ /**
+ * Provides a way to execute multiple assertions without exiting
+ * a test after the first assertion failure. All assertions in the
+ * provided callback will run even if some assertions fail. If an
+ * assertion fails, the test will fail and provide details on the
+ * failed assertion(s).
+ *
+ * @example
+ * Assert.withSoftAssertions(function(soft) {
+ * // Use usual Assert functions via `soft` reference:
+ * soft.equal(conversation.id.length, 12);
+ * soft.ok(Array.isArray(conversation.messages));
+ * soft.ok(!isNaN(conversation.createdDate));
+ * });
+ *
+ * @param {Function} tests - A callback that acceps a `soft` object to run assertions with.
+ */
+ Assert.withSoftAssertions = function (tests) {
+ const soft = new SoftAssert();
+ tests(soft);
+ soft.assertAll();
+ };
+}
+
+softAssertions(Assert);
diff --git a/browser/components/aiwindow/ui/test/xpcshell/test_ChatConversation.js b/browser/components/aiwindow/ui/test/xpcshell/test_ChatConversation.js
@@ -0,0 +1,412 @@
+/* Any copyright is dedicated to the Public Domain.
+ * http://creativecommons.org/publicdomain/zero/1.0/ */
+
+do_get_profile();
+
+const { ChatConversation, MESSAGE_ROLE } = ChromeUtils.importESModule(
+ "moz-src:///browser/components/aiwindow/ui/modules/ChatStore.sys.mjs"
+);
+
+const { UserRoleOpts, AssistantRoleOpts, ToolRoleOpts } =
+ ChromeUtils.importESModule(
+ "moz-src:///browser/components/aiwindow/ui/modules/ChatMessage.sys.mjs"
+ );
+
+add_task(function test_ChatConversation_constructor_defaults() {
+ const conversation = new ChatConversation({});
+
+ Assert.withSoftAssertions(function (soft) {
+ soft.equal(conversation.id.length, 12);
+ soft.ok(Array.isArray(conversation.messages));
+ soft.ok(!isNaN(conversation.createdDate));
+ soft.ok(!isNaN(conversation.updatedDate));
+ soft.strictEqual(conversation.title, undefined);
+ soft.strictEqual(conversation.description, undefined);
+ soft.strictEqual(conversation.pageUrl, undefined);
+ soft.strictEqual(conversation.pageMeta, undefined);
+ });
+});
+
+add_task(function test_ChatConversation_addMessage() {
+ const conversation = new ChatConversation({});
+
+ const content = {
+ type: "text",
+ content: "hello world",
+ };
+
+ conversation.addMessage(
+ MESSAGE_ROLE.USER,
+ content,
+ new URL("https://www.mozilla.com"),
+ 0
+ );
+
+ const message = conversation.messages[0];
+
+ Assert.withSoftAssertions(function (soft) {
+ soft.strictEqual(message.role, MESSAGE_ROLE.USER);
+ soft.strictEqual(message.content, content);
+ soft.strictEqual(message.pageUrl.href, "https://www.mozilla.com/");
+ soft.strictEqual(message.turnIndex, 0);
+ });
+});
+
+add_task(function test_invalidRole_ChatConversation_addMessage() {
+ const conversation = new ChatConversation({});
+
+ const content = {
+ type: "text",
+ content: "hello world",
+ };
+
+ conversation.addMessage(313, content, new URL("https://www.mozilla.com"), 0);
+
+ Assert.equal(conversation.messages.length, 0);
+});
+
+add_task(function test_negativeTurnIndex_ChatConversation_addMessage() {
+ const conversation = new ChatConversation({});
+
+ const content = {
+ type: "text",
+ content: "hello world",
+ };
+
+ conversation.addMessage(
+ MESSAGE_ROLE.USER,
+ content,
+ new URL("https://www.mozilla.com"),
+ -1
+ );
+ const message = conversation.messages[0];
+
+ Assert.equal(message.turnIndex, 0);
+});
+
+add_task(function test_parentMessageId_ChatConversation_addMessage() {
+ const conversation = new ChatConversation({});
+
+ const content = {
+ type: "text",
+ content: "hello world",
+ };
+
+ conversation.addMessage(
+ MESSAGE_ROLE.USER,
+ content,
+ new URL("https://www.mozilla.com"),
+ 0
+ );
+
+ conversation.addMessage(
+ MESSAGE_ROLE.ASSISTANT,
+ content,
+ new URL("https://www.mozilla.com"),
+ 0
+ );
+
+ const userMsg = conversation.messages[0];
+ const assistantMsg = conversation.messages[1];
+
+ Assert.equal(assistantMsg.parentMessageId, userMsg.id);
+});
+
+add_task(function test_ordinal_ChatConversation_addMessage() {
+ const conversation = new ChatConversation({});
+
+ const content = {
+ type: "text",
+ content: "hello world",
+ };
+
+ conversation.addMessage(
+ MESSAGE_ROLE.USER,
+ content,
+ new URL("https://www.mozilla.com"),
+ 0
+ );
+
+ conversation.addMessage(
+ MESSAGE_ROLE.ASSISTANT,
+ content,
+ new URL("https://www.mozilla.com"),
+ 0
+ );
+
+ const userMsg = conversation.messages[0];
+ const assistantMsg = conversation.messages[1];
+
+ Assert.withSoftAssertions(function (soft) {
+ soft.equal(userMsg.ordinal, 1);
+ soft.equal(assistantMsg.ordinal, 2);
+ });
+});
+
+add_task(function test_ChatConversation_addUserMessage() {
+ const conversation = new ChatConversation({});
+
+ const content = "user to assistant msg";
+ conversation.addUserMessage(content, "https://www.mozilla.com", 0);
+
+ const message = conversation.messages[0];
+
+ Assert.withSoftAssertions(function (soft) {
+ soft.equal(message.role, MESSAGE_ROLE.USER);
+ soft.equal(message.turnIndex, 0);
+ soft.deepEqual(message.pageUrl, new URL("https://www.mozilla.com"));
+ soft.deepEqual(message.content, {
+ type: "text",
+ body: "user to assistant msg",
+ });
+ });
+});
+
+add_task(function test_revisionRootMessageId_ChatConversation_addUserMessage() {
+ const conversation = new ChatConversation({});
+
+ const content = "user to assistant msg";
+ conversation.addUserMessage(content, "https://www.firefox.com", 0);
+
+ const message = conversation.messages[0];
+
+ Assert.equal(message.revisionRootMessageId, message.id);
+});
+
+add_task(function test_opts_ChatConversation_addUserMessage() {
+ const conversation = new ChatConversation({});
+
+ const content = "user to assistant msg";
+ conversation.addUserMessage(
+ content,
+ "https://www.firefox.com",
+ 0,
+ new UserRoleOpts("321")
+ );
+
+ const message = conversation.messages[0];
+
+ Assert.equal(message.revisionRootMessageId, "321");
+});
+
+add_task(function test_ChatConversation_addAssistantMessage() {
+ const conversation = new ChatConversation({});
+
+ const content = "response from assistant";
+ conversation.addAssistantMessage("text", content, 0);
+
+ const message = conversation.messages[0];
+
+ Assert.withSoftAssertions(function (soft) {
+ soft.equal(message.role, MESSAGE_ROLE.ASSISTANT);
+ soft.equal(message.turnIndex, 0);
+ soft.deepEqual(message.pageUrl, null);
+ soft.deepEqual(message.content, {
+ type: "text",
+ body: "response from assistant",
+ });
+ soft.strictEqual(message.modelId, null, "modelId should default to false");
+ soft.strictEqual(message.params, null, "params should default to null");
+ soft.strictEqual(message.usage, null, "usage should default to null");
+ soft.strictEqual(
+ message.insightsEnabled,
+ false,
+ "insightsEnabled should default to false"
+ );
+ soft.strictEqual(
+ message.insightsFlagSource,
+ null,
+ "insightsFlagSource should default to null"
+ );
+ soft.deepEqual(
+ message.insightsApplied,
+ [],
+ "insightsApplied should default to emtpy array"
+ );
+ soft.deepEqual(
+ message.webSearchQueries,
+ [],
+ "webSearchQueries should default to emtpy array"
+ );
+ });
+});
+
+add_task(function test_opts_ChatConversation_addAssistantMessage() {
+ const conversation = new ChatConversation({});
+
+ const content = "response from assistant";
+ const assistantOpts = new AssistantRoleOpts(
+ "the-model-id",
+ { some: "params for model" },
+ { usage: "data" },
+ true,
+ 1,
+ ["insight"],
+ ["search"]
+ );
+ conversation.addAssistantMessage("text", content, 0, assistantOpts);
+
+ const message = conversation.messages[0];
+
+ Assert.withSoftAssertions(function (soft) {
+ soft.equal(message.role, MESSAGE_ROLE.ASSISTANT);
+ soft.equal(message.turnIndex, 0);
+ soft.deepEqual(message.pageUrl, null);
+ soft.deepEqual(message.content, {
+ type: "text",
+ body: "response from assistant",
+ });
+ soft.strictEqual(
+ message.modelId,
+ "the-model-id",
+ "modelId should be 'the-model-id'"
+ );
+ soft.deepEqual(
+ message.params,
+ { some: "params for model" },
+ 'params should equal { some: "params for model"}'
+ );
+ soft.deepEqual(
+ message.usage,
+ { usage: "data" },
+ 'usage should equal {"usage": "data"}'
+ );
+ soft.strictEqual(
+ message.insightsEnabled,
+ true,
+ "insightsEnabled should equal true"
+ );
+ soft.strictEqual(
+ message.insightsFlagSource,
+ 1,
+ "insightsFlagSource equal 1"
+ );
+ soft.deepEqual(
+ message.insightsApplied,
+ ["insight"],
+ "insightsApplied should equal ['insight']"
+ );
+ soft.deepEqual(
+ message.webSearchQueries,
+ ["search"],
+ "insightsApplied should equal ['search']"
+ );
+ });
+});
+
+add_task(function test_ChatConversation_addToolCallMessage() {
+ const conversation = new ChatConversation({});
+
+ const content = {
+ random: "tool call specific keys",
+ };
+ conversation.addToolCallMessage(content, 0);
+
+ const message = conversation.messages[0];
+
+ Assert.withSoftAssertions(function (soft) {
+ soft.equal(message.role, MESSAGE_ROLE.TOOL);
+ soft.equal(message.turnIndex, 0);
+ soft.deepEqual(message.pageUrl, null);
+ soft.deepEqual(message.content, {
+ random: "tool call specific keys",
+ });
+ soft.equal(message.modelId, null, "modelId should default to null");
+ });
+});
+
+add_task(function test_opts_ChatConversation_addToolCallMessage() {
+ const conversation = new ChatConversation({});
+
+ const content = {
+ random: "tool call specific keys",
+ };
+ conversation.addToolCallMessage(content, 0, new ToolRoleOpts("the-model-id"));
+
+ const message = conversation.messages[0];
+
+ Assert.withSoftAssertions(function (soft) {
+ soft.equal(message.role, MESSAGE_ROLE.TOOL);
+ soft.equal(message.turnIndex, 0);
+ soft.deepEqual(message.pageUrl, null);
+ soft.deepEqual(message.content, {
+ random: "tool call specific keys",
+ });
+ soft.equal(
+ message.modelId,
+ "the-model-id",
+ "modelId should equal the-model-id"
+ );
+ });
+});
+
+add_task(function test_ChatConversation_addSystemMessage() {
+ const conversation = new ChatConversation({});
+
+ const content = {
+ random: "system call specific keys",
+ };
+ conversation.addSystemMessage("text", content, 0);
+
+ const message = conversation.messages[0];
+
+ Assert.withSoftAssertions(function (soft) {
+ soft.equal(message.role, MESSAGE_ROLE.SYSTEM);
+ soft.equal(message.turnIndex, 0);
+ soft.deepEqual(message.pageUrl, null);
+ soft.deepEqual(message.content, {
+ type: "text",
+ body: { random: "system call specific keys" },
+ });
+ });
+});
+
+add_task(function test_ChatConversation_getSitesList() {
+ const conversation = new ChatConversation({});
+
+ const content = "user to assistant msg";
+ conversation.addUserMessage(content, "https://www.mozilla.com", 0);
+ conversation.addUserMessage(content, "https://www.mozilla.com", 1);
+ conversation.addUserMessage(content, "https://www.firefox.com", 2);
+ conversation.addUserMessage(content, "https://www.cnn.com", 3);
+ conversation.addUserMessage(content, "https://www.espn.com", 4);
+ conversation.addUserMessage(content, "https://www.espn.com", 5);
+
+ const sites = conversation.getSitesList();
+
+ Assert.deepEqual(sites, [
+ URL.parse("https://www.mozilla.com/"),
+ URL.parse("https://www.firefox.com/"),
+ URL.parse("https://www.cnn.com/"),
+ URL.parse("https://www.espn.com/"),
+ ]);
+});
+
+add_task(function test_ChatConversation_getMostRecentPageVisited() {
+ const conversation = new ChatConversation({});
+
+ const content = "user to assistant msg";
+ conversation.addUserMessage(content, "https://www.mozilla.com", 0);
+ conversation.addUserMessage(content, "https://www.mozilla.com", 1);
+ conversation.addUserMessage(content, "https://www.firefox.com", 2);
+ conversation.addUserMessage(content, "https://www.cnn.com", 3);
+ conversation.addUserMessage(content, "https://www.espn.com", 4);
+ conversation.addUserMessage(content, "https://www.espn.com", 5);
+
+ const mostRecentPageVisited = conversation.getMostRecentPageVisited();
+
+ Assert.equal(mostRecentPageVisited, "https://www.espn.com/");
+});
+
+add_task(function test_noBrowsing_ChatConversation_getMostRecentPageVisited() {
+ const conversation = new ChatConversation({});
+
+ const content = "user to assistant msg";
+ conversation.addUserMessage(content, "about:aiwindow", 0);
+ conversation.addUserMessage(content, "", 1);
+ conversation.addUserMessage(content, null, 2);
+
+ const mostRecentPageVisited = conversation.getMostRecentPageVisited();
+
+ Assert.equal(mostRecentPageVisited, null);
+});
diff --git a/browser/components/aiwindow/ui/test/xpcshell/test_ChatMessage.js b/browser/components/aiwindow/ui/test/xpcshell/test_ChatMessage.js
@@ -0,0 +1,113 @@
+/* Any copyright is dedicated to the Public Domain.
+ * http://creativecommons.org/publicdomain/zero/1.0/ */
+
+do_get_profile();
+
+const { ChatMessage } = ChromeUtils.importESModule(
+ "moz-src:///browser/components/aiwindow/ui/modules/ChatStore.sys.mjs"
+);
+
+add_task(function test_ChatConversation_constructor_defaults() {
+ const message = new ChatMessage({
+ ordinal: 0,
+ role: 0,
+ turnIndex: 0,
+ content: "some content",
+ pageUrl: new URL("https://www.mozilla.com"),
+ });
+
+ Assert.withSoftAssertions(function (soft) {
+ soft.equal(message.id.length, 12);
+ soft.equal(message.revisionRootMessageId, message.id);
+ soft.ok(!isNaN(message.createdDate));
+ soft.ok(message.isActiveBranch);
+ const nullFields = [
+ "parentMessageId",
+ "modelId",
+ "params",
+ "usage",
+ "convId",
+ "insightsEnabled",
+ "insightsFlagSource",
+ "insightsApplied",
+ "webSearchQueries",
+ ];
+
+ nullFields.forEach(nullField => {
+ soft.equal(
+ message[nullField],
+ null,
+ `message.${nullField} should default to null`
+ );
+ });
+ });
+});
+
+add_task(function test_pageUrl_as_URL_ChatConversation() {
+ const message = new ChatMessage({
+ ordinal: 0,
+ role: 0,
+ turnIndex: 0,
+ content: "some content",
+ pageUrl: new URL("https://www.mozilla.com"),
+ });
+
+ Assert.withSoftAssertions(function (soft) {
+ soft.ok(URL.isInstance(message.pageUrl));
+ soft.equal(message.pageUrl.href, "https://www.mozilla.com/");
+ });
+});
+
+add_task(function test_pageUrl_as_string_ChatConversation() {
+ const message = new ChatMessage({
+ ordinal: 0,
+ role: 0,
+ turnIndex: 0,
+ content: "some content",
+ pageUrl: "https://www.mozilla.com",
+ });
+
+ Assert.withSoftAssertions(function (soft) {
+ soft.ok(URL.isInstance(message.pageUrl));
+ soft.equal(message.pageUrl.href, "https://www.mozilla.com/");
+ });
+});
+
+add_task(function test_invalid_pageUrl_ChatConversation() {
+ Assert.throws(
+ function () {
+ new ChatMessage({
+ ordinal: 0,
+ role: 0,
+ turnIndex: 0,
+ content: "some content",
+ pageUrl: "www.mozilla.com",
+ });
+ },
+ new RegExp("URL constructor: www.mozilla.com is not a valid URL"),
+ "Did not get correct error message"
+ );
+});
+
+add_task(function test_missing_pageUrl_ChatConversation() {
+ const message = new ChatMessage({
+ ordinal: 0,
+ role: 0,
+ turnIndex: 0,
+ content: "some content",
+ });
+
+ Assert.equal(message.pageUrl, null);
+});
+
+add_task(function test_empty_pageUrl_ChatConversation() {
+ const message = new ChatMessage({
+ ordinal: 0,
+ role: 0,
+ turnIndex: 0,
+ content: "some content",
+ pageUrl: "",
+ });
+
+ Assert.equal(message.pageUrl, null);
+});
diff --git a/browser/components/aiwindow/ui/test/xpcshell/test_ChatStore.js b/browser/components/aiwindow/ui/test/xpcshell/test_ChatStore.js
@@ -0,0 +1,613 @@
+/* Any copyright is dedicated to the Public Domain.
+ * http://creativecommons.org/publicdomain/zero/1.0/ */
+
+do_get_profile();
+
+const lazy = {};
+
+ChromeUtils.defineESModuleGetters(lazy, {
+ sinon: "resource://testing-common/Sinon.sys.mjs",
+});
+
+const { ChatStore, ChatConversation, ChatMessage } = ChromeUtils.importESModule(
+ "moz-src:///browser/components/aiwindow/ui/modules/ChatStore.sys.mjs"
+);
+
+async function addBasicConvoTestData(date, title, updated = null) {
+ const link = "https://www.firefox.com";
+ const updatedDate = updated || date;
+
+ return addConvoWithSpecificTestData(
+ new Date(date),
+ link,
+ link,
+ title,
+ "test content",
+ new Date(updatedDate)
+ );
+}
+
+async function addBasicConvoWithSpecificUpdatedTestData(updatedDate, title) {
+ const link = "https://www.firefox.com";
+ return addConvoWithSpecificTestData(
+ new Date("1/1/2023"),
+ link,
+ link,
+ title,
+ "test content",
+ new Date(updatedDate)
+ );
+}
+
+async function addConvoWithSpecificTestData(
+ createdDate,
+ mainLink,
+ messageLink,
+ title,
+ message = "the message body",
+ updatedDate = false
+) {
+ const conversation = new ChatConversation({
+ createdDate: createdDate.getTime(),
+ updatedDate: updatedDate ? updatedDate.getTime() : createdDate.getTime(),
+ pageUrl: mainLink,
+ });
+ conversation.title = title;
+ conversation.addUserMessage(message, messageLink, 0);
+ await gChatStore.updateConversation(conversation);
+}
+
+async function addConvoWithSpecificCustomContentTestData(
+ createdDate,
+ mainLink,
+ messageLink,
+ title,
+ content,
+ role
+) {
+ const conversation = new ChatConversation({
+ createdDate: createdDate.getTime(),
+ updatedDate: createdDate.getTime(),
+ pageUrl: mainLink,
+ });
+ conversation.title = title;
+ conversation.addMessage(role, content, messageLink, 0);
+ await gChatStore.updateConversation(conversation);
+}
+
+/**
+ * Runs a test atomically so that the clean up code
+ * runs after each test intead of after the entire
+ * list of tasks in the file are done.
+ *
+ * @todo Bug 2005408
+ * Replace add_atomic_task usage when this Bug 1656557 lands
+ *
+ * @param {Function} func - The test function to run
+ */
+function add_atomic_task(func) {
+ return add_task(async function () {
+ await test_ChatStorage_setup();
+
+ try {
+ await func();
+ } finally {
+ await test_cleanUp();
+ }
+ });
+}
+
+let gChatStore, gSandbox;
+
+async function cleanUpDatabase() {
+ if (gChatStore) {
+ await gChatStore.destroyDatabase();
+ gChatStore = null;
+ }
+}
+
+async function test_ChatStorage_setup() {
+ Services.prefs.setBoolPref("browser.aiwindow.removeDatabaseOnStartup", true);
+
+ gChatStore = new ChatStore();
+ await gChatStore.destroyDatabase();
+
+ gSandbox = lazy.sinon.createSandbox();
+}
+
+async function test_cleanUp() {
+ Services.prefs.clearUserPref("browser.aiwindow.removeDatabaseOnStartup");
+
+ await cleanUpDatabase();
+ gSandbox.restore();
+}
+
+add_atomic_task(async function task_ChatStorage_constructor() {
+ gChatStore = new ChatStore();
+
+ Assert.ok(gChatStore, "Should return a ChatStorage instance");
+});
+
+add_atomic_task(async function test_ChatStorage_updateConversation() {
+ let success = true;
+ let errorMessage = "";
+
+ try {
+ gChatStore = new ChatStore();
+ const conversation = new ChatConversation({});
+
+ conversation.addUserMessage("test content", "https://www.firefox.com", 0);
+
+ await gChatStore.updateConversation(conversation);
+ } catch (e) {
+ success = false;
+ errorMessage = e.message;
+ }
+
+ Assert.ok(success, errorMessage);
+});
+
+add_atomic_task(async function test_ChatStorage_findRecentConversations() {
+ gChatStore = new ChatStore();
+
+ await addBasicConvoTestData("1/1/2025", "conversation 1");
+ await addBasicConvoTestData("1/2/2025", "conversation 2");
+ await addBasicConvoTestData("1/3/2025", "conversation 3");
+
+ const recentConversations = await gChatStore.findRecentConversations(2);
+
+ Assert.withSoftAssertions(function (soft) {
+ soft.equal(recentConversations[0].title, "conversation 3");
+ soft.equal(recentConversations[1].title, "conversation 2");
+ });
+});
+
+add_atomic_task(async function test_ChatStorage_findConversationById() {
+ gChatStore = new ChatStore();
+
+ let conversation = new ChatConversation({});
+ conversation.title = "conversation 1";
+ conversation.addUserMessage("test content", "https://www.firefox.com", 0);
+ await gChatStore.updateConversation(conversation);
+
+ const conversationId = conversation.id;
+
+ conversation = await gChatStore.findConversationById(conversationId);
+
+ Assert.withSoftAssertions(function (soft) {
+ soft.equal(conversation.id, conversationId);
+ soft.equal(conversation.title, "conversation 1");
+ });
+});
+
+add_atomic_task(async function test_ChatStorage_findConversationsByDate() {
+ gChatStore = new ChatStore();
+
+ await addBasicConvoWithSpecificUpdatedTestData("1/1/2025", "conversation 1");
+ await addBasicConvoWithSpecificUpdatedTestData("6/1/2025", "conversation 2");
+ await addBasicConvoWithSpecificUpdatedTestData("12/1/2025", "conversation 3");
+
+ const startDate = new Date("5/1/2025").getTime();
+ const endDate = new Date("1/1/2026").getTime();
+ const conversations = await gChatStore.findConversationsByDate(
+ startDate,
+ endDate
+ );
+
+ const errorMessage = `Incorrect message sorting: ${JSON.stringify(conversations)}`;
+
+ Assert.withSoftAssertions(function (soft) {
+ soft.equal(
+ conversations.length,
+ 2,
+ "Incorrect number of conversations received"
+ );
+ soft.equal(conversations[0].title, "conversation 3", errorMessage);
+ soft.equal(conversations[1].title, "conversation 2", errorMessage);
+ });
+});
+
+add_atomic_task(async function test_ChatStorage_findConversationsByURL() {
+ async function addTestData() {
+ await addConvoWithSpecificTestData(
+ new Date("1/1/2025"),
+ new URL("https://www.firefox.com"),
+ new URL("https://www.mozilla.com"),
+ "conversation 1"
+ );
+
+ await addConvoWithSpecificTestData(
+ new Date("1/2/2025"),
+ new URL("https://www.firefox.com"),
+ new URL("https://www.mozilla.org"),
+ "Mozilla.org conversation 1"
+ );
+
+ await addConvoWithSpecificTestData(
+ new Date("1/3/2025"),
+ new URL("https://www.firefox.com"),
+ new URL("https://www.mozilla.org"),
+ "Mozilla.org conversation 2"
+ );
+ }
+
+ gChatStore = new ChatStore();
+
+ await addTestData();
+
+ const conversations = await gChatStore.findConversationsByURL(
+ new URL("https://www.mozilla.org")
+ );
+
+ Assert.withSoftAssertions(function (soft) {
+ soft.equal(conversations.length, 2, "Chat conversations not found");
+ soft.equal(conversations[0].title, "Mozilla.org conversation 2");
+ soft.equal(conversations[1].title, "Mozilla.org conversation 1");
+ });
+});
+
+async function addTestDataForFindMessageByDate() {
+ await gChatStore.updateConversation(
+ new ChatConversation({
+ title: "convo 1",
+ description: "",
+ pageUrl: new URL("https://www.firefox.com"),
+ pageMeta: {},
+ messages: [
+ new ChatMessage({
+ createdDate: new Date("1/1/2025").getTime(),
+ ordinal: 0,
+ role: 0,
+ content: { type: "text", content: "a message" },
+ pageUrl: new URL("https://www.mozilla.com"),
+ }),
+ ],
+ })
+ );
+
+ await gChatStore.updateConversation(
+ new ChatConversation({
+ title: "convo 2",
+ description: "",
+ pageUrl: new URL("https://www.firefox.com"),
+ pageMeta: {},
+ messages: [
+ new ChatMessage({
+ createdDate: new Date("7/1/2025").getTime(),
+ ordinal: 0,
+ role: 0,
+ content: { type: "text", content: "a message in july" },
+ pageUrl: new URL("https://www.mozilla.com"),
+ }),
+ ],
+ })
+ );
+
+ await gChatStore.updateConversation(
+ new ChatConversation({
+ title: "convo 3",
+ description: "",
+ pageUrl: new URL("https://www.firefox.com"),
+ pageMeta: {},
+ messages: [
+ new ChatMessage({
+ createdDate: new Date("8/1/2025").getTime(),
+ ordinal: 0,
+ role: 1,
+ content: { type: "text", content: "a message in august" },
+ pageUrl: new URL("https://www.mozilla.com"),
+ }),
+ ],
+ })
+ );
+}
+
+add_atomic_task(
+ async function test_withoutSpecifiedRole_ChatStorage_findMessagesByDate() {
+ gChatStore = new ChatStore();
+
+ await addTestDataForFindMessageByDate();
+
+ const startDate = new Date("6/1/2025");
+ const endDate = new Date("1/1/2026");
+ const messages = await gChatStore.findMessagesByDate(startDate, endDate);
+
+ Assert.withSoftAssertions(function (soft) {
+ soft.equal(messages.length, 2, "Chat messages not found");
+ soft.equal(messages?.[0]?.content?.content, "a message in august");
+ soft.equal(messages?.[1]?.content?.content, "a message in july");
+ });
+ }
+);
+
+add_atomic_task(async function test_limit_ChatStorage_findMessagesByDate() {
+ gChatStore = new ChatStore();
+
+ await addTestDataForFindMessageByDate();
+
+ const startDate = new Date("6/1/2025");
+ const endDate = new Date("1/1/2026");
+ const messages = await gChatStore.findMessagesByDate(
+ startDate,
+ endDate,
+ -1,
+ 1
+ );
+
+ Assert.withSoftAssertions(function (soft) {
+ soft.equal(messages.length, 1, "Chat messages not found");
+ soft.equal(messages?.[0]?.content?.content, "a message in august");
+ });
+});
+
+add_atomic_task(async function test_skip_ChatStorage_findMessagesByDate() {
+ gChatStore = new ChatStore();
+
+ await addTestDataForFindMessageByDate();
+
+ const startDate = new Date("6/1/2025");
+ const endDate = new Date("1/1/2026");
+ const messages = await gChatStore.findMessagesByDate(
+ startDate,
+ endDate,
+ -1,
+ -1,
+ 1
+ );
+
+ Assert.withSoftAssertions(function (soft) {
+ soft.equal(messages.length, 1, "Chat messages not found");
+ soft.equal(messages?.[0]?.content?.content, "a message in july");
+ });
+});
+
+add_atomic_task(
+ async function test_withSpecifiedRole_ChatStorage_findMessagesByDate() {
+ gChatStore = new ChatStore();
+
+ await addTestDataForFindMessageByDate();
+
+ const startDate = new Date("6/1/2025");
+ const endDate = new Date("1/1/2026");
+ const messages = await gChatStore.findMessagesByDate(startDate, endDate, 0);
+
+ Assert.withSoftAssertions(function (soft) {
+ soft.equal(messages.length, 1, "Chat messages not found");
+ soft.equal(messages?.[0]?.content?.content, "a message in july");
+ });
+ }
+);
+
+add_atomic_task(async function test_ChatStorage_searchContent() {
+ await addConvoWithSpecificTestData(
+ new Date("1/2/2025"),
+ new URL("https://www.firefox.com"),
+ new URL("https://www.mozilla.org"),
+ "Mozilla.org conversation 1",
+ "a random message"
+ );
+
+ await addConvoWithSpecificTestData(
+ new Date("1/2/2025"),
+ new URL("https://www.firefox.com"),
+ new URL("https://www.mozilla.org"),
+ "Mozilla.org conversation 2",
+ "a random message again"
+ );
+
+ await addConvoWithSpecificTestData(
+ new Date("1/2/2025"),
+ new URL("https://www.firefox.com"),
+ new URL("https://www.mozilla.org"),
+ "Mozilla.org conversation 3",
+ "the interesting message"
+ );
+
+ const conversations = await gChatStore.searchContent("body");
+
+ Assert.equal(conversations.length, 3);
+});
+
+add_atomic_task(async function test_deepPath_ChatStorage_searchContent() {
+ async function addTestData() {
+ await addConvoWithSpecificCustomContentTestData(
+ new Date("1/2/2025"),
+ new URL("https://www.firefox.com"),
+ new URL("https://www.mozilla.org"),
+ "Mozilla.org conversation 1",
+ { type: "text", content: "a random message" },
+ 0 // MessageRole.USER
+ );
+
+ await addConvoWithSpecificCustomContentTestData(
+ new Date("1/2/2025"),
+ new URL("https://www.firefox.com"),
+ new URL("https://www.mozilla.org"),
+ "Mozilla.org conversation 2",
+ { type: "text", content: "a random message again" },
+ 0 // MessageRole.USER
+ );
+
+ await addConvoWithSpecificCustomContentTestData(
+ new Date("1/2/2025"),
+ new URL("https://www.firefox.com"),
+ new URL("https://www.mozilla.org"),
+ "Mozilla.org conversation 3",
+ {
+ type: "text",
+ someKey: {
+ deeper: {
+ keyToLookIn: "the interesting message",
+ },
+ },
+ },
+ 0 // MessageRole.USER
+ );
+ }
+
+ await addTestData();
+
+ const conversations = await gChatStore.searchContent(
+ "someKey.deeper.keyToLookIn"
+ );
+
+ const foundConvo = conversations[0];
+ const firstMessage = foundConvo?.messages?.[0];
+ const contentJson = firstMessage?.content;
+
+ Assert.withSoftAssertions(function (soft) {
+ soft.equal(conversations.length, 1);
+ soft.equal(
+ contentJson?.someKey?.deeper?.keyToLookIn,
+ "the interesting message"
+ );
+ });
+});
+
+add_atomic_task(async function test_ChatStorage_search() {
+ async function addTestData() {
+ await addConvoWithSpecificTestData(
+ new Date("1/2/2025"),
+ new URL("https://www.firefox.com"),
+ new URL("https://www.mozilla.org"),
+ "Mozilla.org conversation 1",
+ "a random message"
+ );
+
+ await addConvoWithSpecificTestData(
+ new Date("1/2/2025"),
+ new URL("https://www.firefox.com"),
+ new URL("https://www.mozilla.org"),
+ "Mozilla.org interesting conversation 2",
+ "a random message again"
+ );
+
+ await addConvoWithSpecificTestData(
+ new Date("1/2/2025"),
+ new URL("https://www.firefox.com"),
+ new URL("https://www.mozilla.org"),
+ "Mozilla.org conversation 3",
+ "some other message"
+ );
+ }
+
+ await addTestData();
+
+ const conversations = await gChatStore.search("interesting");
+
+ Assert.withSoftAssertions(function (soft) {
+ soft.equal(conversations.length, 1);
+ soft.equal(
+ conversations[0].title,
+ "Mozilla.org interesting conversation 2"
+ );
+
+ const message = conversations[0].messages[0];
+ soft.equal(message.content.body, "a random message again");
+ });
+});
+
+add_atomic_task(async function test_ChatStorage_deleteConversationById() {
+ await addBasicConvoTestData("1/1/2025", "a conversation");
+
+ let conversations = await gChatStore.findRecentConversations(10);
+
+ Assert.equal(
+ conversations.length,
+ 1,
+ "Test conversation for deleteConversationById() did not save."
+ );
+
+ const conversation = conversations[0];
+
+ await gChatStore.deleteConversationById(conversation.id);
+ conversations = await gChatStore.findRecentConversations(10);
+ Assert.equal(conversations.length, 0, "Test conversation was not deleted");
+});
+
+// TODO: Disabled this test. pruneDatabase() needs some work to switch
+// db file size to be checked via dbstat. Additionally, after switching
+// the last line to `PRAGMA incremental_vacuum;` the disk storage is
+// not immediately freed, so this test is now failing. Will need to
+// revisit this test when pruneDatabase() is updated.
+//
+// add_atomic_task(async function test_ChatStorage_pruneDatabase() {
+// const initialDbSize = await gChatStore.getDatabaseSize();
+//
+// // NOTE: Add enough conversations to increase the SQLite file
+// // by a measurable size
+// for (let i = 0; i < 1000; i++) {
+// await addBasicConvoTestData("1/1/2025", "a conversation");
+// }
+//
+// const dbSizeWithTestData = await gChatStore.getDatabaseSize();
+//
+// Assert.greater(
+// dbSizeWithTestData,
+// initialDbSize,
+// "Test conversations not saved for pruneDatabase() test"
+// );
+//
+// await gChatStore.pruneDatabase(0.5, 100000);
+//
+// const dbSizeAfterPrune = await gChatStore.getDatabaseSize();
+//
+// const proximityToInitialSize = dbSizeAfterPrune - initialDbSize;
+// const proximityToTestDataSize = dbSizeWithTestData - initialDbSize;
+//
+// Assert.less(
+// proximityToInitialSize,
+// proximityToTestDataSize,
+// "The pruned size is not closer to the initial db size than it is to the size with test data in it"
+// );
+// });
+
+add_atomic_task(async function test_applyMigrations_notCalledOnInitialSetup() {
+ lazy.sinon.stub(gChatStore, "CURRENT_SCHEMA_VERSION").returns(0);
+ lazy.sinon.spy(gChatStore, "applyMigrations");
+
+ // Trigger connection to db so file creates and migrations applied
+ await gChatStore.getDatabaseSize();
+
+ Assert.ok(gChatStore.applyMigrations.notCalled);
+});
+
+add_atomic_task(
+ async function test_applyMigrations_calledOnceIfSchemaIsGreaterThanDb() {
+ lazy.sinon.stub(gChatStore, "CURRENT_SCHEMA_VERSION").get(() => 2);
+ lazy.sinon.stub(gChatStore, "getDatabaseSchemaVersion").resolves(1);
+ lazy.sinon.stub(gChatStore, "applyMigrations");
+ lazy.sinon.stub(gChatStore, "setSchemaVersion");
+
+ // Trigger connection to db so file creates and migrations applied
+ await gChatStore.getDatabaseSize();
+
+ Assert.withSoftAssertions(function (soft) {
+ soft.ok(gChatStore.applyMigrations.calledOnce);
+ soft.ok(gChatStore.setSchemaVersion.calledWith(2));
+ });
+ }
+);
+
+add_atomic_task(
+ async function test_applyMigrations_notCalledIfCurrentSchemaIsLessThanDbSchema_dbDowngrades() {
+ lazy.sinon.stub(gChatStore, "CURRENT_SCHEMA_VERSION").get(() => 1);
+ lazy.sinon.stub(gChatStore, "getDatabaseSchemaVersion").resolves(2);
+ lazy.sinon.stub(gChatStore, "applyMigrations");
+ lazy.sinon.stub(gChatStore, "setSchemaVersion");
+
+ // Trigger connection to db so file creates and migrations applied
+ await gChatStore.getDatabaseSize();
+
+ Assert.withSoftAssertions(function (soft) {
+ soft.ok(
+ gChatStore.applyMigrations.notCalled,
+ "applyMigrations was called"
+ );
+ soft.ok(
+ gChatStore.setSchemaVersion.calledWith(1),
+ "setSchemaVersion was not called with 1"
+ );
+ });
+ }
+);
diff --git a/browser/components/aiwindow/ui/test/xpcshell/test_chat-utils.js b/browser/components/aiwindow/ui/test/xpcshell/test_chat-utils.js
@@ -0,0 +1,206 @@
+/* Any copyright is dedicated to the Public Domain.
+ * http://creativecommons.org/publicdomain/zero/1.0/ */
+
+do_get_profile();
+
+const { ChatConversation, ChatMessage } = ChromeUtils.importESModule(
+ "moz-src:///browser/components/aiwindow/ui/modules/ChatStore.sys.mjs"
+);
+const {
+ makeGuid,
+ parseConversationRow,
+ parseMessageRows,
+ parseJSONOrNull,
+ getRoleLabel,
+} = ChromeUtils.importESModule(
+ "moz-src:///browser/components/aiwindow/ui/modules/ChatUtils.sys.mjs"
+);
+
+add_task(function test_makeGuid() {
+ const guid = makeGuid();
+
+ Assert.equal(guid.length, 12);
+});
+
+/**
+ * Mock row
+ */
+class RowStub {
+ constructor(data) {
+ this.data = data;
+ }
+
+ getResultByName(key) {
+ if (Object.hasOwn(this.data, key)) {
+ return this.data[key];
+ }
+
+ throw new Error("NS_ERROR_NOT_AVAILABLE");
+ }
+}
+
+add_task(function test_parseConversationRow() {
+ const now = Date.now();
+ const testRow = new RowStub({
+ conv_id: "123456789012",
+ title: "the title",
+ description: "the description",
+ page_url: "https://www.firefox.com",
+ page_meta: '{"hello": "world"}',
+ created_date: now,
+ updated_date: now,
+ status: "a status",
+ });
+
+ const conversation = parseConversationRow(testRow);
+
+ Assert.withSoftAssertions(function (soft) {
+ // eslint-disable-next-line mozilla/use-isInstance
+ soft.ok(conversation instanceof ChatConversation);
+ soft.equal(conversation.id, "123456789012");
+ soft.equal(conversation.title, "the title");
+ soft.equal(conversation.description, "the description");
+ soft.ok(URL.isInstance(conversation.pageUrl));
+ soft.equal(conversation.pageUrl.href, "https://www.firefox.com/");
+ soft.deepEqual(conversation.pageMeta, { hello: "world" });
+ soft.equal(conversation.createdDate, now);
+ soft.equal(conversation.updatedDate, now);
+ soft.equal(conversation.status, "a status");
+ });
+});
+
+add_task(function test_missingField_parseConversationRow() {
+ const now = Date.now();
+ const testRow = new RowStub({
+ title: "the title",
+ description: "the description",
+ page_url: "https://www.firefox.com",
+ page_meta: '{"hello": "world"}',
+ created_date: now,
+ updated_date: now,
+ status: "a status",
+ });
+
+ Assert.throws(function () {
+ parseConversationRow(testRow);
+ }, /NS_ERROR_NOT_AVAILABLE/);
+});
+
+add_task(function test_parseConversationRow() {
+ const now = Date.now();
+ const testRow = new RowStub({
+ message_id: "123456789012",
+ created_date: now,
+ parent_message_id: "123456",
+ revision_root_message_id: "1234",
+ ordinal: 0,
+ is_active_branch: true,
+ role: 0,
+ model_id: "a model id",
+ params: '{ "some": "data" }',
+ usage: '{ "some": "usage data" }',
+ content: '{ "some": "content data" }',
+ conv_id: "123456789012",
+ page_url: "https://www.firefox.com",
+ turn_index: 0,
+ insights_enabled: true,
+ insights_flag_source: 1,
+ insights_applied: '{ "some": "insights" }',
+ web_search_queries: '{ "some": "web search queries" }',
+ });
+
+ const rows = parseMessageRows([testRow]);
+ const message = rows[0];
+
+ Assert.withSoftAssertions(function (soft) {
+ // eslint-disable-next-line mozilla/use-isInstance
+ soft.ok(message instanceof ChatMessage);
+ soft.equal(message.id, "123456789012");
+ soft.equal(message.createdDate, now);
+ soft.equal(message.parentMessageId, "123456");
+ soft.equal(message.revisionRootMessageId, "1234");
+ soft.equal(message.ordinal, 0);
+ soft.equal(message.isActiveBranch, true);
+ soft.equal(message.role, 0);
+ soft.equal(message.modelId, "a model id");
+ soft.deepEqual(message.params, { some: "data" });
+ soft.deepEqual(message.usage, { some: "usage data" });
+ soft.deepEqual(message.content, { some: "content data" });
+ soft.deepEqual(message.convId, "123456789012");
+ soft.ok(URL.isInstance(message.pageUrl));
+ soft.deepEqual(message.pageUrl.href, "https://www.firefox.com/");
+ soft.equal(message.turnIndex, 0);
+ soft.equal(message.insightsEnabled, true);
+ soft.equal(message.insightsFlagSource, 1);
+ soft.deepEqual(message.insightsApplied, { some: "insights" });
+ soft.deepEqual(message.webSearchQueries, { some: "web search queries" });
+ });
+});
+
+add_task(function test_missingField_parseConversationRow() {
+ const now = Date.now();
+ const testRow = new RowStub({
+ //message_id: "123456789012",
+ created_date: now,
+ parent_message_id: "123456",
+ revision_root_message_id: "1234",
+ ordinal: 0,
+ is_active_branch: true,
+ role: 0,
+ model_id: "a model id",
+ params: '{ "some": "data" }',
+ usage: '{ "some": "usage data" }',
+ content: '{ "some": "content data" }',
+ conv_id: "123456789012",
+ page_url: "https://www.firefox.com",
+ turn_index: 0,
+ insights_enabled: true,
+ insights_flag_source: 1,
+ insights_applied: '{ "some": "insights" }',
+ web_search_queries: '{ "some": "web search queries" }',
+ });
+
+ Assert.throws(function () {
+ parseMessageRows([testRow]);
+ }, /NS_ERROR_NOT_AVAILABLE/);
+});
+
+add_task(function test_parseJSONOrNull() {
+ const json = '{"some": "data"}';
+
+ const parsed = parseJSONOrNull(json);
+
+ Assert.deepEqual(parsed, { some: "data" });
+});
+
+add_task(function test_invalidJson_parseJSONOrNull() {
+ const json = '{some: "data"}';
+
+ const parsed = parseJSONOrNull(json);
+
+ Assert.equal(parsed, null);
+});
+
+add_task(function test_user_getRoleLabel() {
+ const role = getRoleLabel(0);
+
+ Assert.equal(role, "User");
+});
+
+add_task(function test_assistant_getRoleLabel() {
+ const role = getRoleLabel(1);
+
+ Assert.equal(role, "Assistant");
+});
+
+add_task(function test_system_getRoleLabel() {
+ const role = getRoleLabel(2);
+
+ Assert.equal(role, "System");
+});
+
+add_task(function test_user_getRoleLabel() {
+ const role = getRoleLabel(3);
+
+ Assert.equal(role, "Tool");
+});
diff --git a/browser/components/aiwindow/ui/test/xpcshell/xpcshell.toml b/browser/components/aiwindow/ui/test/xpcshell/xpcshell.toml
@@ -0,0 +1,14 @@
+[DEFAULT]
+firefox-appdir = "browser"
+head = ["asserts.js"]
+run-if = [
+ "os != 'android'",
+]
+
+["test_ChatConversation.js"]
+
+["test_ChatMessage.js"]
+
+["test_ChatStore.js"]
+
+["test_chat-utils.js"]
