[AI] Refactor stream parsing to skip intermediate failures (#15960)

Author: andrewheard

FirebaseAI/CHANGELOG.md

@@ -1,3 +1,7 @@
+# 12.12.0
+- [fixed] Updated `GenerativeModelSession` to handle intermediate JSON decoding
+  failures when streaming structured data. (#15960)
+
 # 12.11.0
 - [feature] **Public Preview**: Introduces `GenerativeModelSession` providing
   APIs for generating structured data from Gemini via the same `@Generable` and

FirebaseAI/Sources/GenerativeModelSession.swift

@@ -454,6 +454,8 @@
   public extension GenerativeModelSession {
     /// An asynchronous sequence of snapshots of the model's response.
     struct ResponseStream<Content, PartialContent>: AsyncSequence {
+      // TODO(#15962): Add unit tests for `ResponseStream`.
+
       public typealias Element = Snapshot
 
       /// A snapshot of the model's response at a point in time.
@@ -495,18 +497,52 @@
         @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? {
-          let rawResult = try await rawIterator.next(isolation: actor)
-          return try process(rawResult)
+          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? {
-          let rawResult = try await rawIterator.next()
-          return try process(rawResult)
-        }
+          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
+            }
+          }
 
-        private func process(_ rawResult: RawResult?) throws -> Snapshot? {
-          guard let rawResult else { return nil }
+          // 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)