firefox-main/browser/components/aiwindow/ui/test/xpcshell/test_ChatStoreToolResults.js (file symbol)

Source code

Revision control

Copy as Markdown

Other Tools

Test Info: Warnings

/* Any copyright is dedicated to the Public Domain.
do_get_profile();
const { ChatStore, ChatConversation } = ChromeUtils.importESModule(
"moz-src:///browser/components/aiwindow/ui/modules/ChatStore.sys.mjs"
);
/**
* 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;
async function cleanUpDatabase() {
if (gChatStore) {
await gChatStore.destroyDatabase();
gChatStore = null;
}
}
async function test_ChatStorage_setup() {
Services.prefs.setBoolPref(
"browser.smartwindow.removeDatabaseOnStartup",
true
);
gChatStore = ChatStore;
await gChatStore.destroyDatabase();
}
async function test_cleanUp() {
Services.prefs.clearUserPref("browser.smartwindow.removeDatabaseOnStartup");
await cleanUpDatabase();
}
function makeToolUIData({
toolCallId = "tool-call-1",
uiType = "website-confirmation",
tabs = [{ tabId: "tab-1", label: "Example", href: "https://example.com/" }],
} = {}) {
return {
toolCallId,
timestamp: "2026-05-13T00:00:00.000Z",
updateCount: 0,
uiType,
title: "Close these tabs?",
description: "Select tabs to close",
properties: { tabs },
};
}
add_atomic_task(async function test_toolUIData_insert_round_trip() {
const conversation = new ChatConversation({});
conversation.title = "toolUIData INSERT";
conversation.addUserMessage("Close my tabs", "https://example.com/", 0);
conversation.addAssistantMessage("text", "Here are the tabs I can close:");
const assistant = conversation.messages.at(-1);
const original = makeToolUIData({
tabs: [
{ tabId: "tab-1", label: "Page 1", href: "https://example.com/1" },
{ tabId: "tab-2", label: "Page 2", href: "https://example.com/2" },
],
});
assistant.toolUIData = original;
await gChatStore.updateConversation(conversation);
const reloaded = await gChatStore.findConversationById(conversation.id);
const reloadedAssistant = reloaded.messages.find(m => m.id === assistant.id);
Assert.ok(
reloadedAssistant,
"Reloaded conversation contains the assistant message"
);
Assert.deepEqual(
reloadedAssistant.toolUIData,
original,
"toolUIData roundTrips through the INSERT path"
);
});
add_atomic_task(async function test_toolUIData_update_roundTrip() {
const conversation = new ChatConversation({});
conversation.addUserMessage("Close my tabs", "https://example.com/", 0);
conversation.addAssistantMessage("text", "Pending confirmation");
const assistant = conversation.messages.at(-1);
assistant.toolUIData = makeToolUIData({ uiType: "website-confirmation" });
await gChatStore.updateConversation(conversation);
// Simulate ToolUI.handleUpdate mutating the in-memory object after a click
assistant.toolUIData = {
...assistant.toolUIData,
uiType: "ai-action-result",
updateCount: 1,
properties: {
...assistant.toolUIData.properties,
confirmedData: ["tab-1"],
},
};
await gChatStore.updateConversation(conversation);
const reloaded = await gChatStore.findConversationById(conversation.id);
const reloadedAssistant = reloaded.messages.find(m => m.id === assistant.id);
Assert.withSoftAssertions(soft => {
soft.equal(
reloadedAssistant.toolUIData.uiType,
"ai-action-result",
"uiType reflects the post-confirm mutation"
);
soft.equal(
reloadedAssistant.toolUIData.updateCount,
1,
"updateCount reflects the post-confirm mutation"
);
soft.deepEqual(
reloadedAssistant.toolUIData.properties.confirmedData,
["tab-1"],
"confirmedData persisted through the ON CONFLICT UPDATE branch"
);
});
});
add_atomic_task(async function test_toolUIData_null_roundTrip() {
const conversation = new ChatConversation({});
conversation.addUserMessage("Just a message", "https://example.com/", 0);
conversation.addAssistantMessage("text", "Just a reply");
const assistant = conversation.messages.at(-1);
// toolUIData intentionally not set
await gChatStore.updateConversation(conversation);
const reloaded = await gChatStore.findConversationById(conversation.id);
const reloadedAssistant = reloaded.messages.find(m => m.id === assistant.id);
Assert.strictEqual(
reloadedAssistant.toolUIData,
null,
"Messages without toolUIData reload as null"
);
});
add_atomic_task(async function test_toolUIData_undoDismissed_roundTrip() {
const conversation = new ChatConversation({});
conversation.addUserMessage("Close my tabs", "https://example.com/", 0);
conversation.addAssistantMessage("text", "Closed");
const assistant = conversation.messages.at(-1);
const base = makeToolUIData({ uiType: "ai-action-result" });
assistant.toolUIData = {
...base,
properties: { ...base.properties, undoDismissed: true },
};
await gChatStore.updateConversation(conversation);
const reloaded = await gChatStore.findConversationById(conversation.id);
const reloadedAssistant = reloaded.messages.find(m => m.id === assistant.id);
Assert.strictEqual(
reloadedAssistant.toolUIData.properties.undoDismissed,
true,
"undoDismissed:true survives the ChatStore roundTrip"
);
});
function makeHistoryResults() {
return [
{
title: "Page 1",
visitDate: 1700000000000000,
visitCount: 3,
timestamp: "Yesterday",
},
{
title: "Page 2",
visitDate: 1700000100000000,
visitCount: 1,
timestamp: "Yesterday",
},
];
}
add_atomic_task(async function test_historyResults_insert_round_trip() {
const conversation = new ChatConversation({});
conversation.title = "historyResults INSERT";
conversation.addUserMessage("Show my history", "https://example.com/", 0);
conversation.addAssistantMessage("text", "Here is what I found:");
const assistant = conversation.messages.at(-1);
const original = makeHistoryResults();
assistant.historyResults = original;
await gChatStore.updateConversation(conversation);
const reloaded = await gChatStore.findConversationById(conversation.id);
const reloadedAssistant = reloaded.messages.find(m => m.id === assistant.id);
Assert.ok(
reloadedAssistant,
"Reloaded conversation contains the assistant message"
);
Assert.deepEqual(
reloadedAssistant.historyResults,
original,
"historyResults roundTrips through the INSERT path"
);
});
add_atomic_task(async function test_historyResults_update_roundTrip() {
const conversation = new ChatConversation({});
conversation.addUserMessage("Show my history", "https://example.com/", 0);
conversation.addAssistantMessage("text", "Searching...");
// The message row first persists while still streaming, with no snapshot yet.
const assistant = conversation.messages.at(-1);
await gChatStore.updateConversation(conversation);
// When that same message completes, receiveResponse writes its snapshot,
// re-persisting the existing row through the ON CONFLICT UPDATE branch.
const snapshot = makeHistoryResults();
assistant.historyResults = snapshot;
await gChatStore.updateConversation(conversation);
const reloaded = await gChatStore.findConversationById(conversation.id);
const reloadedAssistant = reloaded.messages.find(m => m.id === assistant.id);
Assert.deepEqual(
reloadedAssistant.historyResults,
snapshot,
"historyResults snapshot persisted through the ON CONFLICT UPDATE branch"
);
});
add_atomic_task(async function test_historyResults_empty_roundTrip() {
const conversation = new ChatConversation({});
conversation.addUserMessage("Just a message", "https://example.com/", 0);
conversation.addAssistantMessage("text", "Just a reply");
const assistant = conversation.messages.at(-1);
// historyResults intentionally left at its default empty array
await gChatStore.updateConversation(conversation);
const reloaded = await gChatStore.findConversationById(conversation.id);
const reloadedAssistant = reloaded.messages.find(m => m.id === assistant.id);
Assert.deepEqual(
reloadedAssistant.historyResults,
[],
"Messages without historyResults reload as an empty array"
);
});
add_atomic_task(async function test_historyResults_rehydrates_pool() {
const conversation = new ChatConversation({});
conversation.addUserMessage("Show my history", "https://example.com/", 0);
conversation.addAssistantMessage("text", "Here is what I found:");
const assistant = conversation.messages.at(-1);
const original = makeHistoryResults();
assistant.historyResults = original;
await gChatStore.updateConversation(conversation);
const reloaded = await gChatStore.findConversationById(conversation.id);
Assert.deepEqual(
reloaded.getHistoryResultsSnapshot(),
original,
"Reloaded conversation rehydrates its history results pool from messages"
);
});
function makeCitations() {
return [
{ url: "https://example.com/1", title: "Source 1", hasFavicon: true },
{ url: "https://example.com/2", title: "Source 2" },
];
}
function makeStrippedCitations() {
return [
{ url: "https://example.com/1", title: "Source 1" },
{ url: "https://example.com/2", title: "Source 2" },
];
}
add_atomic_task(async function test_citations_insert_roundTrip() {
const conversation = new ChatConversation({});
conversation.title = "citations INSERT";
conversation.addUserMessage("Search the web", "https://example.com/", 0);
conversation.addAssistantMessage("text", "Here is what I found:");
const assistant = conversation.messages.at(-1);
assistant.citations = makeCitations();
await gChatStore.updateConversation(conversation);
const reloaded = await gChatStore.findConversationById(conversation.id);
const reloadedAssistant = reloaded.messages.find(m => m.id === assistant.id);
Assert.ok(
reloadedAssistant,
"Reloaded conversation contains the assistant message"
);
Assert.deepEqual(
reloadedAssistant.citations,
makeStrippedCitations(),
"citations roundTrip through the INSERT path without resolved assets"
);
});
add_atomic_task(async function test_citations_update_roundTrip() {
const conversation = new ChatConversation({});
conversation.addUserMessage("Search the web", "https://example.com/", 0);
conversation.addAssistantMessage("text", "Searching...");
const assistant = conversation.messages.at(-1);
await gChatStore.updateConversation(conversation);
assistant.citations = makeCitations();
await gChatStore.updateConversation(conversation);
const reloaded = await gChatStore.findConversationById(conversation.id);
const reloadedAssistant = reloaded.messages.find(m => m.id === assistant.id);
Assert.deepEqual(
reloadedAssistant.citations,
makeStrippedCitations(),
"citations snapshot persisted through the ON CONFLICT UPDATE branch"
);
});
add_atomic_task(async function test_citations_empty_roundTrip() {
const conversation = new ChatConversation({});
conversation.addUserMessage("Just a message", "https://example.com/", 0);
conversation.addAssistantMessage("text", "Just a reply");
const assistantId = conversation.messages.at(-1).id;
await gChatStore.updateConversation(conversation);
const reloaded = await gChatStore.findConversationById(conversation.id);
Assert.deepEqual(
reloaded.messages.find(m => m.id === assistantId).citations,
[],
"Messages without citations reload as an empty array"
);
});
add_atomic_task(async function test_citations_rehydrates_pool() {
const conversation = new ChatConversation({});
conversation.addUserMessage("Search the web", "https://example.com/", 0);
conversation.addAssistantMessage("text", "Here is what I found:");
const assistant = conversation.messages.at(-1);
assistant.citations = makeCitations();
await gChatStore.updateConversation(conversation);
const reloaded = await gChatStore.findConversationById(conversation.id);
// The snapshot only covers URLs read this turn
Assert.deepEqual(
reloaded.getCitationsSnapshot(),
[],
"A reloaded conversation has no pending citations of its own"
);
reloaded.addCitations([{ url: "https://example.com/1" }]);
Assert.deepEqual(
reloaded.getCitationsSnapshot(),
[{ url: "https://example.com/1", title: "Source 1" }],
"The pool rehydrated from message snapshots, so the title carries forward"
);
});
add_atomic_task(
async function test_updateLLMTelemetryRecord_creates_unprocessed_row() {
const conversation = new ChatConversation({});
conversation.title = "conversation with llm telemetry";
conversation.addUserMessage("test content", "https://www.firefox.com");
await gChatStore.updateConversation(conversation);
await gChatStore.updateLLMTelemetryRecord(conversation.id);
const telemetry = await gChatStore.findLLMTelemetryByConversationId(
conversation.id
);
Assert.ok(telemetry, "LLM telemetry row should exist");
Assert.withSoftAssertions(function (soft) {
soft.equal(telemetry.convId, conversation.id);
soft.equal(telemetry.processed, 0);
soft.deepEqual(telemetry.telemetryPrompts, {});
soft.deepEqual(telemetry.telemetryProbabilities, {});
soft.ok(telemetry.processedTime, "processedTime should be set");
});
}
);
add_atomic_task(
async function test_updateLLMTelemetryRecord_creates_processed_row() {
const conversation = new ChatConversation({});
conversation.title = "processed llm telemetry conversation";
conversation.addUserMessage("test content", "https://www.firefox.com");
await gChatStore.updateConversation(conversation);
await gChatStore.updateLLMTelemetryRecord(
conversation.id,
{
"wasSuccessful-v1": 2,
"isLongConvo-v1": 2,
},
{
"wasSuccessful-v1": 0.9,
"isLongConvo-v1": 0.84,
},
0,
1
);
const telemetry = await gChatStore.findLLMTelemetryByConversationId(
conversation.id
);
Assert.ok(telemetry, "LLM telemetry row should exist");
Assert.withSoftAssertions(function (soft) {
soft.equal(telemetry.convId, conversation.id);
soft.equal(telemetry.processed, 1);
soft.deepEqual(telemetry.telemetryPrompts, {
"wasSuccessful-v1": 2,
"isLongConvo-v1": 2,
});
soft.deepEqual(telemetry.telemetryProbabilities, {
"wasSuccessful-v1": 0.9,
"isLongConvo-v1": 0.84,
});
soft.ok(telemetry.processedTime, "processedTime should be set");
});
}
);
add_atomic_task(
async function test_updateLLMTelemetryRecord_merges_prompts_and_probabilities() {
const conversation = new ChatConversation({});
conversation.title = "merged llm telemetry conversation";
conversation.addUserMessage("test content", "https://www.firefox.com");
await gChatStore.updateConversation(conversation);
await gChatStore.updateLLMTelemetryRecord(
conversation.id,
{
"wasSuccessful-v1": 2,
"isLongConvo-v1": 2,
},
{
"wasSuccessful-v1": 0.9,
"isLongConvo-v1": 0.84,
},
0,
0
);
await gChatStore.updateLLMTelemetryRecord(
conversation.id,
{
"isLongConvo-v1": 8,
},
{
"isLongConvo-v1": 0.95,
},
0,
1
);
const telemetry = await gChatStore.findLLMTelemetryByConversationId(
conversation.id
);
Assert.ok(telemetry, "LLM telemetry row should exist");
Assert.withSoftAssertions(function (soft) {
soft.equal(telemetry.convId, conversation.id);
soft.equal(telemetry.processed, 1);
soft.deepEqual(telemetry.telemetryPrompts, {
"wasSuccessful-v1": 2,
"isLongConvo-v1": 8,
});
soft.deepEqual(telemetry.telemetryProbabilities, {
"wasSuccessful-v1": 0.9,
"isLongConvo-v1": 0.84,
});
soft.ok(telemetry.processedTime, "processedTime should be set");
});
}
);
add_atomic_task(
async function test_updateLLMTelemetryRecord_preserves_existing_data_when_marking_unprocessed() {
const conversation = new ChatConversation({});
conversation.title = "unprocessed preserves telemetry";
conversation.addUserMessage("test content", "https://www.firefox.com");
await gChatStore.updateConversation(conversation);
await gChatStore.updateLLMTelemetryRecord(
conversation.id,
{
"wasSuccessful-v1": 2,
"isLongConvo-v1": 8,
},
{
"wasSuccessful-v1": 0.9,
"isLongConvo-v1": 0.95,
},
0,
1
);
await gChatStore.updateLLMTelemetryRecord(conversation.id, {}, {}, 0);
const telemetry = await gChatStore.findLLMTelemetryByConversationId(
conversation.id
);
Assert.ok(telemetry, "LLM telemetry row should exist");
Assert.withSoftAssertions(function (soft) {
soft.equal(telemetry.convId, conversation.id);
soft.equal(telemetry.processed, 0);
soft.deepEqual(telemetry.telemetryPrompts, {
"wasSuccessful-v1": 2,
"isLongConvo-v1": 8,
});
soft.deepEqual(telemetry.telemetryProbabilities, {
"wasSuccessful-v1": 0.9,
"isLongConvo-v1": 0.95,
});
soft.ok(telemetry.processedTime, "processedTime should be set");
});
}
);
add_atomic_task(
async function test_findLLMTelemetryByConversationId_returns_null_for_missing_row() {
const telemetry =
await gChatStore.findLLMTelemetryByConversationId("missing-conv-id");
Assert.equal(
telemetry,
null,
"Should return null when no LLM telemetry row exists"
);
}
);
add_atomic_task(
async function test_updateLLMTelemetryRecord_sets_uniform_sampling_probability() {
const conversation = new ChatConversation({});
conversation.title = "uniform sampling probability conversation";
conversation.addUserMessage("test content", "https://www.firefox.com");
await gChatStore.updateConversation(conversation);
await gChatStore.updateLLMTelemetryRecord(conversation.id, {}, {}, 750, 0);
const telemetry = await gChatStore.findLLMTelemetryByConversationId(
conversation.id
);
Assert.ok(telemetry, "LLM telemetry row should exist");
Assert.equal(telemetry.uniformSamplingProbability, 750);
}
);
add_atomic_task(
async function test_updateLLMTelemetryRecord_preserves_uniform_sampling_probability() {
const conversation = new ChatConversation({});
conversation.title = "uniform sampling probability preserved conversation";
conversation.addUserMessage("test content", "https://www.firefox.com");
await gChatStore.updateConversation(conversation);
await gChatStore.updateLLMTelemetryRecord(conversation.id, {}, {}, 750, 0);
await gChatStore.updateLLMTelemetryRecord(conversation.id, {}, {}, 999, 1);
const telemetry = await gChatStore.findLLMTelemetryByConversationId(
conversation.id
);
Assert.ok(telemetry, "LLM telemetry row should exist");
Assert.equal(
telemetry.uniformSamplingProbability,
750,
"uniform_sampling_probability should not be overwritten on update"
);
}
);
add_atomic_task(
async function test_findConversationById_hydratesUniformSamplingState() {
const conversation = new ChatConversation({});
conversation.title = "hydration conversation";
conversation.addUserMessage("test content", "https://www.firefox.com");
await gChatStore.updateConversation(conversation);
await gChatStore.updateLLMTelemetryRecord(conversation.id, {}, {}, 0.25, 0);
const reloaded = await gChatStore.findConversationById(conversation.id);
Assert.equal(
reloaded._telemetryUniformSample,
true,
"_telemetryUniformSample is rehydrated from llm_telemetry on reload"
);
Assert.equal(
reloaded._telemetryUniformProbability,
0.25,
"_telemetryUniformProbability is rehydrated from llm_telemetry on reload"
);
}
);
add_atomic_task(
async function test_findConversationById_skipsHydrationWhenNotSampled() {
const conversation = new ChatConversation({});
conversation.title = "no-hydration conversation";
conversation.addUserMessage("test content", "https://www.firefox.com");
await gChatStore.updateConversation(conversation);
await gChatStore.updateLLMTelemetryRecord(conversation.id, {}, {}, 0, 0);
const reloaded = await gChatStore.findConversationById(conversation.id);
Assert.notStrictEqual(
reloaded._telemetryUniformSample,
true,
"_telemetryUniformSample stays unset when uniform_sampling_probability is 0"
);
}
);
add_atomic_task(
async function test_findConversationById_skipsHydrationWhenNoTelemetryRow() {
const conversation = new ChatConversation({});
conversation.title = "no-telemetry-row conversation";
conversation.addUserMessage("test content", "https://www.firefox.com");
await gChatStore.updateConversation(conversation);
const reloaded = await gChatStore.findConversationById(conversation.id);
Assert.notStrictEqual(
reloaded._telemetryUniformSample,
true,
"_telemetryUniformSample stays unset when no llm_telemetry row exists"
);
}
);
/**
* Test that messages with website-confirmation toolUIData get isRestored flag
* when loaded from the database
*/
add_atomic_task(async function test_website_confirmation_isRestored_flag() {
const conversation = new ChatConversation({});
conversation.title = "Test isRestored flag";
conversation.addUserMessage("Close some tabs", "https://example.com/", 0);
conversation.addAssistantMessage("text", "I'll help close those tabs");
const assistant = conversation.messages.at(-1);
// Add website-confirmation toolUIData
const toolUIData = makeToolUIData({
uiType: "website-confirmation",
tabs: [
{ id: "tab-1", url: "https://example.com", title: "Example" },
{ id: "tab-2", url: "https://test.com", title: "Test" },
],
});
// Add originalUserPrompt to properties
toolUIData.properties.originalUserPrompt = "Close some tabs";
assistant.toolUIData = toolUIData;
// Save the conversation
await gChatStore.updateConversation(conversation);
// Load it back from the database
const reloaded = await gChatStore.findConversationById(conversation.id);
const reloadedAssistant = reloaded.messages.find(m => m.id === assistant.id);
// Verify the isRestored flag was set
Assert.ok(
reloadedAssistant.isRestored,
"Messages with website-confirmation toolUIData should have isRestored flag set when loaded from DB"
);
// Verify the toolUIData is preserved
Assert.equal(
reloadedAssistant.toolUIData.uiType,
"website-confirmation",
"uiType should be preserved"
);
Assert.equal(
reloadedAssistant.toolUIData.properties.originalUserPrompt,
"Close some tabs",
"originalUserPrompt should be preserved"
);
Assert.equal(
reloadedAssistant.toolUIData.properties.tabs.length,
2,
"tabs array should be preserved"
);
});
/**
* Test that messages with other UI types don't get isRestored flag
*/
add_atomic_task(async function test_other_ui_types_no_isRestored_flag() {
const conversation = new ChatConversation({});
conversation.title = "Test no isRestored flag";
conversation.addUserMessage("Do something", "https://example.com/", 0);
conversation.addAssistantMessage("text", "Task completed");
const assistant = conversation.messages.at(-1);
// Add ai-action-result toolUIData (not website-confirmation)
assistant.toolUIData = makeToolUIData({
uiType: "ai-action-result",
});
// Save the conversation
await gChatStore.updateConversation(conversation);
// Load it back from the database
const reloaded = await gChatStore.findConversationById(conversation.id);
const reloadedAssistant = reloaded.messages.find(m => m.id === assistant.id);
// Verify the isRestored flag was NOT set
Assert.ok(
!reloadedAssistant.isRestored,
"Messages with non-website-confirmation toolUIData should NOT have isRestored flag"
);
// Verify the toolUIData is still preserved
Assert.equal(
reloadedAssistant.toolUIData.uiType,
"ai-action-result",
"uiType should be preserved"
);
});