Source code
Revision control
Copy as Markdown
Other Tools
Test Info: Warnings
- This test runs only with pattern: os != 'android'
- Manifest: browser/components/aiwindow/ui/test/xpcshell/xpcshell.toml
/* Any copyright is dedicated to the Public Domain.
do_get_profile();
const { ChatConversation, MESSAGE_ROLE, ChatMessage } =
ChromeUtils.importESModule(
"moz-src:///browser/components/aiwindow/ui/modules/ChatStore.sys.mjs"
);
const { SecurityProperties } = ChromeUtils.importESModule(
"moz-src:///browser/components/aiwindow/models/SecurityProperties.sys.mjs"
);
const { MEMORIES_FLAG_SOURCE, SYSTEM_PROMPT_TYPE } = ChromeUtils.importESModule(
"moz-src:///browser/components/aiwindow/ui/modules/ChatEnums.sys.mjs"
);
const { UserRoleOpts, AssistantRoleOpts, ToolRoleOpts } =
ChromeUtils.importESModule(
"moz-src:///browser/components/aiwindow/ui/modules/ChatMessage.sys.mjs"
);
const { MemoryStore } = ChromeUtils.importESModule(
"moz-src:///browser/components/aiwindow/services/MemoryStore.sys.mjs"
);
const { MemoriesManager } = ChromeUtils.importESModule(
"moz-src:///browser/components/aiwindow/models/memories/MemoriesManager.sys.mjs"
);
const { EmbeddingsGenerator } = ChromeUtils.importESModule(
"chrome://global/content/ml/EmbeddingsGenerator.sys.mjs"
);
const lazy = {};
ChromeUtils.defineESModuleGetters(lazy, {
sinon: "resource://testing-common/Sinon.sys.mjs",
});
add_task(function test_ChatConversation_constructor_defaults() {
const conversation = new ChatConversation({});
Assert.withSoftAssertions(function (soft) {
soft.equal(conversation.id.length, 36);
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",
userContext: {},
};
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",
userContext: {},
};
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",
userContext: {},
};
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",
userContext: {},
};
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",
userContext: {},
};
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, new URL("https://www.mozilla.com"));
const message = conversation.messages[0];
Assert.withSoftAssertions(function (soft) {
soft.equal(message.role, MESSAGE_ROLE.USER);
soft.equal(message.turnIndex, 1);
soft.deepEqual(message.pageUrl, new URL("https://www.mozilla.com"));
soft.deepEqual(message.content, {
type: "text",
body: "user to assistant msg",
userContext: {},
contextPageUrl: "https://www.mozilla.com/",
});
});
});
add_task(function test_revisionRootMessageId_ChatConversation_addUserMessage() {
const conversation = new ChatConversation({});
const content = "user to assistant msg";
conversation.addUserMessage(content, "https://www.firefox.com");
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,
new UserRoleOpts({ revisionRootMessageId: "321" })
);
const message = conversation.messages[0];
Assert.equal(message.revisionRootMessageId, "321");
});
add_task(async function test_userContext_ChatConversation_addUserMessage() {
const conversation = new ChatConversation({});
const content = "user to assistant msg";
const userContext = { testContextInfo: "123" };
conversation.addUserMessage(
content,
new UserRoleOpts({ revisionRootMessageId: "321" }),
userContext
);
const message = conversation.messages[0];
Assert.deepEqual(message.content.userContext, userContext);
});
add_task(function test_contextPageUrl_ChatConversation_addUserMessage() {
const conversation = new ChatConversation({});
conversation.addUserMessage(
"user msg",
new URL("https://www.mozilla.com/page")
);
const message = conversation.messages[0];
Assert.equal(message.content.contextPageUrl, "https://www.mozilla.com/page");
});
add_task(function test_noContextPageUrl_ChatConversation_addUserMessage() {
const conversation = new ChatConversation({});
conversation.addUserMessage("user msg", null);
const message = conversation.messages[0];
Assert.ok(!("contextPageUrl" in message.content));
});
add_task(function test_ChatConversation_addAssistantMessage() {
const conversation = new ChatConversation({});
const content = "response from assistant";
conversation.addAssistantMessage("text", content);
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.memoriesEnabled,
false,
"memoriesEnabled should default to false"
);
soft.strictEqual(
message.memoriesFlagSource,
null,
"memoriesFlagSource should default to null"
);
soft.deepEqual(
message.memoriesApplied,
[],
"memoriesApplied 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,
["memory"],
["search"]
);
conversation.addAssistantMessage("text", content, 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.memoriesEnabled,
true,
"memoriesEnabled should equal true"
);
soft.strictEqual(
message.memoriesFlagSource,
1,
"memoriesFlagSource equal 1"
);
soft.deepEqual(
message.memoriesApplied,
["memory"],
"memoriesApplied should equal ['memory']"
);
soft.deepEqual(
message.webSearchQueries,
["search"],
"memoriesApplied should equal ['search']"
);
});
});
add_task(function test_ChatConversation_addToolCallMessage() {
const conversation = new ChatConversation({});
const content = {
random: "tool call specific keys",
};
conversation.addToolCallMessage(content);
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, 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);
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, new URL("https://www.mozilla.com"));
conversation.addUserMessage(content, new URL("https://www.mozilla.com"));
conversation.addUserMessage(content, new URL("https://www.firefox.com"));
conversation.addUserMessage(content, new URL("https://www.cnn.com"));
conversation.addUserMessage(content, new URL("https://www.espn.com"));
conversation.addUserMessage(content, new URL("https://www.espn.com"));
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, new URL("https://www.mozilla.com"));
conversation.addUserMessage(content, new URL("https://www.mozilla.com"));
conversation.addUserMessage(content, new URL("https://www.firefox.com"));
conversation.addUserMessage(content, new URL("https://www.cnn.com"));
conversation.addUserMessage(content, new URL("https://www.espn.com"));
conversation.addUserMessage(content, new URL("https://www.espn.com"));
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, new URL("about:aiwindow"));
conversation.addUserMessage(content, null);
conversation.addUserMessage(content, null);
const mostRecentPageVisited = conversation.getMostRecentPageVisited();
Assert.equal(mostRecentPageVisited, null);
});
add_task(function test_ChatConversation_renderState() {
const conversation = new ChatConversation({});
const content = "user to assistant msg";
conversation.addUserMessage(content, "about:aiwindow");
conversation.addToolCallMessage("some content");
conversation.addAssistantMessage("text", "a response");
conversation.addUserMessage(content, "about:aiwindow");
conversation.addSystemMessage("text", "some system message");
conversation.addAssistantMessage("text", "a response");
const renderState = conversation.renderState();
Assert.deepEqual(renderState, [
conversation.messages[0],
conversation.messages[2],
conversation.messages[3],
conversation.messages[5],
]);
});
add_task(function test_ChatConversation_currentTurnIndex() {
const conversation = new ChatConversation({});
const content = "user to assistant msg";
conversation.addSystemMessage("text", "the system prompt");
conversation.addUserMessage(content, "about:aiwindow");
conversation.addAssistantMessage("text", "a response");
conversation.addUserMessage(content, "about:aiwindow");
conversation.addAssistantMessage("text", "a response");
conversation.addUserMessage(content, "about:aiwindow");
conversation.addAssistantMessage("text", "a response");
conversation.addUserMessage(content, "about:aiwindow");
conversation.addAssistantMessage("text", "a response");
conversation.addUserMessage(content, "about:aiwindow");
conversation.addAssistantMessage("text", "a response");
Assert.deepEqual(conversation.currentTurnIndex(), 4);
});
add_task(function test_ChatConversation_helpersTurnIndexing() {
const conversation = new ChatConversation({});
conversation.addSystemMessage("text", "the system prompt");
conversation.addUserMessage("a user's prompt", "https://www.somesite.com");
conversation.addToolCallMessage({ some: "tool call details" });
conversation.addAssistantMessage("text", "the llm response");
conversation.addUserMessage(
"a user's second prompt",
);
conversation.addToolCallMessage({ some: "more tool call details" });
conversation.addAssistantMessage("text", "the second llm response");
Assert.withSoftAssertions(function (soft) {
soft.equal(conversation.messages.length, 7);
soft.equal(conversation.messages[0].turnIndex, 0);
soft.equal(conversation.messages[1].turnIndex, 0);
soft.equal(conversation.messages[2].turnIndex, 0);
soft.equal(conversation.messages[3].turnIndex, 0);
soft.equal(conversation.messages[4].turnIndex, 1);
soft.equal(conversation.messages[5].turnIndex, 1);
soft.equal(conversation.messages[6].turnIndex, 1);
});
});
add_task(function test_ChatConversation_getMessagesInOpenAiFormat() {
const conversation = new ChatConversation({});
conversation.addSystemMessage("text", "the system prompt");
conversation.addUserMessage(
"a user's prompt",
new UserRoleOpts(),
{ testContext: "321" }
);
conversation.addToolCallMessage({
tool_call_id: "123",
name: "tool_1",
body: [1, 2, 3],
});
conversation.addAssistantMessage("text", "the llm response");
conversation.addUserMessage(
"a user's second prompt",
"some question",
new UserRoleOpts(),
{ testContext: "654" }
);
conversation.addToolCallMessage({
tool_call_id: "456",
name: "tool_1",
body: [4, 5, 6],
});
conversation.addAssistantMessage("text", "the second llm response");
const openAiFormat = conversation.getMessagesInOpenAiFormat();
Assert.deepEqual(openAiFormat, [
{ role: "system", content: "the system prompt" },
{ role: "user", content: "a user's prompt" },
{ role: "tool", content: "[1,2,3]", name: "tool_1", tool_call_id: "123" },
{ role: "assistant", content: "the llm response" },
{ role: "user", content: "654" },
{ role: "user", content: "a user's second prompt" },
{ role: "tool", content: "[4,5,6]", name: "tool_1", tool_call_id: "456" },
{ role: "assistant", content: "the second llm response" },
]);
});
add_task(async function test_unrelatedMessage_ChatConversation_retryMessage() {
const conversation = new ChatConversation({});
conversation.addSystemMessage("text", "the system prompt");
conversation.addUserMessage("a user's prompt", "https://www.somesite.com");
const unrelatedMessage = new ChatMessage({
ordinal: 0,
role: MESSAGE_ROLE.USER,
content: "some content",
turnIndex: 0,
});
await Assert.rejects(
conversation.retryMessage(unrelatedMessage),
/Unrelated message/
);
});
add_task(async function test_nonUserMessage_ChatConversation_retryMessage() {
const conversation = new ChatConversation({});
conversation.addSystemMessage("text", "the system prompt");
conversation.addUserMessage("a user's prompt", "https://www.somesite.com");
await Assert.rejects(
conversation.retryMessage(conversation.messages[0]),
/Not a user message/
);
});
add_task(
async function test_ChatConversation_retryMessage_returnsRemovedMessages() {
let sandbox = lazy.sinon.createSandbox();
const conversation = new ChatConversation({});
sandbox.stub(ChatConversation, "getRealTimeInfo").callsFake(() => {
conversation.addSystemMessage(
SYSTEM_PROMPT_TYPE.REAL_TIME,
"real time data"
);
});
sandbox.stub(conversation, "getMemoriesContext").callsFake(() => {
conversation.addSystemMessage(
SYSTEM_PROMPT_TYPE.MEMORIES,
"memories data"
);
});
conversation.addSystemMessage("text", "the system prompt");
conversation.addUserMessage("a user's prompt", "https://www.somesite.com");
conversation.addToolCallMessage({ some: "tool call details" });
conversation.addAssistantMessage("text", "the llm response");
conversation.addUserMessage("a user's second prompt", "some question");
conversation.addToolCallMessage({ some: "more tool call details" });
conversation.addAssistantMessage("text", "the second llm response");
const toDeleteMessages = await conversation.retryMessage(
conversation.messages[1]
);
Assert.withSoftAssertions(function (soft) {
soft.equal(toDeleteMessages.length, 6, "Incorrect number of messages");
soft.equal(toDeleteMessages[0].content.body, "a user's prompt");
soft.equal(toDeleteMessages[1].content.some, "tool call details");
soft.equal(toDeleteMessages[2].content.body, "the llm response");
soft.equal(toDeleteMessages[3].content.body, "a user's second prompt");
soft.equal(toDeleteMessages[4].content.some, "more tool call details");
soft.equal(toDeleteMessages[5].content.body, "the second llm response");
});
sandbox.restore();
}
);
add_task(async function test_filtersEphemeral_ChatConversation_retryMessage() {
const conversation = new ChatConversation({});
conversation.addSystemMessage(SYSTEM_PROMPT_TYPE.TEXT, "the system prompt");
conversation.addSystemMessage(SYSTEM_PROMPT_TYPE.REAL_TIME, "real time data");
conversation.addSystemMessage(SYSTEM_PROMPT_TYPE.MEMORIES, "memories data");
conversation.addUserMessage("a user's prompt", "https://www.somesite.com");
conversation.addAssistantMessage("text", "the llm response");
const retryTarget = conversation.messages.find(
m => m.role === MESSAGE_ROLE.USER
);
const deleted = await conversation.retryMessage(retryTarget);
Assert.withSoftAssertions(function (soft) {
soft.equal(
conversation.messages.length,
1,
"Only the base system prompt should remain"
);
soft.equal(
conversation.messages[0].content.type,
SYSTEM_PROMPT_TYPE.TEXT,
"Remaining message should be the base system prompt"
);
soft.equal(
deleted.length,
4,
"Should return ephemeral messages and spliced messages"
);
soft.equal(
deleted[0].content.type,
SYSTEM_PROMPT_TYPE.REAL_TIME,
"First deleted should be the real time message"
);
soft.equal(
deleted[1].content.type,
SYSTEM_PROMPT_TYPE.MEMORIES,
"Second deleted should be the memories message"
);
soft.equal(
deleted[2].content.body,
"a user's prompt",
"Third deleted should be the retried user message"
);
soft.equal(
deleted[3].content.body,
"the llm response",
"Fourth deleted should be the assistant message"
);
});
});
add_task(
async function test_uniqueOrdinalsWithoutMemories_ChatConversation_retryMessage() {
const conversation = new ChatConversation({});
conversation.addSystemMessage(SYSTEM_PROMPT_TYPE.TEXT, "the system prompt");
conversation.addSystemMessage(
SYSTEM_PROMPT_TYPE.REAL_TIME,
"real time data"
);
conversation.addSystemMessage(SYSTEM_PROMPT_TYPE.MEMORIES, "memories data");
conversation.addUserMessage("a user's prompt", "https://www.somesite.com");
conversation.addAssistantMessage("text", "the llm response");
const retryTarget = conversation.messages.find(
m => m.role === MESSAGE_ROLE.USER
);
const originalUserOrdinal = retryTarget.ordinal;
await conversation.retryMessage(retryTarget);
conversation.addSystemMessage(
SYSTEM_PROMPT_TYPE.REAL_TIME,
"new real time data"
);
conversation.addUserMessage("a user's prompt", "https://www.somesite.com");
conversation.addAssistantMessage("text", "the new llm response");
const ordinals = conversation.messages.map(m => m.ordinal);
const uniqueOrdinals = new Set(ordinals);
Assert.withSoftAssertions(function (soft) {
soft.equal(
ordinals.length,
uniqueOrdinals.size,
"All ordinals should be unique after retry without memories"
);
soft.ok(
conversation.messages
.filter(m => m.role === MESSAGE_ROLE.ASSISTANT)
.at(-1).ordinal > originalUserOrdinal,
"New assistant ordinal must be greater than original user ordinal"
);
});
}
);
add_task(async function test_returnsContent_ChatConversation_getRealTimeInfo() {
console.log(Object.keys(lazy.sinon));
const mockGetRealTimeMapping = lazy.sinon.stub().resolves({
todayDate: "2024-01-15",
url: "https://example.com",
title: "Example",
hasTabInfo: false,
locale: "en-US",
timezone: "America/Los_Angeles",
isoTimestamp: "2024-01-15T10:30:00",
});
const mockEngineInstance = {
loadPrompt: lazy.sinon
.stub()
.resolves("Current date: {todayDate}\nLocale: {locale}"),
};
const realTimeInfo = await ChatConversation.getRealTimeInfo(
mockEngineInstance,
{
getRealTimeMapping: mockGetRealTimeMapping,
}
);
Assert.withSoftAssertions(function (soft) {
soft.ok(
mockEngineInstance.loadPrompt.called,
"loadPrompt should be called"
);
});
Assert.equal(
realTimeInfo,
"Current date: 2024-01-15\nLocale: en-US",
"getRealTimeInfo returns the expected contexutal information"
);
});
add_task(
async function test_returnsNoContent_ChatConversation_getRealTimeInfo() {
const mockEngineInstance = {
loadPrompt: lazy.sinon.stub().resolves("prompt text"),
};
const mockGetRealTimeMapping = lazy.sinon.stub().resolves(null);
const realTimeInfo = await ChatConversation.getRealTimeInfo(
mockEngineInstance,
{
getRealTimeMapping: mockGetRealTimeMapping,
}
);
Assert.equal(
realTimeInfo,
null,
"getRealTimeInfo returns null if constructRealTime returns an empty object"
);
}
);
add_task(
async function test_returnsContent_ChatConversation_getMemoriesContext() {
console.log(Object.keys(lazy.sinon));
const mockEngineInstance = {
loadPrompt: lazy.sinon.stub().resolves("prompt text"),
};
const constructMemories = lazy.sinon
.stub()
.resolves({ content: "memories data" });
const conversation = new ChatConversation({});
const memoriesContext = await conversation.getMemoriesContext(
"hello",
mockEngineInstance,
constructMemories,
new SecurityProperties()
);
Assert.withSoftAssertions(function (soft) {
soft.ok(
constructMemories.calledWith("hello", mockEngineInstance),
"constructMemories should be called with message and engineInstance"
);
});
Assert.equal(
memoriesContext,
"memories data",
"getMemoriesContext returns the expected memories information"
);
}
);
add_task(
async function test_returnsNoContent_ChatConversation_getMemoriesContext() {
console.log(Object.keys(lazy.sinon));
const constructMemories = lazy.sinon.stub().resolves({});
const conversation = new ChatConversation({});
const memoriesContext = await conversation.getMemoriesContext(
"hello",
constructMemories
);
Assert.equal(
memoriesContext,
null,
"getMemoriesContext returns null if constructMemories returns an empty object"
);
}
);
add_task(function test_ChatConversation_renderState_filters_phantom_messages() {
const conversation = new ChatConversation({});
conversation.addUserMessage("What's the weather?", "about:aiwindow");
conversation.addAssistantMessage("text", "");
conversation.addAssistantMessage("function", {
tool_calls: [
{
id: "call_1",
function: {
name: "run_search",
arguments: '{"query":"weather"}',
},
},
],
});
conversation.addAssistantMessage("text", "Here is the weather forecast.");
const renderState = conversation.renderState();
Assert.equal(
renderState.length,
2,
"Should only contain user message and real assistant message"
);
Assert.equal(renderState[0].role, MESSAGE_ROLE.USER);
Assert.equal(renderState[1].role, MESSAGE_ROLE.ASSISTANT);
Assert.equal(renderState[1].content.body, "Here is the weather forecast.");
});
add_task(
async function test_deduplicatesMemoryIds_ChatConversation_receiveResponse() {
let sandbox = lazy.sinon.createSandbox();
const mockMemories = [{ id: "mem-1" }, { id: "mem-2" }];
sandbox.stub(MemoryStore, "getMemories").resolves(mockMemories);
const conversation = new ChatConversation({});
conversation.addAssistantMessage("text", "some response");
const assistantMsg = conversation.messages.at(-1);
assistantMsg._pendingMemoryIds = ["mem-1", "mem-1", "mem-2", "mem-2"];
async function* emptyStream() {}
await conversation.receiveResponse(emptyStream());
Assert.ok(
MemoryStore.getMemories.calledOnce,
"MemoryStore.getMemories should be called exactly once"
);
const { memoryIds } = MemoryStore.getMemories.firstCall.args[0];
Assert.equal(
memoryIds.size,
2,
"memoryIds should be a deduplicated Set of size 2"
);
Assert.ok(memoryIds.has("mem-1"), "memoryIds should contain mem-1");
Assert.ok(memoryIds.has("mem-2"), "memoryIds should contain mem-2");
Assert.deepEqual(
assistantMsg.memoriesApplied,
mockMemories,
"memoriesApplied should be set to the resolved memories"
);
Assert.ok(
!("_pendingMemoryIds" in assistantMsg),
"_pendingMemoryIds should be deleted after processing"
);
sandbox.restore();
}
);
add_task(async function test_addUserMessage_sets_memories_fields() {
const conversation = new ChatConversation({});
const userOpts = new UserRoleOpts({
memoriesEnabled: false,
memoriesFlagSource: MEMORIES_FLAG_SOURCE.CONVERSATION,
});
await conversation.addUserMessage("hello", null, userOpts);
const lastUserMessage = conversation.messages
.filter(m => m.role === MESSAGE_ROLE.USER)
.at(-1);
Assert.ok(lastUserMessage, "Last user message exists");
Assert.equal(
lastUserMessage.memoriesEnabled,
false,
"memoriesEnabled is persisted on the user message"
);
Assert.equal(
lastUserMessage.memoriesFlagSource,
MEMORIES_FLAG_SOURCE.CONVERSATION,
"memoriesFlagSource is persisted on the user message"
);
});
add_task(async function test_generatePrompt_emitsUserMessage() {
const sandbox = lazy.sinon.createSandbox();
const conversation = new ChatConversation({});
const mockEngineInstance = {
loadPrompt: lazy.sinon.stub().resolves("system prompt"),
};
sandbox.stub(ChatConversation, "getRealTimeInfo").resolves(null);
sandbox.stub(conversation, "getMemoriesContext").resolves(null);
let emittedMessage = null;
conversation.on("chat-conversation:message-update", (_, msg) => {
emittedMessage = msg;
});
await conversation.generatePrompt("hello", null, mockEngineInstance);
Assert.ok(emittedMessage, "event should have been emitted");
Assert.equal(emittedMessage.content.body, "hello");
Assert.equal(emittedMessage.role, MESSAGE_ROLE.USER);
sandbox.restore();
});
add_task(async function test_generatePrompt_skipUserDispatch() {
const sandbox = lazy.sinon.createSandbox();
const conversation = new ChatConversation({});
const mockEngineInstance = {
loadPrompt: lazy.sinon.stub().resolves("system prompt"),
};
sandbox.stub(ChatConversation, "getRealTimeInfo").resolves(null);
sandbox.stub(conversation, "getMemoriesContext").resolves(null);
let emitted = false;
conversation.on("chat-conversation:message-update", () => {
emitted = true;
});
await conversation.generatePrompt(
"hello",
null,
mockEngineInstance,
undefined,
true
);
Assert.ok(
!emitted,
"event should not be emitted when skipUserDispatch is true"
);
sandbox.restore();
});
add_task(async function test_generatePrompt_memoriesContextErrorDoesNotThrow() {
let sandbox = lazy.sinon.createSandbox();
// Seed a memory so getRelevantMemories reaches the embeddings path
await MemoryStore.addMemory({
id: "memory-embed-fail",
memory_summary: "User likes hiking",
category: "preference",
intent: "profile",
reasoning: "Test memory",
score: 0.5,
updated_at: Date.now(),
is_deleted: false,
});
MemoriesManager._clearEmbeddingsCache();
sandbox
.stub(EmbeddingsGenerator.prototype, "embedMany")
.rejects(new Error("Failed to download embedding model"));
const conversation = new ChatConversation({});
const mockEngineInstance = {
loadPrompt: sandbox.stub().resolves("system prompt"),
};
sandbox
.stub(ChatConversation, "getRealTimeInfo")
.resolves("real time context");
const result = await conversation.generatePrompt(
"hello",
null,
mockEngineInstance,
{ memoriesEnabled: true }
);
Assert.ok(result, "generatePrompt should resolve successfully");
const userMessage = conversation.messages.find(
m => m.role === MESSAGE_ROLE.USER
);
Assert.ok(
EmbeddingsGenerator.prototype.embedMany.calledOnce,
"embedMany should have been called exactly once"
);
Assert.equal(
userMessage.content.userContext.realTimeContext,
"real time context",
"realTimeContext should still be set despite embeddings failure"
);
Assert.ok(
!("memoriesContext" in userMessage.content.userContext),
"memoriesContext should not be set when embedMany rejects"
);
await MemoryStore.hardDeleteMemory("memory-embed-fail", "other");
MemoriesManager._clearEmbeddingsCache();
sandbox.restore();
});
add_task(
async function test_generatePrompt_userContextPopulatedBeforeResolving() {
const sandbox = lazy.sinon.createSandbox();
const conversation = new ChatConversation({});
const mockEngineInstance = {
loadPrompt: lazy.sinon.stub().resolves("system prompt"),
};
sandbox
.stub(ChatConversation, "getRealTimeInfo")
.resolves("real time context");
sandbox
.stub(conversation, "getMemoriesContext")
.resolves("memories context");
await conversation.generatePrompt("hello", null, mockEngineInstance, {
memoriesEnabled: true,
});
const userMessage = conversation.messages.find(
m => m.role === MESSAGE_ROLE.USER
);
Assert.withSoftAssertions(function (soft) {
soft.equal(
userMessage.content.userContext.realTimeContext,
"real time context",
"realTimeContext should be set on userContext before generatePrompt resolves"
);
soft.equal(
userMessage.content.userContext.memoriesContext,
"memories context",
"memoriesContext should be set on userContext before generatePrompt resolves"
);
});
sandbox.restore();
}
);
add_task(async function test_getRealTimeInfo_setsPrivateData_when_hasTabInfo() {
const securityProperties = new SecurityProperties();
const mockGetRealTimeMapping = lazy.sinon.stub().resolves({
todayDate: "2024-01-15",
url: "https://example.com",
title: "Example Page",
hasTabInfo: true,
locale: "en-US",
timezone: "America/Los_Angeles",
isoTimestamp: "2024-01-15T10:30:00",
});
const mockEngineInstance = {
loadPrompt: lazy.sinon.stub().resolves("{todayDate}"),
};
await ChatConversation.getRealTimeInfo(mockEngineInstance, {
getRealTimeMapping: mockGetRealTimeMapping,
securityProperties,
});
securityProperties.commit();
Assert.ok(
securityProperties.privateData,
"privateData should be true after commit when hasTabInfo is true"
);
});
add_task(
async function test_getRealTimeInfo_doesNotSetPrivateData_when_noTabInfo() {
const securityProperties = new SecurityProperties();
const mockGetRealTimeMapping = lazy.sinon.stub().resolves({
todayDate: "2024-01-15",
hasTabInfo: false,
locale: "en-US",
timezone: "America/Los_Angeles",
isoTimestamp: "2024-01-15T10:30:00",
});
const mockEngineInstance = {
loadPrompt: lazy.sinon.stub().resolves("{todayDate}"),
};
await ChatConversation.getRealTimeInfo(mockEngineInstance, {
getRealTimeMapping: mockGetRealTimeMapping,
securityProperties,
});
securityProperties.commit();
Assert.ok(
!securityProperties.privateData,
"privateData should remain false when hasTabInfo is false"
);
}
);
add_task(
async function test_getMemoriesContext_setsPrivateData_when_memoriesFound() {
const securityProperties = new SecurityProperties();
const constructMemories = lazy.sinon
.stub()
.resolves({ content: "some memory" });
const mockEngineInstance = {};
const conversation = new ChatConversation({});
await conversation.getMemoriesContext(
"hello",
mockEngineInstance,
constructMemories,
securityProperties
);
securityProperties.commit();
Assert.ok(
securityProperties.privateData,
"privateData should be true after commit when memories were found"
);
}
);
add_task(
async function test_getMemoriesContext_doesNotSetPrivateData_when_noMemories() {
const securityProperties = new SecurityProperties();
const constructMemories = lazy.sinon.stub().resolves(null);
const mockEngineInstance = {};
const conversation = new ChatConversation({});
await conversation.getMemoriesContext(
"hello",
mockEngineInstance,
constructMemories,
securityProperties
);
securityProperties.commit();
Assert.ok(
!securityProperties.privateData,
"privateData should remain false when no memories were found"
);
}
);
add_task(
async function test_generatePrompt_commitsPrivateData_when_hasTabInfo() {
const mockEngineInstance = {
loadPrompt: lazy.sinon.stub().resolves("system prompt"),
};
const conversation = new ChatConversation({});
const sandbox = lazy.sinon.createSandbox();
sandbox
.stub(ChatConversation, "getRealTimeInfo")
.callsFake(async (_, opts) => {
opts.securityProperties?.setPrivateData();
return "real time info";
});
sandbox.stub(conversation, "getMemoriesContext").resolves(null);
await conversation.generatePrompt("hello", null, mockEngineInstance);
Assert.ok(
conversation.securityProperties.privateData,
"privateData should be committed true when getRealTimeInfo stages it"
);
sandbox.restore();
}
);
add_task(
async function test_generatePrompt_commitsPrivateData_when_memoriesEnabled() {
const mockEngineInstance = {
loadPrompt: lazy.sinon.stub().resolves("system prompt"),
};
const conversation = new ChatConversation({});
const sandbox = lazy.sinon.createSandbox();
sandbox.stub(ChatConversation, "getRealTimeInfo").resolves(null);
sandbox
.stub(conversation, "getMemoriesContext")
.callsFake(async (_, _engine, _construct, sp) => {
sp?.setPrivateData();
return "some memories";
});
await conversation.generatePrompt("hello", null, mockEngineInstance, {
memoriesEnabled: true,
});
Assert.ok(
conversation.securityProperties.privateData,
"privateData should be committed true when getMemoriesContext stages it"
);
sandbox.restore();
}
);
add_task(
async function test_generatePrompt_doesNotSetPrivateData_when_noTabOrMemories() {
const mockEngineInstance = {
loadPrompt: lazy.sinon.stub().resolves("system prompt"),
};
const conversation = new ChatConversation({});
const sandbox = lazy.sinon.createSandbox();
sandbox.stub(ChatConversation, "getRealTimeInfo").resolves(null);
sandbox.stub(conversation, "getMemoriesContext").resolves(null);
await conversation.generatePrompt("hello", null, mockEngineInstance);
Assert.ok(
!conversation.securityProperties.privateData,
"privateData should remain false when no private data was staged"
);
sandbox.restore();
}
);
add_task(function test_securityProperties_plainObject_normalization() {
const conversation = new ChatConversation({
securityProperties: { untrustedInput: true },
});
Assert.withSoftAssertions(function (soft) {
soft.ok(
conversation.securityProperties instanceof SecurityProperties,
"securityProperties should be a SecurityProperties instance"
);
soft.equal(
conversation.securityProperties.untrustedInput,
true,
"untrustedInput should be true when explicitly set"
);
soft.equal(
conversation.securityProperties.privateData,
false,
"privateData should default to false when missing from input"
);
});
});