[FAL] Add limited-use token support

FirebaseAI/Sources/AppCheckOptions.swift

@@ -0,0 +1,48 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+/// Configurable options for how App Check is used within a ``FirebaseAI`` instance.
+///
+/// Can be set when creating a ``FirebaseAI.Config``.
+@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+public struct AppCheckOptions: Sendable, Hashable {
+  /// Use `limitedUseTokens`, instead of the standard cached tokens, when sending requests
+  /// to the backend.
+  let requireLimitedUseTokens: Bool
+
+  /// Creates a new ``AppCheckOptions`` value.
+  ///
+  /// - Parameters:
+  ///   - requireLimitedUseTokens: When sending tokens to the backend, this option enables
+  ///     the usage of App Check's `limitedUseTokens` instead of the standard cached tokens.
+  ///
+  ///     A new `limitedUseToken` will be generated for each request; providing a lower attack
+  ///     surface for malicious parties to hijack tokens. When used alongside [replay protection](https://firebase.google.com/docs/app-check/custom-resource-backend#replay-protection),
+  ///     `limitedUseTokens` are also _consumed_ after each request, ensuring they can't be used
+  ///     again.
+  ///
+  ///     _To prevent breakage, 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
+  ///     > `limitedUseTokens`. Because `limitedUseTokens` are backwards compatable, 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 `limitedUseTokens` ahead of time will also allow you to enable replay
+  ///     > protection down the road (when support is added), without breaking your users.
+  public init(requireLimitedUseTokens: Bool = false) {
+    self.requireLimitedUseTokens = requireLimitedUseTokens
+  }
+}

FirebaseAI/Sources/FirebaseAI.swift

@@ -32,13 +32,19 @@ public final class FirebaseAI: Sendable {
   ///     ``FirebaseApp``.
   ///   - backend: The backend API for the Firebase AI SDK; if not specified, uses the default
   ///     ``Backend/googleAI()`` (Gemini Developer API).
+  ///   - config: Configuration options for the Firebase AI SDK that propogate across all models
+  ///     created. Uses default options when not specified, see the ``FirebaseAI.Config``
+  ///     documentation for more information.
   /// - Returns: A `FirebaseAI` instance, configured with the custom `FirebaseApp`.
   public static func firebaseAI(app: FirebaseApp? = nil,
-                                backend: Backend = .googleAI()) -> FirebaseAI {
+                                backend: Backend = .googleAI(),
+                                config: FirebaseAI
+                                  .Config = .config(appCheck: AppCheckOptions())) -> FirebaseAI {
     let instance = createInstance(
       app: app,
       location: backend.location,
-      apiConfig: backend.apiConfig
+      apiConfig: backend.apiConfig,
+      aiConfig: config
     )
     // Verify that the `FirebaseAI` instance is always configured with the production endpoint since
     // this is the public API surface for creating an instance.
@@ -90,7 +96,8 @@ public final class FirebaseAI: Sendable {
       tools: tools,
       toolConfig: toolConfig,
       systemInstruction: systemInstruction,
-      requestOptions: requestOptions
+      requestOptions: requestOptions,
+      aiConfig: aiConfig
     )
   }
 
@@ -126,21 +133,39 @@ public final class FirebaseAI: Sendable {
       apiConfig: apiConfig,
       generationConfig: generationConfig,
       safetySettings: safetySettings,
-      requestOptions: requestOptions
+      requestOptions: requestOptions,
+      aiConfig: aiConfig
     )
   }
 
   /// Class to enable FirebaseAI to register via the Objective-C based Firebase component system
   /// to include FirebaseAI in the userAgent.
   @objc(FIRVertexAIComponent) class FirebaseVertexAIComponent: NSObject {}
 
+  /// Configuration options for ``FirebaseAI``, which persists across all models.
+  @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
+  public struct Config: Sendable, Hashable {
+    /// Options for App Check specific behavior within a ``FirebaseAI`` instance.
+    let appCheck: AppCheckOptions
+
+    /// Creates a new ``FirebaseAI.Config`` value.
+    ///
+    /// - Parameters:
+    ///   - appCheck: Optionally configure certain behavior with how App Check is used.
+    public static func config(appCheck: AppCheckOptions = AppCheckOptions()) -> Config {
+      Config(appCheck: appCheck)
+    }
+  }
+
   // MARK: - Private
 
   /// Firebase data relevant to Firebase AI.
   let firebaseInfo: FirebaseInfo
 
   let apiConfig: APIConfig
 
+  let aiConfig: FirebaseAI.Config
+
   /// 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] = [:]
@@ -156,7 +181,7 @@ public final class FirebaseAI: Sendable {
   )
 
   static func createInstance(app: FirebaseApp?, location: String?,
-                             apiConfig: APIConfig) -> FirebaseAI {
+                             apiConfig: APIConfig, aiConfig: FirebaseAI.Config) -> FirebaseAI {
     guard let app = app ?? FirebaseApp.app() else {
       fatalError("No instance of the default Firebase app was found.")
     }
@@ -166,16 +191,26 @@ public final class FirebaseAI: Sendable {
     // Unlock before the function returns.
     defer { os_unfair_lock_unlock(&instancesLock) }
 
-    let instanceKey = InstanceKey(appName: app.name, location: location, apiConfig: apiConfig)
+    let instanceKey = InstanceKey(
+      appName: app.name,
+      location: location,
+      apiConfig: apiConfig,
+      aiConfig: aiConfig
+    )
     if let instance = instances[instanceKey] {
       return instance
     }
-    let newInstance = FirebaseAI(app: app, location: location, apiConfig: apiConfig)
+    let newInstance = FirebaseAI(
+      app: app,
+      location: location,
+      apiConfig: apiConfig,
+      aiConfig: aiConfig
+    )
     instances[instanceKey] = newInstance
     return newInstance
   }
 
-  init(app: FirebaseApp, location: String?, apiConfig: APIConfig) {
+  init(app: FirebaseApp, location: String?, apiConfig: APIConfig, aiConfig: FirebaseAI.Config) {
     guard let projectID = app.options.projectID else {
       fatalError("The Firebase app named \"\(app.name)\" has no project ID in its configuration.")
     }
@@ -195,6 +230,7 @@ public final class FirebaseAI: Sendable {
     )
     self.apiConfig = apiConfig
     self.location = location
+    self.aiConfig = aiConfig
   }
 
   func modelResourceName(modelName: String) -> String {
@@ -249,5 +285,6 @@ public final class FirebaseAI: Sendable {
     let appName: String
     let location: String?
     let apiConfig: APIConfig
+    let aiConfig: FirebaseAI.Config
   }
 }

FirebaseAI/Sources/GenerativeAIService.swift

@@ -30,9 +30,12 @@ struct GenerativeAIService {
 
   private let urlSession: URLSession
 
-  init(firebaseInfo: FirebaseInfo, urlSession: URLSession) {
+  private let aiConfig: FirebaseAI.Config
+
+  init(firebaseInfo: FirebaseInfo, urlSession: URLSession, aiConfig: FirebaseAI.Config) {
     self.firebaseInfo = firebaseInfo
     self.urlSession = urlSession
+    self.aiConfig = aiConfig
   }
 
   func loadRequest<T: GenerativeAIRequest>(request: T) async throws -> T.Response {
@@ -177,7 +180,7 @@ struct GenerativeAIService {
     urlRequest.setValue("application/json", forHTTPHeaderField: "Content-Type")
 
     if let appCheck = firebaseInfo.appCheck {
-      let tokenResult = await appCheck.getToken(forcingRefresh: false)
+      let tokenResult = await fetchAppCheckToken(appCheck: appCheck)
       urlRequest.setValue(tokenResult.token, forHTTPHeaderField: "X-Firebase-AppCheck")
       if let error = tokenResult.error {
         AILog.error(
@@ -207,6 +210,23 @@ struct GenerativeAIService {
     return urlRequest
   }
 
+  private func fetchAppCheckToken(appCheck: AppCheckInterop) async
+    -> FIRAppCheckTokenResultInterop {
+    if aiConfig.appCheck.requireLimitedUseTokens {
+      if let token = await appCheck.getLimitedUseToken?() {
+        return token
+      }
+
+      AILog.error(
+        code: .appCheckTokenFetchFailed,
+        "Missing getLimitedUseToken() function, but requireLimitedUseTokens was enabled."
+      )
+      // falls back to standard token
+    }
+
+    return await appCheck.getToken(forcingRefresh: false)
+  }
+
   private func httpResponse(urlResponse: URLResponse) throws -> HTTPURLResponse {
     // The following condition should always be true: "Whenever you make HTTP URL load requests, any
     // response objects you get back from the URLSession, NSURLConnection, or NSURLDownload class

FirebaseAI/Sources/GenerativeModel.swift

@@ -76,6 +76,7 @@ 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`.
+  ///   - aiConfig: Configuration for various behavior shared across models.
   init(modelName: String,
        modelResourceName: String,
        firebaseInfo: FirebaseInfo,
@@ -86,13 +87,14 @@ public final class GenerativeModel: Sendable {
        toolConfig: ToolConfig? = nil,
        systemInstruction: ModelContent? = nil,
        requestOptions: RequestOptions,
-       urlSession: URLSession = GenAIURLSession.default) {
+       urlSession: URLSession = GenAIURLSession.default, aiConfig: FirebaseAI.Config) {
     self.modelName = modelName
     self.modelResourceName = modelResourceName
     self.apiConfig = apiConfig
     generativeAIService = GenerativeAIService(
       firebaseInfo: firebaseInfo,
-      urlSession: urlSession
+      urlSession: urlSession,
+      aiConfig: aiConfig
     )
     self.generationConfig = generationConfig
     self.safetySettings = safetySettings

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

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