Migrate to parameter instead of config struct

Author: daymxn

FirebaseAI/Sources/AppCheckOptions.swift

@@ -35,15 +35,33 @@ public final class FirebaseAI: Sendable {
   ///   - 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.
+  ///   - 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.
+  ///
+  ///     _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`.
   public static func firebaseAI(app: FirebaseApp? = nil,
                                 backend: Backend = .googleAI(),
-                                config: FirebaseAI.Config = .config()) -> FirebaseAI {
+                                useLimitedUseAppCheckTokens: Bool = false) -> FirebaseAI {
     let instance = createInstance(
       app: app,
       location: backend.location,
       apiConfig: backend.apiConfig,
-      aiConfig: config
+      useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens
     )
     // Verify that the `FirebaseAI` instance is always configured with the production endpoint since
     // this is the public API surface for creating an instance.
@@ -96,7 +114,7 @@ public final class FirebaseAI: Sendable {
       toolConfig: toolConfig,
       systemInstruction: systemInstruction,
       requestOptions: requestOptions,
-      aiConfig: aiConfig
+      useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens
     )
   }
 
@@ -133,37 +151,22 @@ public final class FirebaseAI: Sendable {
       generationConfig: generationConfig,
       safetySettings: safetySettings,
       requestOptions: requestOptions,
-      aiConfig: aiConfig
+      useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens
     )
   }
 
   /// 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, Equatable {
-    /// 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
+  let useLimitedUseAppCheckTokens: Bool
 
   /// A map of active `FirebaseAI` instances keyed by the `FirebaseApp` name and the `location`,
   /// in the format `appName:location`.
@@ -180,7 +183,7 @@ public final class FirebaseAI: Sendable {
   )
 
   static func createInstance(app: FirebaseApp?, location: String?,
-                             apiConfig: APIConfig, aiConfig: FirebaseAI.Config) -> FirebaseAI {
+                             apiConfig: APIConfig, useLimitedUseAppCheckTokens: Bool) -> FirebaseAI {
     guard let app = app ?? FirebaseApp.app() else {
       fatalError("No instance of the default Firebase app was found.")
     }
@@ -194,7 +197,7 @@ public final class FirebaseAI: Sendable {
       appName: app.name,
       location: location,
       apiConfig: apiConfig,
-      aiConfig: aiConfig
+      useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens
     )
     if let instance = instances[instanceKey] {
       return instance
@@ -203,13 +206,13 @@ public final class FirebaseAI: Sendable {
       app: app,
       location: location,
       apiConfig: apiConfig,
-      aiConfig: aiConfig
+      useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens
     )
     instances[instanceKey] = newInstance
     return newInstance
   }
 
-  init(app: FirebaseApp, location: String?, apiConfig: APIConfig, aiConfig: FirebaseAI.Config) {
+  init(app: FirebaseApp, location: String?, apiConfig: APIConfig, useLimitedUseAppCheckTokens: Bool) {
     guard let projectID = app.options.projectID else {
       fatalError("The Firebase app named \"\(app.name)\" has no project ID in its configuration.")
     }
@@ -229,7 +232,7 @@ public final class FirebaseAI: Sendable {
     )
     self.apiConfig = apiConfig
     self.location = location
-    self.aiConfig = aiConfig
+    self.useLimitedUseAppCheckTokens = useLimitedUseAppCheckTokens
   }
 
   func modelResourceName(modelName: String) -> String {
@@ -284,6 +287,6 @@ public final class FirebaseAI: Sendable {
     let appName: String
     let location: String?
     let apiConfig: APIConfig
-    let aiConfig: FirebaseAI.Config
+    let useLimitedUseAppCheckTokens: Bool
   }
 }

FirebaseAI/Sources/FirebaseAI.swift

@@ -30,12 +30,12 @@ struct GenerativeAIService {
 
   private let urlSession: URLSession
 
-  private let aiConfig: FirebaseAI.Config
+  private let useLimitedUseAppCheckTokens: Bool
 
-  init(firebaseInfo: FirebaseInfo, urlSession: URLSession, aiConfig: FirebaseAI.Config) {
+  init(firebaseInfo: FirebaseInfo, urlSession: URLSession, useLimitedUseAppCheckTokens: Bool) {
     self.firebaseInfo = firebaseInfo
     self.urlSession = urlSession
-    self.aiConfig = aiConfig
+    self.useLimitedUseAppCheckTokens = useLimitedUseAppCheckTokens
   }
 
   func loadRequest<T: GenerativeAIRequest>(request: T) async throws -> T.Response {
@@ -212,7 +212,7 @@ struct GenerativeAIService {
 
   private func fetchAppCheckToken(appCheck: AppCheckInterop) async throws
     -> FIRAppCheckTokenResultInterop {
-    if aiConfig.appCheck.requireLimitedUseTokens {
+    if useLimitedUseAppCheckTokens {
       if let token = await appCheck.getLimitedUseToken?() {
         return token
       }

FirebaseAI/Sources/GenerativeAIService.swift

@@ -76,7 +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.
+  ///   - useLimitedUseAppCheckTokens: Use App Check's limited-use tokens instead of the standard cached tokens.
   init(modelName: String,
        modelResourceName: String,
        firebaseInfo: FirebaseInfo,
@@ -88,14 +88,14 @@ public final class GenerativeModel: Sendable {
        systemInstruction: ModelContent? = nil,
        requestOptions: RequestOptions,
        urlSession: URLSession = GenAIURLSession.default,
-       aiConfig: FirebaseAI.Config = .config()) {
+       useLimitedUseAppCheckTokens: Bool) {
     self.modelName = modelName
     self.modelResourceName = modelResourceName
     self.apiConfig = apiConfig
     generativeAIService = GenerativeAIService(
       firebaseInfo: firebaseInfo,
       urlSession: urlSession,
-      aiConfig: aiConfig
+      useLimitedUseAppCheckTokens: useLimitedUseAppCheckTokens
     )
     self.generationConfig = generationConfig
     self.safetySettings = safetySettings

FirebaseAI/Sources/GenerativeModel.swift

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

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

@@ -156,13 +156,13 @@ class VertexComponentTests: XCTestCase {
       app: VertexComponentTests.app,
       location: location,
       apiConfig: APIConfig(service: .vertexAI(endpoint: .firebaseProxyProd), version: .v1beta),
-      aiConfig: .config()
+      useLimitedUseAppCheckTokens: false
     )
     let vertex2 = FirebaseAI.createInstance(
       app: VertexComponentTests.app,
       location: location,
       apiConfig: APIConfig(service: .vertexAI(endpoint: .firebaseProxyProd), version: .v1),
-      aiConfig: .config()
+      useLimitedUseAppCheckTokens: false
     )
 
     // Ensure they are different instances.
@@ -215,7 +215,7 @@ class VertexComponentTests: XCTestCase {
       app: app,
       location: nil,
       apiConfig: apiConfig,
-      aiConfig: .config()
+      useLimitedUseAppCheckTokens: false
     )
     let model = "test-model-name"
 
@@ -234,7 +234,7 @@ class VertexComponentTests: XCTestCase {
       app: app,
       location: nil,
       apiConfig: apiConfig,
-      aiConfig: .config()
+      useLimitedUseAppCheckTokens: false
     )
     let model = "test-model-name"
     let projectID = vertex.firebaseInfo.projectID
@@ -270,7 +270,7 @@ class VertexComponentTests: XCTestCase {
       app: app,
       location: nil,
       apiConfig: apiConfig,
-      aiConfig: .config()
+      useLimitedUseAppCheckTokens: false
     )
     let modelResourceName = vertex.modelResourceName(modelName: modelName)
     let expectedSystemInstruction = ModelContent(role: nil, parts: systemInstruction.parts)