[AI] Lazily initialize sessions in HybridModelSession (#16119)

Author: andrewheard

FirebaseAI/Sources/Types/Internal/HybridModelSession.swift

@@ -13,20 +13,38 @@
 // limitations under the License.
 
 #if compiler(>=6.2.3)
+  private import FirebaseCoreInternal
+
   final class HybridModelSession: _ModelSession {
-    private let primary: any _ModelSession
-    private let secondary: any _ModelSession
+    private let primaryModel: any LanguageModel
+    private let secondaryModel: any LanguageModel
+    private let tools: [any ToolRepresentable]?
+    private let instructions: String?
+
+    typealias SessionState = (primary: (any _ModelSession)?, secondary: (any _ModelSession)?)
+    private let lock: UnfairLock<SessionState>
+
+    init(primaryModel: any LanguageModel, secondaryModel: any LanguageModel,
+         tools: [any ToolRepresentable]?, instructions: String?) {
+      self.primaryModel = primaryModel
+      self.secondaryModel = secondaryModel
+      self.tools = tools
+      self.instructions = instructions
+      lock = UnfairLock((primary: nil, secondary: nil))
+    }
 
-    init(primary: any _ModelSession, secondary: any _ModelSession) {
-      self.primary = primary
-      self.secondary = secondary
+    enum SessionModel {
+      case primaryModel
+      case secondaryModel
     }
 
     /// Returns `true` if the session has history (i.e., it has already had one or more chat turns).
     ///
     /// > Important: This property is for **internal use only** and may change at any time.
     var _hasHistory: Bool {
-      return primary._hasHistory || secondary._hasHistory
+      return lock.withLock { state in
+        state.primary?._hasHistory == true || state.secondary?._hasHistory == true
+      }
     }
 
     /// Sends a prompt to the model and returns a ``_ModelSessionResponse``.
@@ -45,8 +63,12 @@
       return try await TaskLocals.$isHybridRequest.withValue(true) {
         // If the secondary session contains history then a previous fallback occurred.
         // Stick with the secondary session to maintain conversation consistency.
-        if secondary._hasHistory {
-          return try await secondary._respond(
+        let useSecondary = lock.withLock { state in
+          state.secondary?._hasHistory == true
+        }
+        if useSecondary {
+          let secondarySession = try getSession(for: .secondaryModel)
+          return try await secondarySession._respond(
             to: prompt,
             schema: schema,
             includeSchemaInPrompt: includeSchemaInPrompt,
@@ -56,20 +78,29 @@
 
         do {
           // First try the primary session.
-          return try await primary._respond(
+          let primarySession = try getSession(for: .primaryModel)
+          return try await primarySession._respond(
             to: prompt,
             schema: schema,
             includeSchemaInPrompt: includeSchemaInPrompt,
             options: options
           )
         } catch {
+          if Task.isCancelled || error is CancellationError {
+            throw error
+          }
+
           // Do not fallback to second session if the primary session contains history.
-          if primary._hasHistory {
+          let primaryHasHistory = lock.withLock { state in
+            state.primary?._hasHistory == true
+          }
+          if primaryHasHistory {
             throw error
           }
 
           // Fallback to the second session if the first fails or is unavailable.
-          return try await secondary._respond(
+          let secondarySession = try getSession(for: .secondaryModel)
+          return try await secondarySession._respond(
             to: prompt,
             schema: schema,
             includeSchemaInPrompt: includeSchemaInPrompt,
@@ -95,49 +126,95 @@
       return TaskLocals.$isHybridRequest.withValue(true) {
         // If the secondary session contains history then a previous fallback occurred.
         // Stick with the secondary session to maintain conversation consistency.
-        if secondary._hasHistory {
-          return secondary._streamResponse(
-            to: prompt,
-            schema: schema,
-            includeSchemaInPrompt: includeSchemaInPrompt,
-            options: options
-          )
+        let useSecondary = lock.withLock { state in
+          state.secondary?._hasHistory == true
         }
-
-        return AsyncThrowingStream { continuation in
-          let task = Task {
-            // First try the primary session.
-            let stream = primary._streamResponse(
+        if useSecondary {
+          do {
+            let secondarySession = try getSession(for: .secondaryModel)
+            return secondarySession._streamResponse(
               to: prompt,
               schema: schema,
               includeSchemaInPrompt: includeSchemaInPrompt,
               options: options
             )
+          } catch {
+            return AsyncThrowingStream { continuation in
+              continuation.finish(throwing: error)
+            }
+          }
+        }
 
-            var didYield = false
+        return AsyncThrowingStream { continuation in
+          let task = Task {
             do {
-              for try await snapshot in stream {
-                didYield = true
-                continuation.yield(snapshot)
+              // First try the primary session.
+              let primarySession = try self.getSession(for: .primaryModel)
+              let stream = primarySession._streamResponse(
+                to: prompt,
+                schema: schema,
+                includeSchemaInPrompt: includeSchemaInPrompt,
+                options: options
+              )
+
+              var didYield = false
+              do {
+                for try await snapshot in stream {
+                  didYield = true
+                  continuation.yield(snapshot)
+                }
+                continuation.finish()
+              } catch {
+                if Task.isCancelled || error is CancellationError {
+                  continuation.finish(throwing: error)
+                  return
+                }
+
+                // Do not fallback to second session if the primary session contains history or has
+                // already yielded data.
+                let primaryHasHistory = self.lock.withLock { state in
+                  state.primary?._hasHistory == true
+                }
+                if didYield || primaryHasHistory {
+                  continuation.finish(throwing: error)
+                  return
+                }
+
+                // Fallback to the second session if the first fails or is unavailable.
+                let secondarySession = try self.getSession(for: .secondaryModel)
+                let stream = secondarySession._streamResponse(
+                  to: prompt,
+                  schema: schema,
+                  includeSchemaInPrompt: includeSchemaInPrompt,
+                  options: options
+                )
+
+                do {
+                  for try await snapshot in stream {
+                    continuation.yield(snapshot)
+                  }
+                  continuation.finish()
+                } catch {
+                  continuation.finish(throwing: error)
+                }
               }
-              continuation.finish()
             } catch {
-              // Do not fallback to second session if the primary session contains history or has
-              // already yielded data.
-              if didYield || primary._hasHistory {
+              if Task.isCancelled || error is CancellationError {
                 continuation.finish(throwing: error)
                 return
               }
 
+              // Failure to create primary session.
               // Fallback to the second session if the first fails or is unavailable.
-              let stream = secondary._streamResponse(
-                to: prompt,
-                schema: schema,
-                includeSchemaInPrompt: includeSchemaInPrompt,
-                options: options
-              )
-
               do {
+                let secondarySession = try self.getSession(for: .secondaryModel)
+                let stream = secondarySession._streamResponse(
+                  to: prompt,
+                  schema: schema,
+                  includeSchemaInPrompt: includeSchemaInPrompt,
+                  options: options
+                )
+
                 for try await snapshot in stream {
                   continuation.yield(snapshot)
                 }
@@ -151,5 +228,32 @@
         }
       }
     }
+
+    private func getSession(for model: SessionModel) throws -> any _ModelSession {
+      let languageModel = (model == .primaryModel) ? primaryModel : secondaryModel
+
+      // 1. Check if it exists under lock.
+      let existing = lock.withLock { state in
+        (model == .primaryModel) ? state.primary : state.secondary
+      }
+      if let existing {
+        return existing
+      }
+
+      // 2. Create it outside the lock.
+      let session = try languageModel._startSession(tools: tools, instructions: instructions)
+
+      // 3. Try to store it under lock.
+      return lock.withLock { state in
+        if model == .primaryModel {
+          if let existing = state.primary { return existing }
+          state.primary = session
+        } else {
+          if let existing = state.secondary { return existing }
+          state.secondary = session
+        }
+        return session
+      }
+    }
   }
 #endif // compiler(>=6.2.3)

FirebaseAI/Sources/Types/Public/HybridModel.swift

@@ -46,41 +46,12 @@
     /// > Important: This method is for **internal use only** and may change at any time.
     public func _startSession(tools: [any ToolRepresentable]?,
                               instructions: String?) throws -> any _ModelSession {
-      var primarySession: (any _ModelSession)?
-      var primaryError: Error?
-      var secondarySession: (any _ModelSession)?
-      var secondaryError: Error?
-
-      do {
-        primarySession = try primary._startSession(tools: tools, instructions: instructions)
-      } catch {
-        primaryError = error
-      }
-
-      do {
-        secondarySession = try secondary._startSession(tools: tools, instructions: instructions)
-      } catch {
-        secondaryError = error
-      }
-
-      if let primarySession, let secondarySession {
-        return HybridModelSession(primary: primarySession, secondary: secondarySession)
-      } else if let primarySession {
-        return primarySession
-      } else if let secondarySession {
-        return secondarySession
-      } else {
-        let context = GenerativeModelSession.GenerationError.Context(
-          debugDescription: """
-          Failed to start session for model "\(primary._modelName)" with error: \
-          \(primaryError?.localizedDescription ?? "unknown error");
-          Failed to start session for model "\(secondary._modelName)" with error: \
-          \(secondaryError?.localizedDescription ?? "unknown error")
-          """
-        )
-
-        throw GenerativeModelSession.GenerationError.assetsUnavailable(context)
-      }
+      return HybridModelSession(
+        primaryModel: primary,
+        secondaryModel: secondary,
+        tools: tools,
+        instructions: instructions
+      )
     }
   }
 #endif // compiler(>=6.2.3)

FirebaseAI/Tests/TestApp/Tests/Integration/GenerativeModelSessionHybridTests.swift

@@ -208,6 +208,41 @@
       #expect(response.rawResponse.modelVersion == validModel._modelName)
     }
 
+    @Test(arguments: [InstanceConfig.vertexAI_v1beta_global])
+    @available(iOS 26.0, macOS 26.0, visionOS 26.0, *)
+    @available(tvOS, unavailable)
+    @available(watchOS, unavailable)
+    func streamResponseText_fallbackOnFoundationModelsError(_ config: InstanceConfig) async throws {
+      let firebaseAI = FirebaseAI.componentInstance(config)
+      let systemModel = FirebaseAI.SystemLanguageModel.default
+      let geminiModel = firebaseAI.geminiModel(name: ModelNames.gemini2_5_FlashLite)
+      let session = firebaseAI.generativeModelSession(
+        model: HybridModel(primary: systemModel, secondary: geminiModel)
+      )
+      let prompt = "In one sentence, why is the sky blue?"
+
+      let stream = session.streamResponse(to: prompt)
+
+      var receivedTexts = [String]()
+      var isComplete = false
+      for try await snapshot in stream {
+        let partial = snapshot.content
+        receivedTexts.append(partial)
+        isComplete = snapshot.rawContent.isComplete
+      }
+      #expect(isComplete)
+
+      let response = try await stream.collect()
+      let content = response.content
+      #expect(!content.isEmpty)
+
+      if await foundationModelsIsAvailable() {
+        #expect(response.rawResponse.modelVersion == systemModel._modelName)
+      } else {
+        #expect(response.rawResponse.modelVersion == geminiModel._modelName)
+      }
+    }
+
     /// Returns `true` if `FoundationModels.SystemLanguageModel` is available.
     ///
     /// This is a workaround for `SystemLanguageModel.isAvailable`, which returns `true` if *any*