[FAL] Move useLimitedUseAppCheckTokens property into FirebaseInfo + Update docs (#15234)

Author: daymxn

FirebaseAI/CHANGELOG.md

@@ -1,13 +1,10 @@
 # 12.2.0
 - [feature] Added support for returning thought summaries, which are synthesized
   versions of a model's internal reasoning process. (#15096)
-- [feature] Added a new configuration option to use limited-use App
-  Check tokens for attesting Firebase AI Logic requests. This enhances
-  security against replay attacks. To use this feature, configure it
-  explicitly via the new `useLimitedUseAppCheckTokens` parameter when
-  initializing `FirebaseAI`. We recommend migrating to limited-use
-  tokens now, so your app will be ready to take advantage of replay
-  protection when it becomes available for Firebase AI Logic. (#15099)
+- [feature] Added support for limited-use tokens with Firebase App Check. These short-lived tokens
+  provide greater protection for the APIs that give you access to Gemini and Imagen models. Learn
+  how to [enable usage of limited-use tokens](https://firebase.google.com/docs/ai-logic/app-check).
+  (#15099)
 
 # 12.0.0
 - [feature] Added support for Grounding with Google Search. (#15014)

FirebaseAI/Sources/FirebaseAI.swift

@@ -33,21 +33,12 @@ public final class FirebaseAI: Sendable {
   ///   - backend: The backend API for the Firebase AI SDK; if not specified, uses the default
   ///     ``Backend/googleAI()`` (Gemini Developer API).
   ///   - useLimitedUseAppCheckTokens: When sending tokens to the backend, this option enables
-  ///     the usage of App Check's limited-use tokens instead of the standard cached tokens.
-  ///
-  ///     A new limited-use tokens will be generated for each request; providing a smaller attack
-  ///     surface for malicious parties to hijack tokens. When used alongside replay protection,
-  ///     limited-use tokens are also _consumed_ after each request, ensuring they can't be used
-  ///     again.
+  ///     the usage of App Check's limited-use tokens instead of the standard cached tokens. Learn
+  ///     more about [limited-use tokens](https://firebase.google.com/docs/ai-logic/app-check),
+  ///     including their nuances, when to use them, and best practices for integrating them into
+  ///     your app.
   ///
   ///     _This flag is set to `false` by default._
-  ///
-  ///     > Important: Replay protection is not currently supported for the FirebaseAI backend.
-  ///     > While this feature is being developed, you can still migrate to using
-  ///     > limited-use tokens. Because limited-use tokens are backwards compatible, you can still
-  ///     > use them without replay protection. Due to their shorter TTL over standard App Check
-  ///     > tokens, they still provide a security benefit.
-  ///     >
   ///   > Migrating to limited-use tokens sooner minimizes disruption when support for replay
   ///   > protection is added.
   /// - Returns: A `FirebaseAI` instance, configured with the custom `FirebaseApp`.
@@ -110,8 +101,7 @@ public final class FirebaseAI: Sendable {
       tools: tools,
       toolConfig: toolConfig,
       systemInstruction: systemInstruction,
-      requestOptions: requestOptions,
-      useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens
+      requestOptions: requestOptions
     )
   }
 
@@ -147,8 +137,7 @@ public final class FirebaseAI: Sendable {
       apiConfig: apiConfig,
       generationConfig: generationConfig,
       safetySettings: safetySettings,
-      requestOptions: requestOptions,
-      useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens
+      requestOptions: requestOptions
     )
   }
 
@@ -163,8 +152,6 @@ public final class FirebaseAI: Sendable {
 
   let apiConfig: APIConfig
 
-  let useLimitedUseAppCheckTokens: Bool
-
   /// A map of active `FirebaseAI` instances keyed by the `FirebaseApp` name and the `location`,
   /// in the format `appName:location`.
   private nonisolated(unsafe) static var instances: [InstanceKey: FirebaseAI] = [:]
@@ -227,11 +214,11 @@ public final class FirebaseAI: Sendable {
       projectID: projectID,
       apiKey: apiKey,
       firebaseAppID: app.options.googleAppID,
-      firebaseApp: app
+      firebaseApp: app,
+      useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens
     )
     self.apiConfig = apiConfig
     self.location = location
-    self.useLimitedUseAppCheckTokens = useLimitedUseAppCheckTokens
   }
 
   func modelResourceName(modelName: String) -> String {

FirebaseAI/Sources/FirebaseInfo.swift

@@ -27,19 +27,22 @@ struct FirebaseInfo: Sendable {
   let projectID: String
   let apiKey: String
   let firebaseAppID: String
+  let useLimitedUseAppCheckTokens: Bool
   let app: FirebaseApp
 
   init(appCheck: AppCheckInterop? = nil,
        auth: AuthInterop? = nil,
        projectID: String,
        apiKey: String,
        firebaseAppID: String,
-       firebaseApp: FirebaseApp) {
+       firebaseApp: FirebaseApp,
+       useLimitedUseAppCheckTokens: Bool) {
     self.appCheck = appCheck
     self.auth = auth
     self.projectID = projectID
     self.apiKey = apiKey
     self.firebaseAppID = firebaseAppID
+    self.useLimitedUseAppCheckTokens = useLimitedUseAppCheckTokens
     app = firebaseApp
   }
 }

FirebaseAI/Sources/GenerativeAIService.swift

@@ -30,12 +30,9 @@ struct GenerativeAIService {
 
   private let urlSession: URLSession
 
-  private let useLimitedUseAppCheckTokens: Bool
-
-  init(firebaseInfo: FirebaseInfo, urlSession: URLSession, useLimitedUseAppCheckTokens: Bool) {
+  init(firebaseInfo: FirebaseInfo, urlSession: URLSession) {
     self.firebaseInfo = firebaseInfo
     self.urlSession = urlSession
-    self.useLimitedUseAppCheckTokens = useLimitedUseAppCheckTokens
   }
 
   func loadRequest<T: GenerativeAIRequest>(request: T) async throws -> T.Response {
@@ -212,7 +209,7 @@ struct GenerativeAIService {
 
   private func fetchAppCheckToken(appCheck: AppCheckInterop) async throws
     -> FIRAppCheckTokenResultInterop {
-    if useLimitedUseAppCheckTokens {
+    if firebaseInfo.useLimitedUseAppCheckTokens {
       if let token = await getLimitedUseAppCheckToken(appCheck: appCheck) {
         return token
       }
@@ -242,7 +239,7 @@ struct GenerativeAIService {
       Never
     >) in
       guard
-        useLimitedUseAppCheckTokens,
+        firebaseInfo.useLimitedUseAppCheckTokens,
         // `getLimitedUseToken(completion:)` is an optional protocol method. Optional binding
         // is performed to make sure `continuation` is called even if the method’s not implemented.
         let limitedUseTokenClosure = appCheck.getLimitedUseToken

FirebaseAI/Sources/GenerativeModel.swift

@@ -76,8 +76,6 @@ public final class GenerativeModel: Sendable {
   ///     only text content is supported.
   ///   - requestOptions: Configuration parameters for sending requests to the backend.
   ///   - urlSession: The `URLSession` to use for requests; defaults to `URLSession.shared`.
-  ///   - useLimitedUseAppCheckTokens: Use App Check's limited-use tokens instead of the standard
-  /// cached tokens.
   init(modelName: String,
        modelResourceName: String,
        firebaseInfo: FirebaseInfo,
@@ -88,15 +86,13 @@ public final class GenerativeModel: Sendable {
        toolConfig: ToolConfig? = nil,
        systemInstruction: ModelContent? = nil,
        requestOptions: RequestOptions,
-       urlSession: URLSession = GenAIURLSession.default,
-       useLimitedUseAppCheckTokens: Bool = false) {
+       urlSession: URLSession = GenAIURLSession.default) {
     self.modelName = modelName
     self.modelResourceName = modelResourceName
     self.apiConfig = apiConfig
     generativeAIService = GenerativeAIService(
       firebaseInfo: firebaseInfo,
-      urlSession: urlSession,
-      useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens
+      urlSession: urlSession
     )
     self.generationConfig = generationConfig
     self.safetySettings = safetySettings

FirebaseAI/Sources/Types/Public/Imagen/ImagenModel.swift

@@ -53,14 +53,12 @@ public final class ImagenModel {
        generationConfig: ImagenGenerationConfig?,
        safetySettings: ImagenSafetySettings?,
        requestOptions: RequestOptions,
-       urlSession: URLSession = GenAIURLSession.default,
-       useLimitedUseAppCheckTokens: Bool = false) {
+       urlSession: URLSession = GenAIURLSession.default) {
     self.modelResourceName = modelResourceName
     self.apiConfig = apiConfig
     generativeAIService = GenerativeAIService(
       firebaseInfo: firebaseInfo,
-      urlSession: urlSession,
-      useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens
+      urlSession: urlSession
     )
     self.generationConfig = generationConfig
     self.safetySettings = safetySettings

FirebaseAI/Tests/Unit/ChatTests.swift

@@ -68,7 +68,8 @@ final class ChatTests: XCTestCase {
           projectID: "my-project-id",
           apiKey: "API_KEY",
           firebaseAppID: "My app ID",
-          firebaseApp: app
+          firebaseApp: app,
+          useLimitedUseAppCheckTokens: false
         ),
         apiConfig: FirebaseAI.defaultVertexAIAPIConfig,
         tools: nil,

FirebaseAI/Tests/Unit/GenerativeModelVertexAITests.swift

@@ -507,13 +507,13 @@ final class GenerativeModelVertexAITests: XCTestCase {
       modelName: testModelName,
       modelResourceName: testModelResourceName,
       firebaseInfo: GenerativeModelTestUtil.testFirebaseInfo(
-        appCheck: AppCheckInteropFake(token: appCheckToken)
+        appCheck: AppCheckInteropFake(token: appCheckToken),
+        useLimitedUseAppCheckTokens: true
       ),
       apiConfig: apiConfig,
       tools: nil,
       requestOptions: RequestOptions(),
-      urlSession: urlSession,
-      useLimitedUseAppCheckTokens: true
+      urlSession: urlSession
     )
     MockURLProtocol
       .requestHandler = try GenerativeModelTestUtil.httpRequestHandler(

FirebaseAI/Tests/Unit/TestUtilities/GenerativeModelTestUtil.swift

@@ -103,7 +103,8 @@ enum GenerativeModelTestUtil {
 
   static func testFirebaseInfo(appCheck: AppCheckInterop? = nil,
                                auth: AuthInterop? = nil,
-                               privateAppID: Bool = false) -> FirebaseInfo {
+                               privateAppID: Bool = false,
+                               useLimitedUseAppCheckTokens: Bool = false) -> FirebaseInfo {
     let app = FirebaseApp(instanceWithName: "testApp",
                           options: FirebaseOptions(googleAppID: "ignore",
                                                    gcmSenderID: "ignore"))
@@ -114,7 +115,8 @@ enum GenerativeModelTestUtil {
       projectID: "my-project-id",
       apiKey: "API_KEY",
       firebaseAppID: "My app ID",
-      firebaseApp: app
+      firebaseApp: app,
+      useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens
     )
   }
 }