Bug 2021177 - Adds FxA implementation of the MlpaTokenProvider

mobile/android/android-components/components/lib/llm-mlpa/src/main/java/mozilla/components/lib/llm/mlpa/service/FetchClientMlpaService.kt

@@ -74,7 +74,7 @@ class FetchClientMlpaService(
             url = "${config.baseUrl}/v1/chat/completions",
             method = Request.Method.POST,
             headers = MutableHeaders(
-                "authorization" to authorizationToken.value,
+                "authorization" to "Bearer ${authorizationToken.value}",
                 "content-type" to "application/json",
                 "service-type" to "s2s",
             ),

mobile/android/fenix/app/src/main/java/org/mozilla/fenix/components/Components.kt

@@ -49,6 +49,8 @@ import org.mozilla.fenix.components.appstate.AppAction
 import org.mozilla.fenix.components.appstate.AppState
 import org.mozilla.fenix.components.appstate.setup.checklist.SetupChecklistState
 import org.mozilla.fenix.components.appstate.setup.checklist.getSetupChecklistCollection
+import org.mozilla.fenix.components.llm.Llm
+import org.mozilla.fenix.components.llm.ext.accessTokenProvider
 import org.mozilla.fenix.components.metrics.MetricsMiddleware
 import org.mozilla.fenix.crashes.CrashReportingAppMiddleware
 import org.mozilla.fenix.crashes.SettingsCrashReportCache
@@ -410,7 +412,12 @@ class Components(private val context: Context) {
     }
 
     val llm: Llm by lazyMonitored {
-        Llm(core.client, integrityClient, clientUUID)
+        Llm(
+            client = core.client,
+            fxaTokenProvider = backgroundServices.accountManager.accessTokenProvider,
+            integrityClient = integrityClient,
+            userIdProvider = clientUUID,
+        )
     }
 
     val clientUUID by lazyMonitored { ClientUUID.build(context) }

mobile/android/fenix/app/src/main/java/org/mozilla/fenix/components/Llm.kt → mobile/android/fenix/app/src/main/java/org/mozilla/fenix/components/llm/FenixMlpaService.kt

@@ -2,22 +2,15 @@
  * License, v. 2.0. If a copy of the MPL was not distributed with this
  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
 
-package org.mozilla.fenix.components
+package org.mozilla.fenix.components.llm
 
 import mozilla.components.concept.fetch.Client
-import mozilla.components.concept.integrity.IntegrityClient
-import mozilla.components.lib.llm.mlpa.MlpaLlmProvider
-import mozilla.components.lib.llm.mlpa.MlpaTokenProvider
-import mozilla.components.lib.llm.mlpa.UserIdProvider
 import mozilla.components.lib.llm.mlpa.service.AuthenticationService
 import mozilla.components.lib.llm.mlpa.service.AuthorizationToken
 import mozilla.components.lib.llm.mlpa.service.ChatService
 import mozilla.components.lib.llm.mlpa.service.FetchClientMlpaService
 import mozilla.components.lib.llm.mlpa.service.MlpaConfig
 import mozilla.components.lib.llm.mlpa.service.MlpaService
-import mozilla.components.lib.llm.mlpa.service.PackageName
-import org.mozilla.fenix.BuildConfig
-import org.mozilla.fenix.perf.lazyMonitored
 
 /**
  * Temporary class for toggling between prod and nonprod MLPA environment
@@ -36,26 +29,3 @@ class FenixMlpaService(
         request: ChatService.Request,
     ) = service.completion(authorizationToken, request)
 }
-
-/**
- * Component group for LLM services.
- */
-class Llm(
-    private val client: Client,
-    private val integrityClient: IntegrityClient,
-    private val userIdProvider: UserIdProvider,
-) {
-
-    val fenixMlpaService by lazyMonitored { FenixMlpaService(client) }
-    val mlpaProvider: MlpaLlmProvider by lazyMonitored {
-        MlpaLlmProvider(
-            MlpaTokenProvider.mlpaIntegrityHandshake(
-                integrityClient = integrityClient,
-                authenticationService = fenixMlpaService,
-                userIdProvider = userIdProvider,
-                packageName = PackageName(BuildConfig.APPLICATION_ID),
-            ),
-            fenixMlpaService,
-        )
-    }
-}

mobile/android/fenix/app/src/main/java/org/mozilla/fenix/components/llm/Llm.kt

@@ -0,0 +1,45 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+package org.mozilla.fenix.components.llm
+
+import mozilla.components.concept.fetch.Client
+import mozilla.components.concept.integrity.IntegrityClient
+import mozilla.components.lib.llm.mlpa.MlpaLlmProvider
+import mozilla.components.lib.llm.mlpa.MlpaTokenProvider
+import mozilla.components.lib.llm.mlpa.UserIdProvider
+import mozilla.components.lib.llm.mlpa.service.PackageName
+import org.mozilla.fenix.BuildConfig
+import org.mozilla.fenix.components.llm.ext.FxaAccessTokenProvider
+import org.mozilla.fenix.components.llm.ext.choose
+import org.mozilla.fenix.components.llm.ext.fxaTokenProvider
+import org.mozilla.fenix.perf.lazyMonitored
+
+/**
+ * Component group for LLM services.
+ */
+class Llm(
+    private val client: Client,
+    private val fxaTokenProvider: FxaAccessTokenProvider,
+    private val integrityClient: IntegrityClient,
+    private val userIdProvider: UserIdProvider,
+) {
+
+    val fenixMlpaService by lazyMonitored { FenixMlpaService(client) }
+
+    val mlpaProvider: MlpaLlmProvider by lazyMonitored {
+        MlpaLlmProvider(
+            MlpaTokenProvider.choose(
+                MlpaTokenProvider.fxaTokenProvider(fxaTokenProvider),
+                MlpaTokenProvider.mlpaIntegrityHandshake(
+                    integrityClient = integrityClient,
+                    authenticationService = fenixMlpaService,
+                    userIdProvider = userIdProvider,
+                    packageName = PackageName(BuildConfig.APPLICATION_ID),
+                ),
+            ),
+            fenixMlpaService,
+        )
+    }
+}

mobile/android/fenix/app/src/main/java/org/mozilla/fenix/components/llm/ext/FxaAccountManager.kt

@@ -0,0 +1,14 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+package org.mozilla.fenix.components.llm.ext
+
+import mozilla.components.service.fxa.manager.FxaAccountManager
+import mozilla.components.service.fxa.manager.SCOPE_PROFILE
+
+internal val FxaAccountManager.accessTokenProvider get() = FxaAccessTokenProvider {
+    authenticatedAccount()
+        ?.getAccessToken(SCOPE_PROFILE)
+        ?.token
+}

mobile/android/fenix/app/src/main/java/org/mozilla/fenix/components/llm/ext/MlpaTokenProvider.kt

@@ -0,0 +1,37 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+package org.mozilla.fenix.components.llm.ext
+
+import mozilla.components.lib.llm.mlpa.MlpaTokenProvider
+import mozilla.components.lib.llm.mlpa.service.AuthorizationToken
+
+internal class AllTokenProvidersFailed : IllegalStateException("All token providers failed to retrieve a token.")
+internal class FxaMissingAccessToken : IllegalStateException("Unable to get access token from FxaAccessTokenProvider")
+
+/** Convenience interface for getting an fxa access token. */
+fun interface FxaAccessTokenProvider {
+    /** Returns an access token or null */
+    suspend fun provide(): String?
+}
+
+/** Implementation of [MlpaTokenProvider] that takes the first successful token it receives.
+ * @param tokenProviders a list of [MlpaTokenProvider].
+ * @return an [MlpaTokenProvider].
+ */
+fun MlpaTokenProvider.Companion.choose(vararg tokenProviders: MlpaTokenProvider) = MlpaTokenProvider {
+    tokenProviders.firstNotNullOfOrNull { provider ->
+        provider.fetchToken().takeIf { it.isSuccess }
+    } ?: Result.failure(AllTokenProvidersFailed())
+}
+
+/** Implementation of [MlpaTokenProvider] that tries to fetch an fxa access token.
+ * @param tokenProvider a list of [FxaAccessTokenProvider].
+ * @return an [MlpaTokenProvider].
+ */
+fun MlpaTokenProvider.Companion.fxaTokenProvider(tokenProvider: FxaAccessTokenProvider) = MlpaTokenProvider {
+    tokenProvider.provide()?.let {
+        Result.success(AuthorizationToken(it))
+    } ?: Result.failure(FxaMissingAccessToken())
+}

mobile/android/fenix/app/src/main/java/org/mozilla/fenix/debugsettings/llm/LlmTools.kt

@@ -25,7 +25,7 @@ import mozilla.components.concept.llm.Llm
 import mozilla.components.concept.llm.Prompt
 import mozilla.components.lib.llm.mlpa.MlpaLlmProvider
 import org.mozilla.fenix.R
-import org.mozilla.fenix.components.Llm as LlmComponent
+import org.mozilla.fenix.components.llm.Llm as LlmComponent
 
 /**
  * Debug drawer view to test an [IntegrityClient].

mobile/android/fenix/app/src/main/java/org/mozilla/fenix/debugsettings/navigation/DebugDrawerRoute.kt

@@ -13,7 +13,7 @@ import mozilla.components.concept.storage.CreditCardsAddressesStorage
 import mozilla.components.concept.storage.LoginsStorage
 import org.mozilla.fenix.R
 import org.mozilla.fenix.components.ClientUUID
-import org.mozilla.fenix.components.Llm
+import org.mozilla.fenix.components.llm.Llm
 import org.mozilla.fenix.debugsettings.addons.ui.AddonsDebugToolsScreen
 import org.mozilla.fenix.debugsettings.addresses.AddressesDebugRegionRepository
 import org.mozilla.fenix.debugsettings.addresses.AddressesTools

mobile/android/fenix/app/src/main/java/org/mozilla/fenix/debugsettings/ui/FenixOverlay.kt

@@ -29,7 +29,7 @@ import mozilla.components.concept.storage.LoginsStorage
 import mozilla.telemetry.glean.Glean
 import org.mozilla.fenix.R
 import org.mozilla.fenix.components.ClientUUID
-import org.mozilla.fenix.components.Llm
+import org.mozilla.fenix.components.llm.Llm
 import org.mozilla.fenix.debugsettings.addresses.AddressesDebugRegionRepository
 import org.mozilla.fenix.debugsettings.addresses.AddressesTools
 import org.mozilla.fenix.debugsettings.addresses.FakeAddressesDebugRegionRepository
@@ -229,6 +229,6 @@ private fun FenixOverlayPreview() {
         creditCardsAddressesStorage = FakeCreditCardsAddressesStorage(),
         clientUUID = FakeClientUUID(),
         integrityClient = IntegrityClient.testSuccess,
-        llm = Llm(FakeClient(), FakeIntegrityClient(), FakeUserIdProvider()),
+        llm = Llm(FakeClient(), { null }, FakeIntegrityClient(), FakeUserIdProvider()),
     )
 }