Source code
Revision control
Copy as Markdown
Other Tools
Test Info:
/* Any copyright is dedicated to the Public Domain.
"use strict";
/**
* End-to-end responsiveness of Smart Window with mocked MLPA.
*
* All flows include any on-device ML calls that could occur during a Smart Window
* user interaction. To support this and avoid fetching models and Remote Settings
* in every subtest, this reuses the taskcluster fetches from the toolkit/component/ml
* perftest suite.
*
* The metric is an interval read from the "SmartWindow" profiler markers the
* product emits, with the LLM calls stubbed to answer instantly. It tracks the
* overhead introduced in the "hot" path of Smart Window:
*
* ttft-overhead: from message submitted to first token, excluding MLPA calls
*
* It is measured in various scenarios, as some algorithms like the memory
* indexing have execution times that depend on the user profile state. The scenarios
* exercise various settings for:
*
* Window state: cold | warm
* A cold window has no ML engines running and nothing cached (like the memory embeddings)
* A warm window has warm engines and some cache is possible
* Profile:
* visits: number of entries in the history
* memories: number of memories in store
*
* Every scenario runs one unrecorded iteration before the measured ones, so
* the one-time costs of reaching it do not land in the sample.
*
* The profiler is started here unless the harness was asked for a profile, so
* `--gecko-profile` still yields a usable one. See harnessWantsProfile.
*
* To run locally, serve the models from a local hub like the ML perftests do
*
* ./mach python toolkit/components/ml/tests/tools/create_local_hub.py \
* --fetches-dir ~/ml-fetches \
* --model xenova-all-minilm-l6-v2 --model mozilla-query-intent
* MOZ_ML_LOCAL_DIR=~/ml-fetches ./mach perftest \
* --hooks toolkit/components/ml/tests/tools/hooks_local_hub.py \
* --mochitest-extra-args=headless \
* browser/components/aiwindow/ui/test/browser_perf/browser_smartwindow_perf.js
*/
const perfMetadata = {
owner: "GenAI Team",
name: "browser_smartwindow_perf.js",
description:
"User-perceived responsiveness of Smart Window across window states and profile sizes",
options: {
default: {
perfherder: true,
perfherder_metrics: [
{ name: "ttft-overhead", unit: "ms", shouldAlert: true },
],
verbose: true,
manifest: "perftest.toml",
manifest_flavor: "browser-chrome",
try_platform: ["linux", "mac", "win"],
},
},
};
requestLongerTimeout(30);
const { EngineProcess } = ChromeUtils.importESModule(
"chrome://global/content/ml/EngineProcess.sys.mjs"
);
const { MemoryStore } = ChromeUtils.importESModule(
"moz-src:///browser/components/aiwindow/services/MemoryStore.sys.mjs"
);
const { ProfilerTestUtils } = ChromeUtils.importESModule(
);
const { PlacesTestUtils } = ChromeUtils.importESModule(
);
// Lazy because mozperftest evaluates this file in node to read perfMetadata.
const lazy = {};
ChromeUtils.defineLazyGetter(
lazy,
"PlacesFrecencyRecalculator",
() =>
Cc["@mozilla.org/places/frecency-recalculator;1"].getService(Ci.nsIObserver)
.wrappedJSObject
);
const ITERATIONS = 5;
const CHUNKS = ["The ", "quick ", "brown ", "fox."];
const PROMPT = "what should I look at next?";
const LONG_PROMPT = "Summarize the following passage. ".repeat(300);
const DEEP_CONVERSATION_TURNS = 5;
const QUERY_CONTEXT_TIMEOUT_MS = 30000;
const MESSAGE_COMPLETE_EVENT = "chat-conversation:message-complete";
const METRIC = "ttft-overhead";
const TTFT_MARKER = "Time to first token (TTFT)";
const TURNAROUND_MARKER = "Total turnaround time";
const SERVER_MARKER = "ServerE2E";
const PROFILES = {
"fresh-profile": { visits: 0, memories: 0 },
"medium-profile": { visits: 2000, memories: 20 },
"large-profile": { visits: 20000, memories: 200 },
};
// Grouped by profile, which is seeded once per group.
const SCENARIOS = [
{ profile: "fresh-profile", windowState: "cold" },
{ profile: "fresh-profile", windowState: "warm" },
{ profile: "medium-profile", windowState: "cold" },
{ profile: "medium-profile", windowState: "warm" },
{ profile: "medium-profile", windowState: "warm", prompt: "long-prompt" },
{ profile: "medium-profile", windowState: "warm", prompt: "deep-chat" },
{ profile: "large-profile", windowState: "cold" },
{ profile: "large-profile", windowState: "warm" },
];
/**
* @param {object} scenario - Row from SCENARIOS.
* @returns {string} The subtest name, e.g. "ttft-overhead-cold-large-profile".
*/
function metricName(scenario) {
const parts = [METRIC, scenario.windowState, scenario.profile];
if (scenario.prompt) {
parts.push(scenario.prompt);
}
return parts.join("-");
}
/** @type {Map<string, number[]>} Subtest name to its per-iteration values. */
const journal = new Map();
/**
* Fills history and the memory store to the size the profile describes.
*
* @param {string} name - Key into PROFILES.
*/
async function seedProfile(name) {
const { visits, memories } = PROFILES[name];
await PlacesUtils.history.clear();
for (let start = 0; start < visits; start += 1000) {
const pages = [];
for (let i = start; i < Math.min(start + 1000, visits); i++) {
pages.push({
title: `${name} history entry ${i}`,
visits: [{ transition: PlacesUtils.history.TRANSITIONS.LINK }],
});
}
await PlacesUtils.history.insertMany(pages);
}
// Make sure Places is ready to be used after the insertion.
await lazy.PlacesFrecencyRecalculator.recalculateAnyOutdatedFrecencies();
await PlacesTestUtils.promiseAsyncUpdates();
for (const memory of await MemoryStore.getMemories({
includeSoftDeleted: true,
})) {
await MemoryStore.hardDeleteMemory(memory.id);
}
for (let i = 0; i < memories; i++) {
await MemoryStore.addMemory({
id: `perf-mem-${i}`,
memory_summary: `Reads about ${name} topic number ${i} regularly`,
});
}
}
/**
* @param {string} windowState - "cold" or "warm".
*/
async function applyWindowState(windowState) {
if (windowState !== "cold") {
return;
}
await EngineProcess.destroyMLEngine();
await TestUtils.waitForCondition(
() => EngineProcess.areAllEnginesTerminated(),
"Wait for the inference engines to terminate"
);
MemoryStore._clearEmbeddingsCache();
}
/**
* @param {MozBrowser} browser - The browser hosting ai-window.
* @returns {Element} The ai-window element, which runs in the parent process.
*/
function aiWindowElement(browser) {
return browser.contentDocument.querySelector("ai-window");
}
/**
* @param {ChatConversation} conversation - The conversation being answered.
* @returns {Promise<void>} Resolves when the next assistant message completes.
*/
function promiseMessageComplete(conversation) {
return new Promise(resolve => {
const onComplete = () => {
conversation.off(MESSAGE_COMPLETE_EVENT, onComplete);
resolve();
};
conversation.on(MESSAGE_COMPLETE_EVENT, onComplete);
});
}
/**
* Types a prompt, selects the chat action, submits it, and waits for the turn
* to fully process. The profiler is not read here, so serialization stays out
* of the measured path.
*
* @param {MozBrowser} browser - The browser hosting ai-window.
* @param {string} text - The prompt to submit.
*/
async function runTurn(browser, text) {
await SpecialPowers.spawn(
browser,
[text, QUERY_CONTEXT_TIMEOUT_MS],
async (prompt, timeoutMs) => {
const aiWindow = content.document.querySelector("ai-window");
const bar = await ContentTaskUtils.waitForCondition(() => {
const el = aiWindow.shadowRoot.querySelector("#ai-window-smartbar");
return el && el.inputField && el;
}, "Wait for the smartbar");
bar.value = prompt.slice(0, -1);
bar.inputField.focus();
EventUtils.sendString(prompt.slice(-1), content);
// This promise never rejects on its own, so it needs a timeout.
let timer;
try {
await Promise.race([
bar.lastQueryContextPromise,
new Promise((resolve, reject) => {
timer = content.setTimeout(
() =>
reject(new Error("Timed out waiting for the query context")),
timeoutMs
);
}),
]);
} finally {
content.clearTimeout(timer);
}
}
);
await AIWindowTestUtils.selectExplicitSmartbarAction(browser, "chat");
await AIWindowTestUtils.waitForSmartbarAction(browser, "chat");
// The action dropdown keeps focus, so the input is refocused before Enter.
await SpecialPowers.spawn(browser, [], async () => {
const aiWindow = content.document.querySelector("ai-window");
const bar = aiWindow.shadowRoot.querySelector("#ai-window-smartbar");
bar.inputField.focus();
await ContentTaskUtils.waitForCondition(
() => bar.matches(":focus-within"),
"Wait for the smartbar to accept focus"
);
});
const aiWindow = aiWindowElement(browser);
const messageComplete = promiseMessageComplete(aiWindow.conversation);
await SpecialPowers.spawn(browser, [], () => {
EventUtils.synthesizeKey("KEY_Enter", {}, content);
});
await messageComplete;
// The turnaround marker is recorded after the message completes and before
// the window stops generating.
await TestUtils.waitForCondition(
() => !aiWindow.isGenerating,
"Wait for the turn to finish"
);
}
/**
* Collects the "SmartWindow" markers of the parent process.
*
* @returns {Array<{label: string, duration: number}>}
*/
function smartWindowMarkers() {
const into = [];
for (const thread of Services.profiler.getProfileData().threads) {
const { schema, data } = thread.markers;
for (const row of data) {
if (thread.stringTable[row[schema.name]] !== "SmartWindow") {
continue;
}
const payload = row[schema.data];
let label = payload && payload.name;
if (typeof label === "number") {
label = thread.stringTable[label];
}
into.push({
label: label || "",
duration: row[schema.endTime] - row[schema.startTime],
});
}
}
return into;
}
function markerDurations(markers, label) {
return markers
.filter(marker => marker.label === label)
.map(marker => marker.duration);
}
/**
* @returns {boolean} True when a profile was requested, so the harness owns
* the profiler session. IsActive() is always true under mochitest.
*/
function harnessWantsProfile() {
return Services.env.exists("MOZ_PROFILER_SHUTDOWN");
}
/**
* @returns {object} Label to the number of markers already in the buffer.
*/
function markerBaseline() {
const markers = smartWindowMarkers();
const baseline = {};
for (const label of [TTFT_MARKER, TURNAROUND_MARKER, SERVER_MARKER]) {
baseline[label] = markerDurations(markers, label).length;
}
return baseline;
}
/**
* @param {Array<{label: string, duration: number}>} markers - Markers to read.
* @param {string} label - Marker label.
* @param {object} baseline - From markerBaseline.
* @returns {number[]} The durations, oldest first.
*/
function newMarkerDurations(markers, label, baseline) {
return markerDurations(markers, label).slice(baseline[label]);
}
/**
* Runs one iteration of a scenario in its own window.
*
* @param {object} scenario - Row from SCENARIOS.
* @returns {Promise<number>} The measured TTFT overhead, in ms.
*/
async function runScenario(scenario) {
await applyWindowState(scenario.windowState);
const ownsProfiler = !harnessWantsProfile();
const win = await AIWindowTestUtils.openReadyAIWindow();
try {
const browser = win.gBrowser.selectedBrowser;
await AIWindowTestUtils.getAichatBrowser(browser);
const prompt = scenario.prompt === "long-prompt" ? LONG_PROMPT : PROMPT;
let markers;
let turns = 0;
if (ownsProfiler) {
await ProfilerTestUtils.startProfiler({
features: ["nostacksampling"],
threads: ["GeckoMain"],
});
}
const baseline = markerBaseline();
await AIWindowTestUtils.withServer({ streamChunks: CHUNKS }, async () => {
// A warm window has already answered something.
if (scenario.windowState === "warm") {
await runTurn(browser, "hello there, what is this?");
turns++;
}
if (scenario.prompt === "deep-chat") {
for (let i = 0; i < DEEP_CONVERSATION_TURNS; i++) {
await runTurn(browser, `tell me more about topic ${i}`);
turns++;
}
}
if (scenario.windowState === "cold") {
Assert.ok(
EngineProcess.areAllEnginesTerminated(),
"The engines are still cold when the measured turn starts"
);
}
await runTurn(browser, prompt);
turns++;
markers = smartWindowMarkers();
});
Assert.ok(
!EngineProcess.areAllEnginesTerminated(),
"The on-device ML engines ran during the turn"
);
if (PROFILES[scenario.profile].memories > 0) {
Assert.ok(
!!MemoryStore.memoryEmbeddingsCache,
"The memories embedder produced embeddings"
);
}
Assert.ok(
newMarkerDurations(markers, SERVER_MARKER, baseline).length,
"The turns went through the mocked MLPA server"
);
const ttfts = newMarkerDurations(markers, TTFT_MARKER, baseline);
Assert.equal(ttfts.length, turns, "Every turn emitted a TTFT marker");
Assert.equal(
newMarkerDurations(markers, TURNAROUND_MARKER, baseline).length,
turns,
"Every turn emitted a turnaround marker"
);
// The measured turn is the last one.
return Math.round(ttfts[turns - 1]);
} finally {
if (ownsProfiler && Services.profiler.IsActive()) {
await Services.profiler.StopProfiler();
}
await BrowserTestUtils.closeWindow(win);
}
}
add_setup(async function () {
const prefs = [
["browser.smartwindow.firstrun.modelChoice", "0"],
// The memories flows are only reachable with ToS consent.
["browser.smartwindow.tos.consentTime", 123],
];
const modelHubRootUrl = Services.env.get("MOZ_MODELS_HUB");
if (modelHubRootUrl) {
prefs.push(["browser.ml.modelHubRootUrl", modelHubRootUrl]);
}
await SpecialPowers.pushPrefEnv({ set: prefs });
registerCleanupFunction(async () => {
await PlacesUtils.history.clear();
await SpecialPowers.popPrefEnv();
});
});
for (const profile of Object.keys(PROFILES)) {
const rows = SCENARIOS.filter(scenario => scenario.profile === profile);
add_task({ name: `smartwindow_${profile}` }, async function () {
await seedProfile(profile);
for (const scenario of rows) {
// Unrecorded warmup iteration.
await runScenario(scenario);
const values = [];
for (let i = 0; i < ITERATIONS; i++) {
values.push(await runScenario(scenario));
}
journal.set(metricName(scenario), values);
}
});
}
add_task(function report() {
const metrics = [...journal].map(([name, values]) => ({
name,
values,
unit: "ms",
shouldAlert: true,
}));
Assert.equal(
metrics.length,
SCENARIOS.length,
"Every scenario reported its metric"
);
for (const { name, values } of metrics) {
Assert.equal(
values.length,
ITERATIONS,
`${name} was measured on every iteration`
);
}
info(`perfMetrics | ${JSON.stringify(metrics)}`);
});