GenerativeModelSession.swift

// Copyright 2026 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.

// TODO: Remove the `#if compiler(>=6.2.3)` when Xcode 26.2 is the minimum supported version.
#if compiler(>=6.2.3)
  import Foundation
  #if canImport(FoundationModels)
    import FoundationModels
  #endif // canImport(FoundationModels)

  /// A session that handles multi-turn interactions with a generative model, similar to ``Chat``.
  ///
  /// A `GenerativeModelSession` retains history between requests. For single-turn requests to a
  /// model, use `generativeModelSession(model:tools:instructions:)` to start a new session.
  /// `GenerativeModelSession` is particularly useful for generating structured data.
  ///
  /// **Public Preview**: This API is a public preview and may be subject to change.
  ///
  /// Example usage:
  /// ```swift
  /// @Generable
  /// struct UserProfile {
  ///   @Guide(description: "A unique username for the user.")
  ///   var username: String
  ///
  ///   @Guide(description: "A short bio about the user, no more than 100 characters.")
  ///   var bio: String
  ///
  ///   @Guide(description: "A list of the user's favorite topics.", .count(3))
  ///   var favoriteTopics: [String]
  /// }
  ///
  /// let firebaseAI = // ... a `FirebaseAI` instance
  /// let session = firebaseAI.generativeModelSession(model: "gemini-model-name")
  /// let prompt = "Generate a user profile for a cat lover who enjoys hiking."
  /// let response = try await session.respond(to: prompt, generating: UserProfile.self)
  ///
  /// print("Username: \(response.content.username)")
  /// print("Bio: \(response.content.bio)")
  /// print("Favorite Topics: \(response.content.favoriteTopics.joined(separator: ", "))")
  /// ```
  public final class GenerativeModelSession: Sendable {
    let session: Chat
    let functionDeclarations: [String: FunctionDeclaration]

    // The maximum number of automatic back-and-forth turns the session will perform to resolve
    // function calls.
    //
    // This prevents infinite looping if the model consistently requests one or more function calls.
    //
    // TODO: Add ability to configure this setting.
    static let maxAutoFunctionCallTurns = 10

    /// Creates a new `GenerativeModelSession` with the given model.
    ///
    /// **Public Preview**: This API is a public preview and may be subject to change.
    /// - Parameter model: The `GenerativeModel` to use for generating content.
    init(model: GenerativeModel) {
      session = model.startChat()
      functionDeclarations = model.functionDeclarationsByName()
    }

    /// Sends a new prompt to the model and returns a `Response` containing the generated content as
    /// a `String`.
    ///
    /// **Public Preview**: This API is a public preview and may be subject to change.
    /// - Parameter prompt: The content to send to the model.
    /// - Parameter options: An optional `GenerationConfig` to override the model's default
    /// generation configuration.
    /// - Returns: A `Response` containing the generated content as a `String`.
    /// - Throws: A `GenerationError` if the model fails to generate a response.
    @discardableResult
    public nonisolated(nonsending)
    func respond(to prompt: PartsRepresentable..., options: GenerationConfig? = nil) async throws
      -> GenerativeModelSession.Response<String> {
      let model = FirebaseAI.SystemLanguageModel()

      // Hardcoded Prefer On-Device for Testing
      if model.isAvailable {
        let session = FirebaseAI.LanguageModelSession(model: model)
        return try await session.respond(
          to: prompt,
          schema: nil as FirebaseAI.GenerationSchema?,
          generating: String.self,
          includeSchemaInPrompt: false,
          options: options
        )
      } else {
        return try await respond(
          to: prompt,
          schema: nil as FirebaseAI.GenerationSchema?,
          generating: String.self,
          includeSchemaInPrompt: false,
          options: options
        )
      }
    }

    #if canImport(FoundationModels)
      /// Sends a new prompt to the model and returns a `Response` containing the generated content
      /// as
      /// `GeneratedContent`.
      ///
      /// **Public Preview**: This API is a public preview and may be subject to change.
      /// - Parameter prompt: The content to send to the model.
      /// - Parameter schema: The `GenerationSchema` to use for generating the content.
      /// - Parameter includeSchemaInPrompt: Whether to include the schema in the prompt to the
      /// model.
      /// - Parameter options: An optional `GenerationConfig` to override the model's default
      /// generation configuration.
      /// - Returns: A `Response` containing the generated content as `GeneratedContent`.
      /// - Throws: A `GenerationError` if the model fails to generate a response.
      @available(iOS 26.0, macOS 26.0, *)
      @available(tvOS, unavailable)
      @available(watchOS, unavailable)
      @discardableResult
      nonisolated(nonsending)
      func respond(to prompt: PartsRepresentable..., schema: FoundationModels.GenerationSchema,
                   includeSchemaInPrompt: Bool = true,
                   options: GenerationConfig? = nil) async throws
        -> GenerativeModelSession.Response<FirebaseAI.GeneratedContent> {
        // TODO: Replace `FoundationModels.GenerationSchema` and make this method public when
        // `FirebaseAI.GenerationSchema`'s public API is ready.
        return try await respond(
          to: prompt,
          schema: FirebaseAI.GenerationSchema(schema),
          generating: FirebaseAI.GeneratedContent.self,
          includeSchemaInPrompt: includeSchemaInPrompt,
          options: options
        )
      }

      /// Sends a new prompt to the model and returns a `Response` containing the generated content
      /// as
      /// a `Generable` type.
      ///
      /// **Public Preview**: This API is a public preview and may be subject to change.
      /// - Parameter prompt: The content to send to the model.
      /// - Parameter type: The `Generable` type to decode the response into.
      /// - Parameter includeSchemaInPrompt: Whether to include the schema in the prompt to the
      /// model.
      /// - Parameter options: An optional `GenerationConfig` to override the model's default
      /// generation configuration.
      /// - Returns: A `Response` containing the generated content as the specified `Generable`
      /// type.
      /// - Throws: A `GenerationError` if the model fails to generate a response.
      @available(iOS 26.0, macOS 26.0, *)
      @available(tvOS, unavailable)
      @available(watchOS, unavailable)
      @discardableResult
      public nonisolated(nonsending)
      func respond<Content>(to prompt: PartsRepresentable...,
                            generating type: Content.Type = Content.self,
                            includeSchemaInPrompt: Bool = true,
                            options: GenerationConfig? = nil) async throws
        -> GenerativeModelSession.Response<Content> where Content: Generable {
        let model = FirebaseAI.SystemLanguageModel()

        // Hardcoded Prefer On-Device for Testing
        if model.isAvailable {
          let session = FirebaseAI.LanguageModelSession(model: model)
          return try await session.respond(
            to: prompt,
            schema: FirebaseAI.GenerationSchema(Content.generationSchema),
            generating: type,
            includeSchemaInPrompt: includeSchemaInPrompt,
            options: options
          )
        } else {
          return try await respond(
            to: prompt,
            schema: FirebaseAI.GenerationSchema(Content.generationSchema),
            generating: type,
            includeSchemaInPrompt: includeSchemaInPrompt,
            options: options
          )
        }
      }

      /// Streams the model's response as `GeneratedContent`.
      ///
      /// **Public Preview**: This API is a public preview and may be subject to change.
      /// - Parameter prompt: The content to send to the model.
      /// - Parameter schema: The `GenerationSchema` to use for generating the content.
      /// - Parameter includeSchemaInPrompt: Whether to include the schema in the prompt to the
      /// model.
      /// - Parameter options: An optional `GenerationConfig` to override the model's default
      /// generation configuration.
      /// - Returns: A `ResponseStream` that yields snapshots of the generated content.
      @available(iOS 26.0, macOS 26.0, *)
      @available(tvOS, unavailable)
      @available(watchOS, unavailable)
      func streamResponse(to prompt: PartsRepresentable...,
                          schema: FoundationModels.GenerationSchema,
                          includeSchemaInPrompt: Bool = true, options: GenerationConfig? = nil)
        -> sending GenerativeModelSession.ResponseStream<
          FirebaseAI.GeneratedContent, FirebaseAI.GeneratedContent
        > {
        // TODO: Replace `FoundationModels.GenerationSchema` and make this method public when
        // `FirebaseAI.GenerationSchema`'s public API is ready.
        return streamResponse(
          to: prompt,
          schema: FirebaseAI.GenerationSchema(schema),
          generating: FirebaseAI.GeneratedContent.self,
          includeSchemaInPrompt: includeSchemaInPrompt,
          options: options
        )
      }

      /// Streams the model's response as a `Generable` type.
      ///
      /// **Public Preview**: This API is a public preview and may be subject to change.
      /// - Parameter prompt: The content to send to the model.
      /// - Parameter type: The `Generable` type to decode the response into.
      /// - Parameter includeSchemaInPrompt: Whether to include the schema in the prompt to the
      /// model.
      /// - Parameter options: An optional `GenerationConfig` to override the model's default
      /// generation configuration.
      /// - Returns: A `ResponseStream` that yields snapshots of the generated content.
      @available(iOS 26.0, macOS 26.0, *)
      @available(tvOS, unavailable)
      @available(watchOS, unavailable)
      public func streamResponse<Content>(to prompt: PartsRepresentable...,
                                          generating type: Content.Type = Content.self,
                                          includeSchemaInPrompt: Bool = true,
                                          options: GenerationConfig? = nil)
        -> sending GenerativeModelSession.ResponseStream<Content, Content.PartiallyGenerated>
        where Content: Generable {
        return streamResponse(
          to: prompt,
          schema: FirebaseAI.GenerationSchema(type.generationSchema),
          generating: type,
          includeSchemaInPrompt: includeSchemaInPrompt,
          options: options
        )
      }
    #endif // canImport(FoundationModels)

    /// Streams the model's response as a `String`.
    ///
    /// **Public Preview**: This API is a public preview and may be subject to change.
    /// - Parameter prompt: The content to send to the model.
    /// - Parameter options: An optional `GenerationConfig` to override the model's default
    /// generation configuration.
    /// - Returns: A `ResponseStream` that yields snapshots of the generated content.
    @available(macOS 12.0, watchOS 8.0, *)
    public func streamResponse(to prompt: PartsRepresentable..., options: GenerationConfig? = nil)
      -> sending GenerativeModelSession.ResponseStream<String, String> {
      return streamResponse(
        to: prompt,
        schema: nil,
        generating: String.self,
        includeSchemaInPrompt: false,
        options: options
      )
    }

    // MARK: - Internal

    private nonisolated(nonsending)
    func respond<Content>(to prompt: [PartsRepresentable], schema: FirebaseAI.GenerationSchema?,
                          generating type: Content.Type, includeSchemaInPrompt: Bool,
                          options: GenerationConfig?) async throws
      -> GenerativeModelSession.Response<Content> {
      let parts = [ModelContent(parts: prompt)]
      let config = try buildConfig(
        options: options,
        schema: schema,
        includeSchemaInPrompt: includeSchemaInPrompt
      )

      var response = try await session.sendMessage(parts, generationConfig: config)

      var autoFunctionCallTurns = 0
      while !response.functionCalls.isEmpty {
        guard autoFunctionCallTurns < GenerativeModelSession.maxAutoFunctionCallTurns else {
          throw GenerationError.internalError(
            GenerationError.Context(
              debugDescription: """
              The model exceeded the maximum allowed automatic function call iterations \
              (\(GenerativeModelSession.maxAutoFunctionCallTurns)).
              """
            ),
            underlyingError: FunctionCallingError.maxFunctionCallTurnsExceeded
          )
        }

        let functionResponses = try await execute(functionCalls: response.functionCalls)

        guard !functionResponses.isEmpty else { break }
        response = try await session.sendMessage(
          [ModelContent(role: "user", parts: functionResponses)],
          generationConfig: config
        )

        autoFunctionCallTurns += 1
      }

      let text: String
      if let responseText = response.text {
        text = responseText
      } else if let parts = response.candidates.first?.content.parts, !parts.isEmpty {
        text = ""
      } else {
        throw GenerationError.decodingFailure(
          GenerationError.Context(debugDescription: "No parts in response: \(response)")
        )
      }
      let generationID = response.responseID.map {
        #if canImport(FoundationModels) && IS_FOUNDATION_MODELS_SUPPORTED_PLATFORM
          if #available(iOS 26.0, macOS 26.0, visionOS 26.0, *) {
            return FirebaseAI.GenerationID(responseID: $0, generationID: GenerationID())
          }
        #endif // canImport(FoundationModels) && IS_FOUNDATION_MODELS_SUPPORTED_PLATFORM

        return FirebaseAI.GenerationID(responseID: $0, generationID: nil)
      }

      let rawContent = try Self.makeRawContent(
        from: text,
        generationID: generationID,
        hasSchema: schema != nil,
        isComplete: true
      )
      let content: Content = try Self.resolveContent(from: rawContent)

      return GenerativeModelSession.Response(
        content: content, rawContent: rawContent, rawResponse: response
      )
    }

    @available(macOS 12.0, watchOS 8.0, *)
    private func streamResponse<Content, PartialContent>(to prompt: [PartsRepresentable],
                                                         schema: FirebaseAI.GenerationSchema?,
                                                         generating type: Content.Type,
                                                         includeSchemaInPrompt: Bool,
                                                         options: GenerationConfig?)
      -> sending GenerativeModelSession.ResponseStream<Content, PartialContent> {
      let initialParts = [ModelContent(parts: prompt)]
      return GenerativeModelSession.ResponseStream { context in
        do {
          let config = try self.buildConfig(
            options: options,
            schema: schema,
            includeSchemaInPrompt: includeSchemaInPrompt
          )

          var currentParts = initialParts
          var generationID: FirebaseAI.GenerationID?
          var autoFunctionCallTurns = 0

          functionCallingLoop: while true {
            let stream = try self.session.sendMessageStream(currentParts, generationConfig: config)

            var streamedText = ""
            var functionCalls = [FunctionCallPart]()

            // 1. Create a buffer to hold the previous iteration's data in order to differentiate
            //    the last chunk to accurately set `isComplete`.
            var pendingChunkData: (
              text: String,
              id: FirebaseAI.GenerationID?,
              response: GenerateContentResponse
            )?

            for try await chunk in stream {
              functionCalls.append(contentsOf: chunk.functionCalls)

              let text: String
              if let responseText = chunk.text {
                text = responseText
              } else if let parts = chunk.candidates.first?.content.parts, !parts.isEmpty {
                text = ""
              } else {
                throw GenerationError.decodingFailure(
                  GenerationError.Context(debugDescription: "No parts in response: \(chunk)")
                )
              }

              // 2. If we have pending data, we now know it wasn't the last chunk.
              if let pending = pendingChunkData, !pending.text.isEmpty {
                let rawContent = try Self.makeRawContent(
                  from: pending.text,
                  generationID: pending.id,
                  hasSchema: schema != nil,
                  isComplete: false
                )
                let rawResult = GenerativeModelSession.ResponseStream<Content, PartialContent>
                  .RawResult(
                    rawContent: rawContent,
                    rawResponse: pending.response
                  )
                await context.yield(rawResult)
              }

              // 3. Update our cumulative state for the current chunk
              streamedText.append(text)
              if generationID == nil {
                generationID = chunk.responseID.map {
                  #if canImport(FoundationModels) && IS_FOUNDATION_MODELS_SUPPORTED_PLATFORM
                    if #available(iOS 26.0, macOS 26.0, visionOS 26.0, *) {
                      return FirebaseAI.GenerationID(
                        responseID: $0, generationID: FoundationModels.GenerationID()
                      )
                    }
                  #endif // canImport(FoundationModels) && IS_FOUNDATION_MODELS_SUPPORTED_PLATFORM

                  return FirebaseAI.GenerationID(responseID: $0, generationID: nil)
                }
              }

              // 4. Save the current state as the new pending chunk.
              pendingChunkData = (text: streamedText, id: generationID, response: chunk)
            }

            // Stream for the current turn finished. Check if there are function calls to handle.
            if !functionCalls.isEmpty {
              guard autoFunctionCallTurns < GenerativeModelSession.maxAutoFunctionCallTurns else {
                throw GenerationError.internalError(
                  GenerationError.Context(
                    debugDescription: """
                    The model exceeded the maximum allowed automatic function call iterations \
                    (\(GenerativeModelSession.maxAutoFunctionCallTurns)).
                    """
                  ),
                  underlyingError: FunctionCallingError.maxFunctionCallTurnsExceeded
                )
              }
              let functionResponses = try await self.execute(functionCalls: functionCalls)

              if !functionResponses.isEmpty {
                // Yield any pending text if it's not empty, but mark it as NOT complete yet.
                if let pending = pendingChunkData, !pending.text.isEmpty {
                  let rawContent = try Self.makeRawContent(
                    from: pending.text,
                    generationID: pending.id,
                    hasSchema: schema != nil,
                    isComplete: false
                  )
                  let rawResult = GenerativeModelSession.ResponseStream<Content, PartialContent>
                    .RawResult(
                      rawContent: rawContent,
                      rawResponse: pending.response
                    )
                  await context.yield(rawResult)
                }

                currentParts = [ModelContent(role: "user", parts: functionResponses)]
                autoFunctionCallTurns += 1
                continue functionCallingLoop
              }
            }

            // 5. The remaining pending chunk is the final one.
            if let finalChunk = pendingChunkData {
              let rawContent = try Self.makeRawContent(
                from: finalChunk.text,
                generationID: finalChunk.id,
                hasSchema: schema != nil,
                isComplete: true
              )
              let rawResult = GenerativeModelSession.ResponseStream<Content, PartialContent>
                .RawResult(
                  rawContent: rawContent,
                  rawResponse: finalChunk.response
                )
              await context.yield(rawResult)
            }

            break functionCallingLoop
          }

          await context.finish()
        } catch {
          await context.finish(throwing: error)
        }
      }
    }

    private func execute(functionCalls: [FunctionCallPart]) async throws -> [FunctionResponsePart] {
      var functionResponses = [FunctionResponsePart]()
      for functionCall in functionCalls {
        guard let functionDeclaration = functionDeclarations[functionCall.name] else {
          throw GenerationError.internalError(
            GenerationError.Context(debugDescription: """
            No function named "\(functionCall.name)" was declared.
            """),
            underlyingError: FunctionCallingError.invalidFunctionCall
          )
        }

        switch functionDeclaration.kind {
        case .manual:
          continue
        case let .foundationModels(tool):
          #if canImport(FoundationModels) && IS_FOUNDATION_MODELS_SUPPORTED_PLATFORM
            if #available(iOS 26.0, macOS 26.0, visionOS 26.0, *) {
              guard let tool = tool as? (any FoundationModels.Tool) else {
                assertionFailure("The value '\(tool)' is not a Foundation Models `Tool`.")
                throw TypeConversionError(
                  from: (any Sendable).self, to: (any FoundationModels.Tool).self
                )
              }
              try functionResponses.append(await FunctionDeclaration.call(
                tool: tool,
                functionCall: functionCall
              ))
              continue
            }
          #endif // canImport(FoundationModels) && IS_FOUNDATION_MODELS_SUPPORTED_PLATFORM
          assertionFailure("""
          A Foundation Models `Tool` '\(tool)' was provided but not running on a supported platform.
          """)
        }
      }

      return functionResponses
    }

    private func buildConfig(options: GenerationConfig?,
                             schema: FirebaseAI.GenerationSchema?,
                             includeSchemaInPrompt: Bool) throws -> GenerationConfig {
      var config = GenerationConfig.merge(
        session.generationConfig, with: options
      ) ?? GenerationConfig()

      if let schema {
        config.responseMIMEType = "application/json"
        config.responseJSONSchema = includeSchemaInPrompt ? try schema.toGeminiJSONSchema() : nil
        config.responseSchema = nil // `responseSchema` must not be set with `responseJSONSchema`
      }

      config.responseModalities = nil // Override to the default (text only)
      config.candidateCount = nil // Override to the default (one candidate)

      return config
    }

    private static func makeRawContent(from text: String, generationID: FirebaseAI.GenerationID?,
                                       hasSchema: Bool, isComplete: Bool) throws
      -> FirebaseAI.GeneratedContent {
      if hasSchema {
        return try FirebaseAI.GeneratedContent(json: text, id: generationID, isComplete: isComplete)
      }

      #if canImport(FoundationModels) && IS_FOUNDATION_MODELS_SUPPORTED_PLATFORM
        if #available(iOS 26.0, macOS 26.0, visionOS 26.0, *) {
          return FirebaseAI
            .GeneratedContent(
              kind: FoundationModels.GeneratedContent.Kind.string(text),
              id: generationID,
              isComplete: isComplete
            )
        }
      #endif // canImport(FoundationModels) && IS_FOUNDATION_MODELS_SUPPORTED_PLATFORM

      return FirebaseAI.GeneratedContent(
        kind: FirebaseAI.GeneratedContent.Kind.string(text),
        id: generationID,
        isComplete: isComplete
      )
    }

    static func resolveContent<T>(from rawContent: FirebaseAI.GeneratedContent) throws -> T {
      if let content = rawContent as? T {
        return content
      }

      #if canImport(FoundationModels) && IS_FOUNDATION_MODELS_SUPPORTED_PLATFORM
        if #available(iOS 26.0, macOS 26.0, visionOS 26.0, *), let contentMetatype = T
          .self as? (any FoundationModels.ConvertibleFromGeneratedContent.Type),
          let content = try contentMetatype.init(rawContent) as? T {
          return content
        }
      #endif // canImport(FoundationModels) && IS_FOUNDATION_MODELS_SUPPORTED_PLATFORM

      if let contentMetatype = T.self as? (any FirebaseAI.ConvertibleFromGeneratedContent.Type),
         let content = try contentMetatype.init(rawContent) as? T {
        return content
      }

      assertionFailure("Unsupported type: \(T.self).")
      // In release builds we throw an error instead of crashing but this state should be
      // unreachable based on the public API.
      throw GenerativeModelSession.TypeConversionError(
        from: type(of: rawContent), to: T.self
      )
    }
  }

  // MARK: - Response Types

  public extension GenerativeModelSession {
    /// The response from a `respond` call.
    struct Response<Content> {
      /// The generated content, decoded into the requested `Generable` type.
      public let content: Content
      /// The raw, undecoded `GeneratedContent` from the model.
      public let rawContent: FirebaseAI.GeneratedContent
      /// The raw `GenerateContentResponse` from the model.
      public let rawResponse: GenerateContentResponse
    }
  }

  public extension GenerativeModelSession {
    /// An asynchronous sequence of snapshots of the model's response.
    struct ResponseStream<Content, PartialContent>: AsyncSequence {
      public typealias Element = Snapshot

      /// A snapshot of the model's response at a point in time.
      public struct Snapshot {
        /// The partially generated content, decoded into the requested `Generable`'s partial type.
        public let content: PartialContent
        /// The raw, undecoded `GeneratedContent` from the model.
        public let rawContent: FirebaseAI.GeneratedContent
        /// The raw `GenerateContentResponse` from the model.
        public let rawResponse: GenerateContentResponse
      }

      private let rawStream: AsyncThrowingStream<RawResult, Error>
      private let context: StreamContext

      init(_ builder: @escaping @Sendable (StreamContext) async -> Void) {
        var extractedContinuation: AsyncThrowingStream<RawResult, Error>.Continuation!
        let stream = AsyncThrowingStream(RawResult.self) { continuation in
          extractedContinuation = continuation
        }
        rawStream = stream

        let context = StreamContext(continuation: extractedContinuation)
        self.context = context

        Task {
          await builder(context)
        }
      }

      /// An iterator that provides snapshots of the model's response.
      public struct AsyncIterator: AsyncIteratorProtocol {
        private var rawIterator: AsyncThrowingStream<RawResult, Error>.Iterator

        init(rawIterator: AsyncThrowingStream<RawResult, Error>.Iterator) {
          self.rawIterator = rawIterator
        }

        @available(iOS 18.0, macOS 15.0, macCatalyst 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *)
        public mutating func next(isolation actor: isolated (any Actor)?) async throws
          -> Snapshot? {
          var lastDecodingError: Error? = nil

          while let rawResult = try await rawIterator.next(isolation: actor) {
            do {
              // If it parses successfully, return the snapshot and discard any errors from previous
              // loop iterations.
              return try process(rawResult)
            } catch {
              // Intermediate failure (e.g., incomplete JSON that could not be parsed).
              // Hold onto the error and let the loop fetch the next chunk.
              lastDecodingError = error
            }
          }

          // If the last chunk processed resulted in an error, throw it.
          if let lastDecodingError {
            throw lastDecodingError
          }

          return nil
        }

        public mutating func next() async throws -> Snapshot? {
          var lastDecodingError: Error? = nil

          while let rawResult = try await rawIterator.next() {
            do {
              // If it parses successfully, return the snapshot and discard any errors from previous
              // loop iterations.
              return try process(rawResult)
            } catch {
              // Intermediate failure (e.g., incomplete JSON that could not be parsed).
              // Hold onto the error and let the loop fetch the next chunk.
              lastDecodingError = error
            }
          }

          // If the last chunk processed resulted in an error, throw it.
          if let lastDecodingError {
            throw lastDecodingError
          }

          return nil
        }

        private func process(_ rawResult: RawResult) throws -> Snapshot {
          let partialContent: PartialContent = try GenerativeModelSession
            .resolveContent(from: rawResult.rawContent)

          return Snapshot(
            content: partialContent,
            rawContent: rawResult.rawContent,
            rawResponse: rawResult.rawResponse
          )
        }
      }

      public func makeAsyncIterator() -> AsyncIterator {
        return AsyncIterator(rawIterator: rawStream.makeAsyncIterator())
      }

      /// Collects the entire streamed response into a single `Response`.
      /// - Returns: The final `Response` containing the fully generated content.
      /// - Throws: A `GenerationError` if the model fails to generate a response.
      public nonisolated(nonsending)
      func collect() async throws -> sending GenerativeModelSession.Response<Content> {
        let finalResult = try await context.value

        let content: Content = try GenerativeModelSession
          .resolveContent(from: finalResult.rawContent)
        return GenerativeModelSession.Response(
          content: content,
          rawContent: finalResult.rawContent,
          rawResponse: finalResult.rawResponse
        )
      }
    }
  }

  extension GenerativeModelSession.ResponseStream {
    struct RawResult: Sendable {
      let rawContent: FirebaseAI.GeneratedContent
      let rawResponse: GenerateContentResponse
    }

    actor StreamContext {
      private let continuation: AsyncThrowingStream<RawResult, Error>.Continuation
      private var finalResult: Result<RawResult, Error>?
      private var waitingContinuations: [CheckedContinuation<RawResult, Error>] = []
      private var latestRaw: RawResult?

      init(continuation: AsyncThrowingStream<RawResult, Error>.Continuation) {
        self.continuation = continuation
      }

      func yield(_ rawResult: RawResult) {
        // Prevent yielding new values if the stream has already been finalized
        guard finalResult == nil else { return }

        latestRaw = rawResult
        continuation.yield(rawResult)
      }

      func finish() {
        continuation.finish()
        finalize(with: nil)
      }

      func finish(throwing error: Error) {
        // TODO: Wrap `FoundationModels.GenerationError` errors into equivalent
        //       `GenerativeModelSession.GenerationError` values.
        continuation.finish(throwing: error)
        finalize(with: error)
      }

      var value: RawResult {
        get async throws {
          // 1. Return immediately if we already have the final result.
          if let result = finalResult {
            return try result.get()
          }

          // 2. Cancellation check: bail out early if the calling task was cancelled.
          try Task.checkCancellation()

          // 3. Suspend and wait.
          return try await withTaskCancellationHandler {
            try await withCheckedThrowingContinuation { continuation in
              self.add(continuation)
            }
          } onCancel: {
            // If the task calling `collect()` is cancelled, fail the stream immediately.
            // This guarantees the continuation is resumed with a CancellationError.
            Task {
              await self.finish(throwing: CancellationError())
            }
          }
        }
      }

      private func add(_ continuation: CheckedContinuation<RawResult, Error>) {
        if let result = finalResult {
          continuation.resume(with: result)
        } else {
          waitingContinuations.append(continuation)
        }
      }

      private func finalize(with error: Error?) {
        // Guards against resuming continuations multiple times, which would crash.
        guard finalResult == nil else { return }

        let result: Result<RawResult, Error>

        if let error = error {
          result = .failure(error)
        } else if let last = latestRaw {
          result = .success(last)
        } else {
          result = .failure(
            GenerativeModelSession.GenerationError.decodingFailure(
              GenerativeModelSession.GenerationError.Context(
                debugDescription: "No content generated in stream."
              )
            )
          )
        }

        finalResult = result

        for continuation in waitingContinuations {
          continuation.resume(with: result)
        }
        waitingContinuations.removeAll()
      }
    }
  }

  extension GenerativeModelSession {
    enum ErrorCodes: Int {
      // Generation Errors
      case decodingFailure = 1000

      // Internal Errors
      case typeConversionFailed = 2000
    }

    static let errorDomain = "\(Constants.baseErrorDomain).\(GenerativeModelSession.self)"

    /// An error that occurs during content generation.
    @nonexhaustive
    public enum GenerationError: Error, LocalizedError {
      /// A context providing more information about the generation error.
      public struct Context: Sendable {
        /// A debug description of the error.
        public let debugDescription: String

        init(debugDescription: String) {
          self.debugDescription = debugDescription
        }
      }

      /// The model's response could not be decoded.
      case decodingFailure(GenerativeModelSession.GenerationError.Context)

      case internalError(GenerativeModelSession.GenerationError.Context, underlyingError: any Error)
    }

    enum FunctionCallingError: Error, LocalizedError {
      case invalidFunctionCall
      case maxFunctionCallTurnsExceeded
    }

    struct TypeConversionError: CustomDebugStringConvertible, CustomNSError {
      public static var errorDomain: String { GenerativeModelSession.errorDomain }

      public var errorCode: Int { ErrorCodes.typeConversionFailed.rawValue }

      public var errorUserInfo: [String: Any] { [NSLocalizedDescriptionKey: debugDescription] }

      let debugDescription: String

      init(from fromType: Any.Type, to toType: Any.Type) {
        debugDescription = "Failed to convert from type '\(fromType)' to '\(toType)'."
      }
    }
  }
#endif // compiler(>=6.2.3)