commit c47db865e2e826e5bbc7e46d1452c691910831d2
parent 99020e68b307489f5e68deaae9f12d66fa5b9d5b
Author: Nick Grato <ngrato@gmail.com>
Date: Mon, 17 Nov 2025 23:05:51 +0000
Bug 1998270 - hook up mozilla account + token passing api key flow r=pdahiya,firefox-ai-ml-reviewers
For first iteration just adding token as an option to the openai pipeline and calling it through the smart assist feature.
Differential Revision: https://phabricator.services.mozilla.com/D271303
Diffstat:
6 files changed, 175 insertions(+), 6 deletions(-)
diff --git a/browser/components/genai/SmartAssistEngine.sys.mjs b/browser/components/genai/SmartAssistEngine.sys.mjs
@@ -11,6 +11,11 @@ ChromeUtils.defineESModuleGetters(lazy, {
/* eslint-disable-next-line mozilla/reject-import-system-module-from-non-system */
import { createEngine } from "chrome://global/content/ml/EngineProcess.sys.mjs";
+import { getFxAccountsSingleton } from "resource://gre/modules/FxAccounts.sys.mjs";
+import {
+ OAUTH_CLIENT_ID,
+ SCOPE_PROFILE,
+} from "resource://gre/modules/FxAccountsCommon.sys.mjs";
const toolsConfig = [
{
@@ -72,6 +77,20 @@ export const SmartAssistEngine = {
_createEngine: createEngine,
+ async _getFxAccountToken() {
+ try {
+ const fxAccounts = getFxAccountsSingleton();
+ const token = await fxAccounts.getOAuthToken({
+ scope: SCOPE_PROFILE,
+ client_id: OAUTH_CLIENT_ID,
+ });
+ return token;
+ } catch (error) {
+ console.error("Error obtaining FxA token:", error);
+ throw error;
+ }
+ },
+
/**
* Creates an OpenAI engine instance configured with Smart Assists preferences.
*
@@ -107,6 +126,7 @@ export const SmartAssistEngine = {
*/
async *fetchWithHistory(messages) {
const engineInstance = await this.createOpenAIEngine();
+ const fxAccountToken = await this._getFxAccountToken();
// We'll mutate a local copy of the thread as we loop
let convo = Array.isArray(messages) ? [...messages] : [];
@@ -115,6 +135,7 @@ export const SmartAssistEngine = {
const streamModelResponse = () =>
engineInstance.runWithGenerator({
streamOptions: { enabled: true },
+ fxAccountToken,
tool_choice: "auto",
tools: toolsConfig,
args: convo,
diff --git a/browser/components/genai/content/smart-assist.mjs b/browser/components/genai/content/smart-assist.mjs
@@ -14,6 +14,8 @@ ChromeUtils.defineESModuleGetters(lazy, {
SmartAssistEngine:
"moz-src:///browser/components/genai/SmartAssistEngine.sys.mjs",
PrivateBrowsingUtils: "resource://gre/modules/PrivateBrowsingUtils.sys.mjs",
+ SpecialMessageActions:
+ "resource://messaging-system/lib/SpecialMessageActions.sys.mjs",
});
const FULL_PAGE_URL = "chrome://browser/content/genai/smartAssistPage.html";
@@ -226,6 +228,25 @@ export class SmartAssist extends MozLitElement {
this._applyNewTabOverride(isChecked);
}
+ /**
+ * Initiates the Firefox Account sign-in flow for MLPA authentication.
+ */
+
+ _signIn() {
+ lazy.SpecialMessageActions.handleAction(
+ {
+ type: "FXA_SIGNIN_FLOW",
+ data: {
+ entrypoint: "ai-window",
+ extraParams: {
+ service: "ai-window",
+ },
+ },
+ },
+ window.browsingContext.topChromeWindow.gBrowser.selectedBrowser
+ );
+ }
+
render() {
const iconSrc = this.showLog
? "chrome://global/skin/icons/arrow-down.svg"
@@ -315,7 +336,17 @@ export class SmartAssist extends MozLitElement {
>
${this.inputAction.label}
</moz-button>
+ <hr/>
+ <h3>The following Elements are for testing purposes</h3>
+ <p>Sign in for MLPA authentication.</p>
+ <moz-button
+ type="primary"
+ size="small"
+ @click=${this._signIn}
+ >
+ Sign in
+ </moz-button>
<!-- Footer - New Tab Override -->
${
this.mode === "sidebar"
diff --git a/browser/components/genai/tests/xpcshell/test_smart_assist_engine.js b/browser/components/genai/tests/xpcshell/test_smart_assist_engine.js
@@ -83,6 +83,7 @@ add_task(async function test_fetchWithHistory_streams_and_forwards_args() {
};
sb.stub(SmartAssistEngine, "_createEngine").resolves(fakeEngine);
+ sb.stub(SmartAssistEngine, "_getFxAccountToken").resolves("mock_token");
const messages = [
{ role: "system", content: "You are helpful" },
@@ -121,6 +122,7 @@ add_task(
try {
const err = new Error("creation failed (generic)");
const stub = sb.stub(SmartAssistEngine, "_createEngine").rejects(err);
+ sb.stub(SmartAssistEngine, "_getFxAccountToken").resolves("mock_token");
const messages = [{ role: "user", content: "Hi" }];
// Must CONSUME the async generator to trigger the rejection
@@ -157,6 +159,7 @@ add_task(async function test_fetchWithHistory_propagates_stream_error() {
},
};
sb.stub(SmartAssistEngine, "_createEngine").resolves(fakeEngine);
+ sb.stub(SmartAssistEngine, "_getFxAccountToken").resolves("mock_token");
const consume = async () => {
let acc = "";
diff --git a/toolkit/components/ml/content/backends/OpenAIPipeline.mjs b/toolkit/components/ml/content/backends/OpenAIPipeline.mjs
@@ -327,9 +327,16 @@ export class OpenAIPipeline {
lazy.console.debug("Running OpenAI pipeline");
try {
const { baseURL, apiKey, modelId } = this.#options;
+ const fxAccountToken = request.fxAccountToken
+ ? request.fxAccountToken
+ : null;
+ const defaultHeaders = fxAccountToken
+ ? { Authorization: `Bearer ${fxAccountToken}` }
+ : undefined;
const client = new OpenAIPipeline.OpenAILib.OpenAI({
baseURL: baseURL ? baseURL : "http://localhost:11434/v1",
apiKey: apiKey || "ollama",
+ ...(defaultHeaders ? { defaultHeaders } : {}),
});
const stream = request.streamOptions?.enabled || false;
const tools = request.tools || [];
@@ -349,11 +356,9 @@ export class OpenAIPipeline {
port,
};
- if (stream) {
- return await this.#handleStreamingResponse(args);
- }
-
- return await this.#handleNonStreamingResponse(args);
+ return stream
+ ? await this.#handleStreamingResponse(args)
+ : await this.#handleNonStreamingResponse(args);
} catch (error) {
const backendError = this.#errorFactory(error);
port?.postMessage({ done: true, ok: false, error: backendError });
diff --git a/toolkit/components/ml/tests/browser/browser_ml_openai.js b/toolkit/components/ml/tests/browser/browser_ml_openai.js
@@ -251,3 +251,104 @@ add_task(async function test_openai_client_tools_streaming() {
await stopMockOpenAI(mockServer);
}
});
+
+/**
+ * Test that fxAccountToken is properly passed to the OpenAI client
+ */
+add_task(async function test_openai_fxaccount_token() {
+ const records = [
+ {
+ ...BASE_ENGINE_OPTIONS,
+ id: "74a71cfd-1734-44e6-85c0-69cf3e874138",
+ },
+ ];
+ const { cleanup } = await setup({ records });
+
+ // Mock server that checks for fxAccount token in headers
+ let capturedFxaHeader = null;
+ const { server: mockServer, port } = startMockOpenAI({
+ echo: "Response with FxA token",
+ onRequest: req => {
+ try {
+ if (req.hasHeader("authorization")) {
+ capturedFxaHeader = req.getHeader("authorization");
+ }
+ } catch (e) {
+ info("Failed to get authorization header: " + e);
+ }
+ },
+ });
+
+ const fxAccountToken = "test_fxa_token_12345";
+
+ const engineInstance = await createEngine({
+ ...BASE_ENGINE_OPTIONS,
+ apiKey: "test-api-key",
+ baseURL: `http://localhost:${port}/v1`,
+ backend: "openai",
+ });
+
+ const request = {
+ args: [
+ {
+ role: "user",
+ content: "Test request with FxA token",
+ },
+ ],
+ fxAccountToken,
+ };
+
+ try {
+ info("Run the inference with fxAccountToken");
+ const result = await engineInstance.run(request);
+
+ Assert.equal(
+ result.finalOutput,
+ "This is a mock summary for testing end-to-end flow.",
+ "Should get expected response"
+ );
+
+ // Verify that the FxA token was sent in the request headers
+ const expectedValue = `Bearer ${fxAccountToken}`;
+
+ Assert.equal(
+ capturedFxaHeader,
+ expectedValue,
+ `FxA token should be included in request headers. Expected: ${expectedValue}, Got: ${capturedFxaHeader}`
+ );
+
+ info("Test without fxAccountToken - should not include header");
+
+ // Create another engine without fxAccountToken
+ const engineWithoutToken = await createEngine({
+ ...BASE_ENGINE_OPTIONS,
+ apiKey: "test-api-key",
+ baseURL: `http://localhost:${port}/v1`,
+ backend: "openai",
+ // No fxAccountToken
+ });
+
+ capturedFxaHeader = null; // Reset captured header
+ const requestWithoutToken = {
+ args: [
+ {
+ role: "user",
+ content: "Test request with FxA token",
+ },
+ ],
+ // No fxAccountToken
+ };
+ await engineWithoutToken.run(requestWithoutToken);
+
+ // Verify Authorization header is the API key when no FxA token is provided
+ Assert.equal(
+ capturedFxaHeader,
+ "Bearer test-api-key",
+ "Authorization should fall back to API key when no FxA token is provided"
+ );
+ } finally {
+ await EngineProcess.destroyMLEngine();
+ await cleanup();
+ await stopMockOpenAI(mockServer);
+ }
+});
diff --git a/toolkit/components/ml/tests/browser/head.js b/toolkit/components/ml/tests/browser/head.js
@@ -783,12 +783,20 @@ function readRequestBody(request) {
});
}
-function startMockOpenAI({ echo = "This gets echoed." } = {}) {
+function startMockOpenAI({
+ echo = "This gets echoed.",
+ onRequest = null,
+} = {}) {
const server = new HttpServer();
server.registerPathHandler("/v1/chat/completions", (request, response) => {
info("GET /v1/chat/completions");
+ // Call the onRequest callback if provided to allow test inspection
+ if (onRequest) {
+ onRequest(request);
+ }
+
let bodyText = "";
if (request.method === "POST") {
try {
