tor-browser

commit 4507193fa852a804507e2019737bb0de2ea06200
parent f1dc23ae227b7948e8cf9c12eae098df74971144
Author: Punam Dahiya <pdahiya@mozilla.com>
Date:   Thu, 11 Dec 2025 21:37:37 +0000

Bug 2005345 - Add service type in OpenAIPipeline request header r=ai-ondevice-reviewers,ngrato

Differential Revision: https://phabricator.services.mozilla.com/D275879

Diffstat:
M
toolkit/components/ml/content/backends/OpenAIPipeline.mjs
 | 
5
++++-
M
toolkit/components/ml/tests/browser/browser_ml_openai.js
 | 
21
+++++++++++++++++++--

2 files changed, 23 insertions(+), 3 deletions(-)
diff --git a/toolkit/components/ml/content/backends/OpenAIPipeline.mjs b/toolkit/components/ml/content/backends/OpenAIPipeline.mjs
@@ -331,7 +331,10 @@ export class OpenAIPipeline {
         ? request.fxAccountToken
         : null;
       const defaultHeaders = fxAccountToken
-        ? { Authorization: `Bearer ${fxAccountToken}` }
+        ? {
+            Authorization: `Bearer ${fxAccountToken}`,
+            "service-type": "ai",
+          }
         : undefined;
       const client = new OpenAIPipeline.OpenAILib.OpenAI({
         baseURL: baseURL ? baseURL : "http://localhost:11434/v1",
diff --git a/toolkit/components/ml/tests/browser/browser_ml_openai.js b/toolkit/components/ml/tests/browser/browser_ml_openai.js
@@ -266,6 +266,7 @@ add_task(async function test_openai_fxaccount_token() {
 
   // Mock server that checks for fxAccount token in headers
   let capturedFxaHeader = null;
+  let capturedServiceTypeHeader = null;
   const { server: mockServer, port } = startMockOpenAI({
     echo: "Response with FxA token",
     onRequest: req => {
@@ -273,8 +274,11 @@ add_task(async function test_openai_fxaccount_token() {
         if (req.hasHeader("authorization")) {
           capturedFxaHeader = req.getHeader("authorization");
         }
+        if (req.hasHeader("service-type")) {
+          capturedServiceTypeHeader = req.getHeader("service-type");
+        }
       } catch (e) {
-        info("Failed to get authorization header: " + e);
+        info("Failed to get header: " + e);
       }
     },
   });
@@ -317,6 +321,12 @@ add_task(async function test_openai_fxaccount_token() {
       `FxA token should be included in request headers. Expected: ${expectedValue}, Got: ${capturedFxaHeader}`
     );
 
+    Assert.equal(
+      capturedServiceTypeHeader,
+      "ai",
+      "service-type header should be 'ai' when FxA token is provided"
+    );
+
     info("Test without fxAccountToken - should not include header");
 
     // Create another engine without fxAccountToken
@@ -328,7 +338,8 @@ add_task(async function test_openai_fxaccount_token() {
       // No fxAccountToken
     });
 
-    capturedFxaHeader = null; // Reset captured header
+    capturedFxaHeader = null;
+    capturedServiceTypeHeader = null;
     const requestWithoutToken = {
       args: [
         {
@@ -346,6 +357,12 @@ add_task(async function test_openai_fxaccount_token() {
       "Bearer test-api-key",
       "Authorization should fall back to API key when no FxA token is provided"
     );
+
+    Assert.equal(
+      capturedServiceTypeHeader,
+      null,
+      "service-type header should not be present when no FxA token is provided"
+    );
   } finally {
     await EngineProcess.destroyMLEngine();
     await cleanup();