Bug 2044193 - Exposed exceptions from LLM provider chooser

Author: fmasalha

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

@@ -7,7 +7,6 @@ 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. */
@@ -17,13 +16,20 @@ fun interface FxaAccessTokenProvider {
 }
 
 /** Implementation of [MlpaTokenProvider] that takes the first successful token it receives.
+ * When every provider fails, the last provider's failure is propagated as-is so its
+ * provider-specific error reaches logs, telemetry, and error services.
  * @param tokenProviders a list of [MlpaTokenProvider].
  * @return an [MlpaTokenProvider].
  */
 fun MlpaTokenProvider.Companion.choose(vararg tokenProviders: MlpaTokenProvider) = MlpaTokenProvider {
+    var lastResult = Result.failure<AuthorizationToken>(
+        IllegalStateException("choose() called with no token providers"),
+    )
     tokenProviders.firstNotNullOfOrNull { provider ->
-        provider.fetchToken().takeIf { it.isSuccess }
-    } ?: Result.failure(AllTokenProvidersFailed())
+        provider.fetchToken()
+            .also { lastResult = it }
+            .takeIf { it.isSuccess }
+    } ?: lastResult
 }
 
 /** Implementation of [MlpaTokenProvider] that tries to fetch an fxa access token.

mobile/android/fenix/app/src/test/java/org/mozilla/fenix/components/llm/ext/MlpaTokenProviderTest.kt

@@ -0,0 +1,53 @@
+/* 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 kotlinx.coroutines.test.runTest
+import mozilla.components.lib.llm.mlpa.MlpaTokenProvider
+import mozilla.components.lib.llm.mlpa.service.AuthorizationToken
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertSame
+import org.junit.Assert.assertTrue
+import org.junit.Assert.fail
+import org.junit.Test
+
+class MlpaTokenProviderTest {
+
+    @Test
+    fun `GIVEN the first provider succeeds WHEN fetching THEN its token is returned`() = runTest {
+        val expected = AuthorizationToken.Integrity("first")
+        val provider = MlpaTokenProvider.choose(
+             { Result.success(expected) },
+            { Result.failure(IllegalStateException("integrity failed")) },
+        )
+
+        assertEquals(expected, provider.fetchToken().getOrThrow() as AuthorizationToken.Integrity)
+    }
+
+    @Test
+    fun `GIVEN the first provider fails WHEN fetching THEN the next provider's token is returned`() = runTest {
+        val expected = AuthorizationToken.Integrity("second")
+        val provider = MlpaTokenProvider.choose(
+            { Result.failure(FxaMissingAccessToken()) },
+            { Result.success(expected) },
+        )
+
+        assertEquals(expected, provider.fetchToken().getOrThrow())
+    }
+
+    @Test
+    fun `GIVEN all providers fail WHEN fetching THEN the last provider's failure is propagated`() = runTest {
+        val terminal = IllegalStateException("integrity failed")
+        val provider = MlpaTokenProvider.choose(
+            { Result.failure(FxaMissingAccessToken()) },
+            { Result.failure(terminal) },
+        )
+
+        val result = provider.fetchToken()
+
+        assertTrue(result.isFailure)
+        assertSame(terminal, result.exceptionOrNull())
+    }
+}