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 { ChatStore, ChatConversation } = ChromeUtils.importESModule(
"moz-src:///browser/components/aiwindow/ui/modules/ChatStore.sys.mjs"
);
const { createParserState } = ChromeUtils.importESModule(
"moz-src:///browser/components/aiwindow/models/TokenStreamParser.sys.mjs"
);
const { setTimeout } = ChromeUtils.importESModule(
"resource://gre/modules/Timer.sys.mjs"
);
// Must stay above STREAM_PERSIST_INTERVAL_MS in ChatStore.sys.mjs, so a
// coalesced write that was not cancelled has time to land.
const PAST_PERSIST_INTERVAL_MS = 1500;
/**
* See test_ChatStoreToolResults.js.
*
* @param {Function} func - The test function to run
*/
function add_atomic_task(func) {
return add_task(async function () {
Services.prefs.setBoolPref(
"browser.smartwindow.removeDatabaseOnStartup",
true
);
await ChatStore.destroyDatabase();
try {
await func();
} finally {
Services.prefs.clearUserPref(
"browser.smartwindow.removeDatabaseOnStartup"
);
await ChatStore.destroyDatabase();
}
});
}
/**
* Counts the full-conversation writes made while streaming, which is what this
* path exists to avoid.
*/
function countFullWrites() {
const original = ChatStore.updateConversation;
let count = 0;
ChatStore.updateConversation = function (...args) {
count++;
return original.apply(this, args);
};
return {
get count() {
return count;
},
restore() {
delete ChatStore.updateConversation;
},
};
}
/**
* Builds a conversation with a user turn and the empty assistant message a
* reply streams into, matching the state ai-window.mjs sets up before
* fetchWithHistory runs. Nothing is persisted, so the first chunk is also the
* conversation's first write.
*
* @returns {object}
*/
function startTurn() {
const conversation = new ChatConversation({});
conversation.addUserMessage("Tell me about Firefox", null);
conversation.addAssistantMessage("text", "");
return {
conversation,
message: conversation.messages.at(-1),
parserState: createParserState(),
};
}
function streamChunks(turn, chunks) {
for (const chunk of chunks) {
turn.conversation.handleChunk(chunk, turn.message, turn.parserState);
}
}
async function reloadAssistant(conversation, message) {
const reloaded = await ChatStore.findConversationById(conversation.id);
return reloaded?.messages.find(m => m.id === message.id);
}
add_atomic_task(async function test_chunks_coalesce_into_one_full_write() {
const turn = startTurn();
const writes = countFullWrites();
try {
streamChunks(turn, ["Firefox ", "is ", "a ", "browser", "."]);
await ChatStore.endStreamingWrites(turn.conversation.id);
Assert.equal(
writes.count,
1,
"Five chunks write the whole conversation once, not once per chunk"
);
} finally {
writes.restore();
}
const assistant = await reloadAssistant(turn.conversation, turn.message);
Assert.equal(
assistant.content.body,
"Firefox is a browser.",
"The flush persists everything streamed after the first chunk"
);
});
add_atomic_task(async function test_first_chunk_persists_the_whole_turn() {
// The narrow write's foreign keys need the conversation row and the
// streaming message's parent, and nothing persists before the first chunk.
const turn = startTurn();
streamChunks(turn, ["Hello"]);
await ChatStore.endStreamingWrites(turn.conversation.id);
const reloaded = await ChatStore.findConversationById(turn.conversation.id);
Assert.equal(
reloaded.messages.length,
2,
"The user message is persisted alongside the streaming one"
);
Assert.equal(
reloaded.messages.at(-1).content.body,
"Hello",
"The streamed body is persisted"
);
});
add_atomic_task(async function test_flush_keeps_an_interrupted_stream() {
// An aborted stream never reaches receiveResponse's updateConversation, so
// the flush is the only thing that keeps the partial reply.
const turn = startTurn();
streamChunks(turn, ["Partial ", "answer"]);
await ChatStore.endStreamingWrites(turn.conversation.id);
const assistant = await reloadAssistant(turn.conversation, turn.message);
Assert.equal(
assistant.content.body,
"Partial answer",
"An aborted stream keeps what it streamed"
);
});
add_atomic_task(async function test_narrow_write_persists_stream_tokens() {
// addTokens fills webSearchQueries and memoriesApplied as chunks arrive, so
// the narrow write has to carry them too, not just the body.
const turn = startTurn();
streamChunks(turn, ["Looking "]);
turn.message.addTokens([
{ key: "search", value: "firefox release notes" },
{ key: "existing_memory", value: "memory-1" },
]);
streamChunks(turn, ["that up"]);
await ChatStore.endStreamingWrites(turn.conversation.id);
const assistant = await reloadAssistant(turn.conversation, turn.message);
Assert.deepEqual(
assistant.webSearchQueries,
["firefox release notes"],
"Search tokens collected mid-stream are persisted"
);
Assert.deepEqual(
assistant.memoriesApplied,
["memory-1"],
"Memory tokens collected mid-stream are persisted"
);
Assert.equal(assistant.content.body, "Looking that up", "The body is intact");
});
add_atomic_task(async function test_delete_cancels_a_coalesced_write() {
const turn = startTurn();
streamChunks(turn, ["Doomed ", "reply"]);
await ChatStore.deleteConversationById(turn.conversation.id);
Assert.equal(
await ChatStore.findConversationById(turn.conversation.id),
null,
"The conversation is gone once deleted"
);
// What is under test is that no write remains scheduled, so there is no
// event to listen for; the interval has to be allowed to elapse.
// eslint-disable-next-line mozilla/no-arbitrary-setTimeout
await new Promise(resolve => setTimeout(resolve, PAST_PERSIST_INTERVAL_MS));
Assert.equal(
await ChatStore.findConversationById(turn.conversation.id),
null,
"No coalesced write fires afterwards to resurrect it"
);
});
add_atomic_task(async function test_second_message_streams_after_the_first() {
// A turn with tool calls streams into a fresh assistant message, so the
// store has to retire the previous message's coalesced write and start over.
const turn = startTurn();
streamChunks(turn, ["First ", "reply"]);
await ChatStore.endStreamingWrites(turn.conversation.id);
const second = turn.conversation.addAssistantMessage("text", "");
const secondTurn = {
conversation: turn.conversation,
message: second,
parserState: createParserState(),
};
streamChunks(secondTurn, ["Second ", "reply"]);
await ChatStore.endStreamingWrites(turn.conversation.id);
const reloaded = await ChatStore.findConversationById(turn.conversation.id);
Assert.equal(
reloaded.messages.find(m => m.id === turn.message.id).content.body,
"First reply",
"The earlier message keeps its body"
);
Assert.equal(
reloaded.messages.find(m => m.id === second.id).content.body,
"Second reply",
"The later message is persisted too"
);
});