Merge remote-tracking branch 'upstream/main' into JesusRojass/#11903

Author: JesusRojass

FirebaseAI/CHANGELOG.md

@@ -1,4 +1,6 @@
 # Unreleased
+- [fixed] Fixed a decoding failure in `GenerateContentResponse` when the Vertex AI
+  backend returns citation metadata with a missing `endIndex`. (#16328)
 - [fixed] Fixed an issue where `generateContentStream` could stall indefinitely
   on mid-stream network drops. (#16298)
 - [fixed] Fixed a resource leak where background network requests would

FirebaseAI/Sources/GenerateContentResponse.swift

@@ -684,12 +684,10 @@ extension CitationMetadata: Decodable {
   public init(from decoder: any Decoder) throws {
     let container = try decoder.container(keyedBy: CodingKeys.self)
 
-    // Decode for Google API if `citationSources` key is present.
-    if container.contains(.citationSources) {
-      citations = try container.decode([Citation].self, forKey: .citationSources)
-    } else { // Fallback to default Vertex AI decoding.
-      citations = try container.decode([Citation].self, forKey: .citations)
-    }
+    let decodedCitations = try container.decodeIfPresent([Citation].self, forKey: .citationSources)
+      ?? container.decodeIfPresent([Citation].self, forKey: .citations)
+      ?? []
+    citations = decodedCitations.filter { !$0.isEmpty }
   }
 }
 
@@ -703,10 +701,19 @@ extension Citation: Decodable {
     case publicationDate
   }
 
+  var isEmpty: Bool {
+    startIndex == 0 &&
+      endIndex == 0 &&
+      uri == nil &&
+      title == nil &&
+      license == nil &&
+      publicationDate == nil
+  }
+
   public init(from decoder: any Decoder) throws {
     let container = try decoder.container(keyedBy: CodingKeys.self)
     startIndex = try container.decodeIfPresent(Int.self, forKey: .startIndex) ?? 0
-    endIndex = try container.decode(Int.self, forKey: .endIndex)
+    endIndex = try container.decodeIfPresent(Int.self, forKey: .endIndex) ?? startIndex
 
     if let uri = try container.decodeIfPresent(String.self, forKey: .uri), !uri.isEmpty {
       self.uri = uri

FirebaseAI/Tests/Unit/Types/CitationTests.swift

@@ -99,14 +99,16 @@ final class CitationTests: XCTestCase {
     XCTAssertNil(citation.publicationDate)
   }
 
-  func testDecodeCitation_missingEndIndex_throws() throws {
+  func testDecodeCitation_missingEndIndex_success() throws {
     let json = """
     {
       "startIndex" : 10
     }
     """
     let jsonData = try XCTUnwrap(json.data(using: .utf8))
 
-    XCTAssertThrowsError(try decoder.decode(Citation.self, from: jsonData))
+    let citation = try decoder.decode(Citation.self, from: jsonData)
+    XCTAssertEqual(citation.startIndex, 10)
+    XCTAssertEqual(citation.endIndex, 10)
   }
 }

FirebaseAI/Tests/Unit/Types/GenerateContentResponseTests.swift

@@ -343,6 +343,62 @@ final class GenerateContentResponseTests: XCTestCase {
     }
   }
 
+  func testDecodeGenerateContentResponse_missingCitationEndIndex_success() throws {
+    let json = """
+    {
+      "candidates": [
+        {
+          "content": { "role": "model", "parts": [ { "text": "Some text." } ] },
+          "finishReason": "STOP",
+          "citationMetadata": {
+            "citations": [
+              {
+                "uri": "https://example.com/source"
+              }
+            ]
+          }
+        }
+      ]
+    }
+    """
+    let jsonData = try XCTUnwrap(json.data(using: .utf8))
+
+    let response = try jsonDecoder.decode(GenerateContentResponse.self, from: jsonData)
+
+    XCTAssertEqual(response.candidates.count, 1)
+    let candidate = try XCTUnwrap(response.candidates.first)
+    let citation = try XCTUnwrap(candidate.citationMetadata?.citations.first)
+    XCTAssertEqual(citation.uri, "https://example.com/source")
+    XCTAssertEqual(citation.endIndex, 0)
+    XCTAssertEqual(citation.startIndex, 0)
+  }
+
+  func testDecodeGenerateContentResponse_emptyCitation_success() throws {
+    let json = """
+    {
+      "candidates": [
+        {
+          "content": { "role": "model", "parts": [ { "text": "Some text." } ] },
+          "finishReason": "STOP",
+          "citationMetadata": {
+            "citations": [
+              {}
+            ]
+          }
+        }
+      ]
+    }
+    """
+    let jsonData = try XCTUnwrap(json.data(using: .utf8))
+
+    let response = try jsonDecoder.decode(GenerateContentResponse.self, from: jsonData)
+
+    XCTAssertEqual(response.candidates.count, 1)
+    let candidate = try XCTUnwrap(response.candidates.first)
+    let citations = try XCTUnwrap(candidate.citationMetadata?.citations)
+    XCTAssertTrue(citations.isEmpty)
+  }
+
   // MARK: - Candidate.isEmpty
 
   func testCandidateIsEmpty_allEmpty_isTrue() throws {